add support for simplefsdp+ep

Author: ruisizhang123

torchtitan/distributed/expert_parallel.py

@@ -11,6 +11,7 @@
 import torch
 import torch.distributed as dist
 import torch.nn as nn
+from torch.distributed._functional_collectives import all_to_all_single_autograd
 from torch.distributed.tensor import (
     DeviceMesh,
     distribute_module,
@@ -22,40 +23,39 @@
 from torch.distributed.tensor.parallel import ParallelStyle
 from torch.distributed.tensor.placement_types import Placement
 
-
 # from torch.distributed._functional_collectives import all_to_all_single_autograd
 # TODO: there is memory leak issue with AC + all_to_all_single_autograd
-# This is a temporary fix by @rakkit https://github.com/pytorch/torchtitan/issues/1467
-class _A2A(torch.autograd.Function):
-    @staticmethod
-    def forward(ctx, x, out_splits, in_splits, group):
-        if isinstance(out_splits, torch.Tensor):
-            out_splits = out_splits.tolist()
-        if isinstance(in_splits, torch.Tensor):
-            in_splits = in_splits.tolist()
-        T_out = int(sum(out_splits))
-
-        y = x.new_empty((T_out,) + tuple(x.shape[1:]))  # allocate by output splits
-        dist.all_to_all_single(y, x.contiguous(), out_splits, in_splits, group=group)
-
-        ctx.in_splits = in_splits
-        ctx.out_splits = out_splits
-        ctx.group = group
-        return y
-
-    @staticmethod
-    def backward(ctx, grad_y):
-        # grad wrt input has length sum(in_splits)
-        T_in = int(sum(ctx.in_splits))
-        grad_x = grad_y.new_empty((T_in,) + tuple(grad_y.shape[1:]))
-        dist.all_to_all_single(
-            grad_x, grad_y.contiguous(), ctx.in_splits, ctx.out_splits, group=ctx.group
-        )
-        return grad_x, None, None, None
-
-
-def all_to_all_single_autograd(x, out_splits, in_splits, group):
-    return _A2A.apply(x, out_splits, in_splits, group)
+# # This is a temporary fix by @rakkit https://github.com/pytorch/torchtitan/issues/1467
+# class _A2A(torch.autograd.Function):
+#     @staticmethod
+#     def forward(ctx, x, out_splits, in_splits, group):
+#         if isinstance(out_splits, torch.Tensor):
+#             out_splits = out_splits.tolist()
+#         if isinstance(in_splits, torch.Tensor):
+#             in_splits = in_splits.tolist()
+#         T_out = int(sum(out_splits))
+
+#         y = x.new_empty((T_out,) + tuple(x.shape[1:]))  # allocate by output splits
+#         dist.all_to_all_single(y, x.contiguous(), out_splits, in_splits, group=group)
+
+#         ctx.in_splits = in_splits
+#         ctx.out_splits = out_splits
+#         ctx.group = group
+#         return y
+
+#     @staticmethod
+#     def backward(ctx, grad_y):
+#         # grad wrt input has length sum(in_splits)
+#         T_in = int(sum(ctx.in_splits))
+#         grad_x = grad_y.new_empty((T_in,) + tuple(grad_y.shape[1:]))
+#         dist.all_to_all_single(
+#             grad_x, grad_y.contiguous(), ctx.in_splits, ctx.out_splits, group=ctx.group
+#         )
+#         return grad_x, None, None, None
+
+
+# def all_to_all_single_autograd(x, out_splits, in_splits, group):
+#     return _A2A.apply(x, out_splits, in_splits, group)
 
 
 TOKEN_GROUP_ALIGN_SIZE_M = 8
@@ -188,14 +188,12 @@ def _token_dispatch(self, mod, inputs, device_mesh):
                 group=device_mesh.get_group(),
             )
             # NOTE: this would incur a device-to-host sync
-            self.input_splits = (
-                num_tokens_per_expert.view(device_mesh.shape[0], -1).sum(dim=1).tolist()
-            )
-            self.output_splits = (
-                num_tokens_per_expert_group.view(device_mesh.shape[0], -1)
-                .sum(dim=1)
-                .tolist()
-            )
+            self.input_splits = num_tokens_per_expert.view(
+                device_mesh.shape[0], -1
+            ).sum(dim=1)
+            self.output_splits = num_tokens_per_expert_group.view(
+                device_mesh.shape[0], -1
+            ).sum(dim=1)
 
         # perform all-to-all
         routed_input = all_to_all_single_autograd(

torchtitan/experiments/simple_fsdp/README.md

@@ -12,10 +12,18 @@ This folder includes an experimental frontend implementation for [SimpleFSDP: Si
 
 ### Enable SimpleFSDP Training
 
+#### Training Llama3 models
+
 ```bash
 CONFIG_FILE="./torchtitan/models/llama3/train_configs/llama3_8b.toml" ./run_train.sh --model.name llama3_simple_fsdp --training.compile
 ```
 
+#### Training DeepSeek_v3 models
+
+```bash
+CONFIG_FILE="./torchtitan/models/deepseek_v3/train_configs/debug_model.toml" ./run_train.sh --model.name deepseekv3_simple_fsdp --training.compile
+```
+
 ### Composability Support
 
 Some of the features require the updates from PyTorch, with which we are working on providing composability support for the following features:
@@ -28,6 +36,7 @@ Some of the features require the updates from PyTorch, with which we are working
 |Tensor Parallelism| ✅ |
 |Context Parallelism| ✅ |
 |Pipeline Parallelism| ✅ |
+|Expert Parallelism| ✅ |
 |Distributed Checkpointing| ✅ |
 |Float8 Training| 🚧 |
 

torchtitan/experiments/simple_fsdp/__init__.py

@@ -8,14 +8,20 @@
 
 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
+from torchtitan.components.optimizer import (
+    build_optimizers,
+    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 llama3_configs, pipeline_llama
 from torchtitan.protocols.train_spec import register_train_spec, TrainSpec
+from .deepseek_v3_model import SimpleFSDPDeepSeekV3Model
+from .deepseek_v3_parallelize import parallelize_deepseekv3
 
-from .model import SimpleFSDPTransformer
-from .parallelize import parallelize_llama
+from .llama3_model import SimpleFSDPTransformer
+from .llama3_parallelize import parallelize_llama
 
 register_train_spec(
     TrainSpec(
@@ -31,3 +37,19 @@
         build_loss_fn=build_cross_entropy_loss,
     )
 )
+
+
+register_train_spec(
+    TrainSpec(
+        name="deepseekv3_simple_fsdp",
+        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,
+    )
+)

torchtitan/experiments/simple_fsdp/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 disable_data_parallel
+
+
+class SimpleFSDPDeepSeekV3Model(DeepSeekV3Model):
+    def __init__(self, model_args: DeepSeekV3ModelArgs):
+        super().__init__(model_args)
+        self.init_weights()
+
+    def init_weights(self, *args, **kwargs):
+        with disable_data_parallel():
+            super().init_weights(*args, **kwargs)

torchtitan/experiments/simple_fsdp/deepseek_v3_parallelize.py

@@ -0,0 +1,153 @@
+# 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.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}).
+        """
+
+    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,
+            enable_async_tp=False,
+        )
+
+    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
+                else None
+            ),
+        )
+
+    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 = []
+        ep_modules = []
+        ep_shared_experts = []
+        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")
+        for _, transformer_block in model.layers.items():
+            if transformer_block.moe_enabled:
+                ep_modules.append(transformer_block.moe.experts)
+                ep_shared_experts.append(transformer_block.moe.shared_expert)
+
+        if parallel_dims.tp_enabled and parallel_dims.ep_enabled:
+            tp_ep_mesh = world_mesh["ep", "tp"]
+        else:
+            tp_ep_mesh = None
+
+        model = data_parallel(
+            model,
+            dp_mesh,
+            dp_mode,
+            ac_mode=job_config.activation_checkpoint.mode,
+            mp_policy=mp_policy,
+            tp_mesh=world_mesh["tp"] if parallel_dims.tp_enabled else None,
+            tp_ep_mesh=tp_ep_mesh,
+            dp_mod_ep_mesh=world_mesh[tuple(dp_mod_ep_mesh_dim_names)]
+            if parallel_dims.ep_enabled
+            else None,
+            ep_modules=ep_modules,
+            ep_shared_experts=ep_shared_experts,
+        )
+        if parallel_dims.dp_replicate_enabled:
+            logger.info("Applied HSDP to the model")
+        else:
+            logger.info("Applied FSDP to the model")
+
+    if job_config.training.compile:
+        torch._inductor.config.reorder_for_peak_memory = False
+        torch._dynamo.config.capture_scalar_outputs = True
+        model = torch.compile(model, fullgraph=True)
+
+    return model