add support for simplefsdp+ep

ruisizhang123 · ruisizhang123 · commit 6e3773525bd6 · 2025-09-26T16:32:02.000-07:00
diff --git a/.github/workflows/integration_test_8gpu_simple_fsdp.yaml b/.github/workflows/integration_test_8gpu_simple_fsdp.yaml
@@ -47,4 +47,4 @@ jobs:
         python -m pip install --force-reinstall --pre torch --index-url https://download.pytorch.org/whl/nightly/cu126
 
         mkdir artifacts-to-be-uploaded
-        python -m torchtitan.experiments.simple_fsdp.tests.integration_tests artifacts-to-be-uploaded --ngpu 8
+        python -m torchtitan.experiments.simple_fsdp.tests.llama3_integration_tests artifacts-to-be-uploaded --ngpu 8
diff --git a/torchtitan/experiments/__init__.py b/torchtitan/experiments/__init__.py
@@ -4,4 +4,4 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-_supported_experiments = frozenset(["flux", "llama4", "qwen3", "simple_fsdp", "vlm"])
+_supported_experiments = frozenset(["flux", "llama4", "qwen3", "simple_fsdp.llama3", "simple_fsdp.deepseek_v3", "vlm"])
diff --git a/torchtitan/experiments/simple_fsdp/README.md b/torchtitan/experiments/simple_fsdp/README.md
@@ -12,10 +12,18 @@ This folder includes an experimental frontend implementation for [SimpleFSDP: Si
 
 ### Run SimpleFSDP Training on Llama 3
 
+#### Training Llama3 models
+
 ```bash
 CONFIG_FILE="./torchtitan/models/llama3/train_configs/llama3_8b.toml" ./run_train.sh --model.name simple_fsdp --compile.enable
 ```
 
+#### Training DeepSeek_v3 models
+
+```bash
+CONFIG_FILE="./torchtitan/models/deepseek_v3/train_configs/debug_model.toml" ./run_train.sh --model.name simple_fsdp --compile.enable
+```
+
 ### Composability Support
 
 Some of the features require the updates from PyTorch, with which we are working on providing composability support for the following features:
@@ -30,7 +38,9 @@ Some of the features require the updates from PyTorch, with which we are working
 |Pipeline Parallelism| ✅ |
 |Distributed Checkpointing| ✅ |
 |Float8 Training| 🚧 |
-
+|Expert Parallelism | ✅ |
+|Expert Parallelism + Activation Checkpointing| 🚧 |
+|Expert Parallelism + DualPipe| 🚧 |
 
 ### Citation
 
diff --git a/torchtitan/experiments/simple_fsdp/deepseek_v3/__init__.py b/torchtitan/experiments/simple_fsdp/deepseek_v3/__init__.py
@@ -0,0 +1,33 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+#
+# Copyright (c) Meta Platforms, Inc. All Rights Reserved.
+
+from torchtitan.components.loss import build_cross_entropy_loss
+from torchtitan.components.lr_scheduler import build_lr_schedulers
+from torchtitan.components.optimizer import build_optimizers_with_moe_load_balancing
+from torchtitan.components.tokenizer import build_hf_tokenizer
+from torchtitan.datasets.hf_datasets import build_hf_dataloader
+from torchtitan.models.deepseek_v3 import deepseekv3_configs
+from torchtitan.models.llama3 import pipeline_llama
+from torchtitan.protocols.train_spec import TrainSpec
+from .deepseek_v3_model import SimpleFSDPDeepSeekV3Model
+from .deepseek_v3_parallelize import parallelize_deepseekv3
+
+
+def get_train_spec() -> TrainSpec:
+    return TrainSpec(
+        name="simple_fsdp.deepseek_v3",
+        model_cls=SimpleFSDPDeepSeekV3Model,
+        model_args=deepseekv3_configs,
+        parallelize_fn=parallelize_deepseekv3,
+        pipelining_fn=pipeline_llama,
+        build_optimizers_fn=build_optimizers_with_moe_load_balancing,
+        build_lr_schedulers_fn=build_lr_schedulers,
+        build_dataloader_fn=build_hf_dataloader,
+        build_tokenizer_fn=build_hf_tokenizer,
+        build_loss_fn=build_cross_entropy_loss,
+    )
diff --git a/torchtitan/experiments/simple_fsdp/deepseek_v3/deepseek_v3_model.py b/torchtitan/experiments/simple_fsdp/deepseek_v3/deepseek_v3_model.py
@@ -0,0 +1,18 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from torchtitan.models.deepseek_v3 import DeepSeekV3Model, DeepSeekV3ModelArgs
+from ..simple_fsdp import enable_active_parametrization
+
+
+class SimpleFSDPDeepSeekV3Model(DeepSeekV3Model):
+    def __init__(self, model_args: DeepSeekV3ModelArgs):
+        super().__init__(model_args)
+        self.init_weights()
+
+    def init_weights(self, *args, **kwargs):
+        with enable_active_parametrization():
+            super().init_weights(*args, **kwargs)
diff --git a/torchtitan/experiments/simple_fsdp/deepseek_v3/deepseek_v3_parallelize.py b/torchtitan/experiments/simple_fsdp/deepseek_v3/deepseek_v3_parallelize.py
@@ -0,0 +1,181 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+import torch.nn as nn
+from torch.distributed.device_mesh import DeviceMesh
+
+from torchtitan.config import JobConfig, TORCH_DTYPE_MAP
+from torchtitan.distributed import ParallelDims
+from torchtitan.distributed.tensor_parallel import maybe_enable_async_tp
+from torchtitan.experiments.llama4.infra.parallelize import apply_moe_ep_tp
+from torchtitan.models.deepseek_v3.infra.parallelize import apply_non_moe_tp
+from torchtitan.models.llama3.infra.parallelize import apply_ac
+from torchtitan.tools.logging import logger
+
+from ..simple_fsdp import data_parallel, MixedPrecisionPolicy
+
+# Adapted from llama4/infra/parallelize.py
+def parallelize_deepseekv3(
+    model: nn.Module,
+    parallel_dims: ParallelDims,
+    job_config: JobConfig,
+):
+    world_mesh = parallel_dims.world_mesh
+    # TODO: TP currently cannot handle uneven seq_len because we set
+    #       `use_local_output=True` to use plain Tensors for legacy reasons.
+    #       Need to revisit this.
+    assert (
+        job_config.training.seq_len % parallel_dims.seq_len_divisor == 0
+    ), f"""
+        Sequence length {job_config.training.seq_len} must be divisible by the product of TP degree
+        ({parallel_dims.tp}) and 2 * CP degree ({parallel_dims.cp}), i.e. {parallel_dims.seq_len_divisor}.
+        """
+
+    if (
+        job_config.parallelism.context_parallel_degree > 1
+        and model.model_args.use_flex_attn
+    ):
+        raise NotImplementedError("CP support for FlexAttention is still in progress.")
+
+    if parallel_dims.tp_enabled:
+        if job_config.parallelism.enable_async_tensor_parallel:
+            # TODO(jianiw): This branch needs to be tested and enabled
+            raise NotImplementedError(
+                "Currently, async TP is not tested for deepseekv3. \
+                torch.compile is not supported yet, which is required for async TP."
+            )
+
+        enable_float8_linear = "float8" in job_config.model.converters
+        float8_is_rowwise = job_config.float8.recipe_name in (
+            "rowwise",
+            "rowwise_with_gw_hp",
+        )
+
+        enable_float8_tensorwise_tp = enable_float8_linear and not float8_is_rowwise
+        if enable_float8_tensorwise_tp:
+            # TODO(jianiw): This branch needs to be tested and enabled
+            raise NotImplementedError(
+                "Currently, float8 tensorwise TP is not tested for deepseekv3"
+            )
+
+        apply_non_moe_tp(
+            model,
+            world_mesh["tp"],
+            loss_parallel=not job_config.parallelism.disable_loss_parallel,
+            enable_float8_tensorwise_tp=False,
+        )
+        maybe_enable_async_tp(job_config, world_mesh["tp"])
+
+    if parallel_dims.tp_enabled or parallel_dims.ep_enabled:
+        apply_moe_ep_tp(
+            model,
+            tp_mesh=world_mesh["tp"] if parallel_dims.tp_enabled else None,
+            ep_mesh=world_mesh["ep"] if parallel_dims.ep_enabled else None,
+            ep_tp_mesh=(
+                world_mesh["ep", "tp"]
+                if parallel_dims.tp_enabled
+                and parallel_dims.ep_enabled
+                and parallel_dims.etp_enabled
+                else None
+            ),
+            etp_enabled=parallel_dims.etp_enabled,
+        )
+
+    if job_config.activation_checkpoint.mode != "none":
+        apply_ac(model, job_config.activation_checkpoint)
+
+    mp_policy = MixedPrecisionPolicy(
+        param_dtype=TORCH_DTYPE_MAP[job_config.training.mixed_precision_param],
+        reduce_dtype=TORCH_DTYPE_MAP[job_config.training.mixed_precision_reduce],
+    )
+
+    # apply data parallel
+    dp_mesh: DeviceMesh | None = None
+    if (
+        parallel_dims.fsdp_enabled
+        or parallel_dims.ep_enabled
+        or parallel_dims.dp_replicate_enabled
+    ):
+        if parallel_dims.dp_replicate_enabled:
+            if parallel_dims.dp_shard_enabled or parallel_dims.cp_enabled:
+                dp_mesh_dim_names = ("dp_replicate", "dp_shard_cp")
+                dp_mode = "hybrid_shard"
+            else:
+                dp_mesh_dim_names = ("dp_replicate",)
+                dp_mode = "replicate"
+        else:
+            dp_mesh_dim_names = ("dp_shard_cp",)
+            dp_mode = "fully_shard"
+
+        dp_mesh = world_mesh[tuple(dp_mesh_dim_names)]
+        # the mesh dim names of which the MoE params are sharded on via FSDP/HSDP
+        dp_mod_ep_mesh_dim_names = []
+
+        if parallel_dims.ep_enabled:
+            if parallel_dims.dp_replicate_enabled:
+                dp_mod_ep_mesh_dim_names.append("dp_replicate")
+            dp_mod_ep_mesh_dim_names.append("dp_shard_mod_ep")
+
+        # allow users to parse their ep/tp mesh_dim_names to simple_fsdp
+        dp_tp_ep_mesh_name = {"dp": "dp_shard", "ep": "", "tp": ""}
+        if parallel_dims.ep_enabled and not parallel_dims.tp_enabled:
+            dp_tp_ep_mesh_name["ep"] = "ep"
+        elif not parallel_dims.ep_enabled and parallel_dims.tp_enabled:
+            dp_tp_ep_mesh_name["tp"] = "tp"
+        elif parallel_dims.ep_enabled and parallel_dims.tp_enabled:
+            dp_tp_ep_mesh_name["ep"], dp_tp_ep_mesh_name["tp"] = "ep", "tp"
+        shared_expert_dp_shard_dim = 0
+
+        for _, transformer_block in model.layers.items():
+            if transformer_block.moe_enabled and parallel_dims.ep_enabled:
+                if (
+                    world_mesh[tuple(dp_mod_ep_mesh_dim_names)].size() * parallel_dims.ep
+                    > transformer_block.moe.experts.num_experts
+                ):
+                    shared_expert_dp_shard_dim = 1
+
+                transformer_block.moe.experts = data_parallel(
+                    transformer_block.moe.experts,
+                    world_mesh[tuple(dp_mod_ep_mesh_dim_names)],
+                    dp_mode,
+                    ac_mode=job_config.activation_checkpoint.mode,
+                    mp_policy=mp_policy,
+                    dp_tp_ep_mesh_name=dp_tp_ep_mesh_name,
+                    dp_shard_dim=shared_expert_dp_shard_dim,
+                )
+                transformer_block.moe.shared_experts = data_parallel(
+                    transformer_block.moe.shared_experts,
+                    dp_mesh,
+                    dp_mode,
+                    ac_mode=job_config.activation_checkpoint.mode,
+                    mp_policy=mp_policy,
+                    dp_tp_ep_mesh_name=dp_tp_ep_mesh_name,
+                )
+                # TODO(ruisizhang123): support set_gradient_divide_factor in simplefsdp
+                # transformer_block.moe.experts.set_gradient_divide_factor(
+                #     parallel_dims.fsdp_gradient_divide_factor,
+                # )
+
+        model = data_parallel(
+            model,
+            dp_mesh,
+            dp_mode,
+            ac_mode=job_config.activation_checkpoint.mode,
+            mp_policy=mp_policy,
+            dp_tp_ep_mesh_name=dp_tp_ep_mesh_name,
+        )
+
+        logger.info(
+            "Applied Data Parallel (simple_fsdp) (dp mode=%s) to the model", dp_mode
+        )
+
+    if job_config.compile.enable:
+        torch._inductor.config.reorder_for_peak_memory = False
+        torch._dynamo.config.capture_scalar_outputs = True
+        model = torch.compile(model, fullgraph=True)
+
+    return model
diff --git a/torchtitan/experiments/simple_fsdp/llama3/__init__.py b/torchtitan/experiments/simple_fsdp/llama3/__init__.py
@@ -13,14 +13,13 @@
 from torchtitan.datasets.hf_datasets import build_hf_dataloader
 from torchtitan.models.llama3 import llama3_configs, pipeline_llama
 from torchtitan.protocols.train_spec import TrainSpec
-
-from .model import SimpleFSDPTransformer
-from .parallelize import parallelize_llama
+from .llama3_model import SimpleFSDPTransformer
+from .llama3_parallelize import parallelize_llama
 
 
 def get_train_spec() -> TrainSpec:
-    return TrainSpec(
-        name="simple_fsdp",
+    return  TrainSpec(
+        name="simple_fsdp.llama3",
         model_cls=SimpleFSDPTransformer,
         model_args=llama3_configs,
         parallelize_fn=parallelize_llama,
diff --git a/torchtitan/experiments/simple_fsdp/llama3/llama3_model.py b/torchtitan/experiments/simple_fsdp/llama3/llama3_model.py
@@ -5,13 +5,13 @@
 # LICENSE file in the root directory of this source tree.
 
 from torchtitan.models.llama3 import Transformer, TransformerModelArgs
-from .simple_fsdp import disable_data_parallel
+from ..simple_fsdp import enable_active_parametrization
 
 
 class SimpleFSDPTransformer(Transformer):
     def __init__(self, model_args: TransformerModelArgs):
         super().__init__(model_args)
 
     def init_weights(self, *args, **kwargs):
-        with disable_data_parallel():
+        with enable_active_parametrization():
             super().init_weights(*args, **kwargs)
diff --git a/torchtitan/experiments/simple_fsdp/llama3/llama3_parallelize.py b/torchtitan/experiments/simple_fsdp/llama3/llama3_parallelize.py
@@ -14,7 +14,7 @@
 from torchtitan.models.llama3.infra.parallelize import apply_tp
 from torchtitan.tools.logging import logger
 
-from .simple_fsdp import data_parallel, MixedPrecisionPolicy
+from ..simple_fsdp import data_parallel, MixedPrecisionPolicy
 
 
 # for selective op activation checkpointing
@@ -116,7 +116,9 @@ def parallelize_llama(
             ac_mode=job_config.activation_checkpoint.mode,
             mp_policy=mp_policy,
         )
-        logger.info("Applied Data Parallel (dp mode=%s) to the model", dp_mode)
+        logger.info(
+            "Applied Data Parallel (simple_fsdp) (dp mode=%s) to the model", dp_mode
+        )
 
     if job_config.compile.enable and "model" in job_config.compile.components:
         torch._inductor.config.reorder_for_peak_memory = False
diff --git a/torchtitan/experiments/simple_fsdp/simple_fsdp.py b/torchtitan/experiments/simple_fsdp/simple_fsdp.py
diff --git a/torchtitan/experiments/simple_fsdp/tests/llama3_integration_tests.py b/torchtitan/experiments/simple_fsdp/tests/llama3_integration_tests.py
