unify moe implementation for llama4 and deepseek_v3

Author: tianyu-l

torchtitan/components/optimizer.py

@@ -24,6 +24,7 @@
 __all__ = [
     "OptimizersContainer",
     "build_optimizers",
+    "build_optimizers_with_moe_load_balancing",
 ]
 
 
@@ -323,3 +324,55 @@ def build_optimizers(
         )
 
     return OptimizersContainer(model_parts, optimizer_cls, optimizer_kwargs)
+
+
+def build_optimizers_with_moe_load_balancing(
+    model_parts: list[nn.Module],
+    optimizer_config: OptimizerConfig,
+    parallel_dims: ParallelDims,
+    ft_manager: FTManager | None = None,
+) -> OptimizersContainer:
+    optimizers = build_optimizers(
+        model_parts=model_parts,
+        optimizer_config=optimizer_config,
+        parallel_dims=parallel_dims,
+        ft_manager=ft_manager,
+    )
+
+    # for MoE auxiliary-loss-free load balancing
+    def _update_expert_bias(
+        model_parts: list[nn.Module],
+        parallel_dims: ParallelDims,
+    ):
+        dp_cp_mesh = (
+            parallel_dims.world_mesh["dp_cp"] if parallel_dims.dp_cp_enabled else None
+        )
+        # TODO: Currently this sync is blocking (thus exposed) and happens on the
+        # default compute stream. Need to assess if this is OK performance-wise.
+        for model_part in model_parts:
+            for transformer_block in model_part.layers.values():
+                if transformer_block.moe_enabled:
+                    moe = transformer_block.moe
+                    if moe.load_balance_coeff is None:
+                        return
+
+                    if dp_cp_mesh is not None:
+                        torch.distributed.all_reduce(
+                            moe.tokens_per_expert, group=dp_cp_mesh.get_group()
+                        )
+
+                    with torch.no_grad():
+                        expert_bias_delta = moe.load_balance_coeff * torch.sign(
+                            moe.tokens_per_expert.mean() - moe.tokens_per_expert
+                        )
+                        expert_bias_delta = expert_bias_delta - expert_bias_delta.mean()
+                        moe.expert_bias.add_(expert_bias_delta)
+                        moe.tokens_per_expert.zero_()
+
+    optimizers.register_step_pre_hook(
+        lambda *args, **kwargs: _update_expert_bias(
+            model_parts, parallel_dims=parallel_dims
+        )
+    )
+
+    return optimizers

torchtitan/components/quantization/float8.py

@@ -10,9 +10,7 @@
 
 from torchtitan.config.job_config import Float8, JobConfig
 from torchtitan.distributed import ParallelDims
-from torchtitan.experiments.llama4.infra.expert_parallel import (
-    set_token_group_alignment_size_m,
-)
+from torchtitan.distributed.expert_parallel import set_token_group_alignment_size_m
 from torchtitan.protocols.model_converter import (
     ModelConverter,
     register_model_converter,

torchtitan/components/quantization/mx.py

@@ -13,6 +13,7 @@
 
 from torchtitan.config.job_config import JobConfig, MX
 from torchtitan.distributed import ParallelDims
+from torchtitan.distributed.expert_parallel import set_token_group_alignment_size_m
 from torchtitan.protocols.model_converter import (
     ModelConverter,
     register_model_converter,
@@ -58,12 +59,8 @@ def __init__(self, job_config: JobConfig, parallel_dims: ParallelDims):
 
         # For MoE training with mxfp8, token group sizes must be multiples of 32
         if job_config.mx.moe_fqns_prototype:
-            from torchtitan.experiments.llama4.infra.expert_parallel import (
-                set_token_group_alignment_size,
-            )
-
             mxfp8_block_size = 32
-            set_token_group_alignment_size(mxfp8_block_size)
+            set_token_group_alignment_size_m(mxfp8_block_size)
             logger.info(f"Setting token group alignment size to {mxfp8_block_size}")
 
         # Configure MXFP8

torchtitan/experiments/llama4/infra/expert_parallel.py renamed to torchtitan/distributed/expert_parallel.py

@@ -11,7 +11,6 @@
 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,
@@ -24,6 +23,41 @@
 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)
+
+
 TOKEN_GROUP_ALIGN_SIZE_M = 8
 ValidTokenGroupAlignmentSize = Literal[8, 16, 32]
 

torchtitan/distributed/parallel_dims.py

@@ -232,8 +232,3 @@ def seq_len_divisor(self):
         # when load balancing is enabled (by default).
         # https://github.com/pytorch/pytorch/blob/4f62dcc/torch/distributed/tensor/experimental/_attention.py#L1246
         return self.tp * (self.cp * 2)
-
-    @cached_property
-    def dense_params_mesh_ndim(self):
-        # Note: In dp2ep EP, EP params mesh ndim is 1 more due to the 'ep' mesh
-        return self.dp_replicate_enabled + self.fsdp_enabled + self.tp_enabled

torchtitan/distributed/pipeline.py renamed to torchtitan/distributed/pipeline_parallel.py

@@ -16,11 +16,9 @@
 from torch import distributed as dist
 from torch.distributed.device_mesh import DeviceMesh
 from torch.distributed.tensor import DTensor
-from torch.nn.attention import SDPBackend
 
 from torchtitan.config import Comm as CommConfig, TORCH_DTYPE_MAP
 from torchtitan.distributed.parallel_dims import ParallelDims
-from torchtitan.models.attention import ScaledDotProductAttention
 from torchtitan.tools.logging import logger
 from torchtitan.tools.utils import device_module, device_type
 
@@ -202,6 +200,10 @@ def context(cp_context: Generator[None, None, None] | None = None):
                 )
 
             if cp_context is not None:
+                from torch.nn.attention import SDPBackend
+
+                from torchtitan.models.attention import ScaledDotProductAttention
+
                 if SDPBackend.MATH in ScaledDotProductAttention.backends:
                     ScaledDotProductAttention.backends.remove(SDPBackend.MATH)
                 assert (
@@ -319,7 +321,7 @@ def clip_grad_norm_(
     error_if_nonfinite: bool = False,
     foreach: bool | None = None,
     pp_mesh: DeviceMesh | None = None,
-    ep_dense_params_mesh_ndim: int | None = None,
+    ep_enabled: bool = False,
 ) -> torch.Tensor:
     """
     Clip the gradient norm of an iterable of parameters.
@@ -349,15 +351,14 @@ def clip_grad_norm_(
         Total norm of the parameter gradients (viewed as a single vector).
 
     """
-    if ep_dense_params_mesh_ndim is not None:
+    if ep_enabled:
         return _clip_grad_norm_with_ep(
             parameters,
             max_norm,
             norm_type,
             error_if_nonfinite,
             foreach,
             pp_mesh,
-            ep_dense_params_mesh_ndim,
         )
 
     if isinstance(parameters, torch.Tensor):
@@ -401,7 +402,6 @@ def _clip_grad_norm_with_ep(
     error_if_nonfinite: bool,
     foreach: bool | None,
     pp_mesh: DeviceMesh | None,
-    dense_params_mesh_ndim: int,
 ) -> torch.Tensor:
     ep_params = []
     non_ep_params = []
@@ -412,12 +412,12 @@ def _clip_grad_norm_with_ep(
         if p.grad is None:
             continue
         assert isinstance(p, DTensor) and isinstance(p.grad, DTensor)
-        if p.device_mesh.ndim == dense_params_mesh_ndim:
-            non_ep_params.append(p)
-            non_ep_grads.append(p.grad)
-        else:
+        if "ep" in p.device_mesh.mesh_dim_names:
             ep_params.append(p)
             ep_grads.append(p.grad)
+        else:
+            non_ep_params.append(p)
+            non_ep_grads.append(p.grad)
     ep_grads_total_norm = torch.nn.utils.get_total_norm(
         ep_grads, norm_type, error_if_nonfinite, foreach
     ).full_tensor()

torchtitan/distributed/utils.py

@@ -231,11 +231,7 @@ def train_step(
             pp_mesh=(
                 parallel_dims.world_mesh["pp"] if parallel_dims.pp_enabled else None
             ),
-            ep_dense_params_mesh_ndim=(
-                parallel_dims.dense_params_mesh_ndim
-                if parallel_dims.ep_enabled
-                else None
-            ),
+            ep_enabled=parallel_dims.ep_enabled,
         )
         self.checkpointer.maybe_wait_for_staging()
         self.optimizers.step()

torchtitan/experiments/forge/example_train.py

@@ -6,15 +6,16 @@
 
 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.llama3 import pipeline_llama
+from torchtitan.models.moe import MoEArgs
 from torchtitan.protocols.train_spec import register_train_spec, TrainSpec
 
 from .infra.parallelize import parallelize_llama
 from .model.args import TransformerModelArgs
 from .model.model import Transformer
-from .optimizer import build_llama4_optimizers
 
 __all__ = [
     "TransformerModelArgs",
@@ -40,7 +41,7 @@
         multiple_of=2048,
         rope_theta=500000,
         max_seq_len=10485760,
-        num_experts=16,
+        moe_args=MoEArgs(num_experts=16),
         interleave_moe_layer_step=1,
     ),
     "17bx128e": TransformerModelArgs(
@@ -51,7 +52,7 @@
         ffn_dim_multiplier=1.2,
         multiple_of=2048,
         rope_theta=500000,
-        num_experts=128,
+        moe_args=MoEArgs(num_experts=128),
     ),
     "debugmodel_irope": TransformerModelArgs(
         dim=256,
@@ -73,7 +74,7 @@
         multiple_of=2048,
         rope_theta=500000,
         max_seq_len=10485760,
-        num_experts=16,
+        moe_args=MoEArgs(num_experts=16),
         interleave_moe_layer_step=1,
         every_n_layers_nope=4,
         use_flex_attn=True,
@@ -87,7 +88,7 @@
         ffn_dim_multiplier=1.2,
         multiple_of=2048,
         rope_theta=500000,
-        num_experts=128,
+        moe_args=MoEArgs(num_experts=128),
         every_n_layers_nope=4,
         use_flex_attn=True,
         attn_mask_type="block_causal",
@@ -102,7 +103,7 @@
         model_args=llama4_configs,
         parallelize_fn=parallelize_llama,
         pipelining_fn=pipeline_llama,
-        build_optimizers_fn=build_llama4_optimizers,
+        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,

torchtitan/experiments/llama4/__init__.py

@@ -21,16 +21,16 @@
 from torchtitan.config import JobConfig, TORCH_DTYPE_MAP
 from torchtitan.distributed import ParallelDims
 
-from torchtitan.models.llama3.infra.parallelize import apply_ac, apply_ddp
-from torchtitan.tools.logging import logger
-
-from .expert_parallel import (
+from torchtitan.distributed.expert_parallel import (
     ExpertParallel,
     ExpertTensorParallel,
     NoParallel,
     TensorParallel,
 )
 
+from torchtitan.models.llama3.infra.parallelize import apply_ac, apply_ddp
+from torchtitan.tools.logging import logger
+
 
 def parallelize_llama(
     model: nn.Module,