[MoE/EP] apply dim-1 FSDP sharding for routed experts and rewrite shared experts with FFN

Author: tianyu-l

torchtitan/distributed/expert_parallel.py

@@ -320,45 +320,41 @@ def wrapper(
         w2: torch.Tensor,
         w3: torch.Tensor,
         x: torch.Tensor,
-        num_tokens_per_expert: torch.Tensor | None = None,
+        num_tokens_per_expert: torch.Tensor,
     ) -> torch.Tensor:
         global TOKEN_GROUP_ALIGN_SIZE_M
         if isinstance(w1, DTensor):
             w1 = w1.to_local()
             w2 = w2.to_local()
             w3 = w3.to_local()
 
-        if num_tokens_per_expert is not None:
-            from torchtitan.experiments.kernels.moe.indices import (
-                generate_permute_indices,
+        from torchtitan.experiments.kernels.moe.indices import generate_permute_indices
+
+        experts_per_ep_rank = w1.shape[0]
+        num_ep_ranks = num_tokens_per_expert.shape[0] // experts_per_ep_rank
+
+        with torch.no_grad():
+            (
+                permuted_indices,
+                num_tokens_per_expert,
+                _,  # offsets,
+            ) = generate_permute_indices(
+                num_tokens_per_expert,
+                experts_per_ep_rank,
+                num_ep_ranks,
+                x.shape[0] + experts_per_ep_rank * TOKEN_GROUP_ALIGN_SIZE_M,
+                TOKEN_GROUP_ALIGN_SIZE_M,
             )
 
-            experts_per_ep_rank = w1.shape[0]
-            num_ep_ranks = num_tokens_per_expert.shape[0] // experts_per_ep_rank
-
-            with torch.no_grad():
-                (
-                    permuted_indices,
-                    num_tokens_per_expert,
-                    _,  # offsets,
-                ) = generate_permute_indices(
-                    num_tokens_per_expert,
-                    experts_per_ep_rank,
-                    num_ep_ranks,
-                    x.shape[0] + experts_per_ep_rank * TOKEN_GROUP_ALIGN_SIZE_M,
-                    TOKEN_GROUP_ALIGN_SIZE_M,
-                )
-
-            x = torch.vstack((x, x.new_zeros((x.shape[-1]))))
-            input_shape = x.shape
-            x = x[permuted_indices, :]
+        x = torch.vstack((x, x.new_zeros((x.shape[-1]))))
+        input_shape = x.shape
+        x = x[permuted_indices, :]
 
         out = func(w1, w2, w3, x, num_tokens_per_expert)
 
-        if num_tokens_per_expert is not None:
-            out_unpermuted = out.new_empty(input_shape)
-            out_unpermuted[permuted_indices, :] = out
-            out = out_unpermuted[:-1]
+        out_unpermuted = out.new_empty(input_shape)
+        out_unpermuted[permuted_indices, :] = out
+        out = out_unpermuted[:-1]
 
         return out
 

torchtitan/experiments/llama4/infra/parallelize.py

@@ -136,7 +136,7 @@ def parallelize_llama(
             reshard_after_forward_policy=job_config.parallelism.fsdp_reshard_after_forward,
             dp_mod_ep_mesh=(
                 world_mesh[tuple(dp_mod_ep_mesh_dim_names)]
-                if dp_mod_ep_mesh_dim_names
+                if parallel_dims.ep_enabled
                 else None
             ),
             gradient_divide_factor=parallel_dims.fsdp_gradient_divide_factor,
@@ -295,34 +295,43 @@ def apply_fsdp(
     if cpu_offload:
         fsdp_config["offload_policy"] = CPUOffloadPolicy()
 
-    for layer_id, transformer_block in model.layers.items():
-        if reshard_after_forward_policy == "always":
+    match reshard_after_forward_policy:
+        case "always":
             reshard_after_forward = True
-        elif reshard_after_forward_policy == "never":
+        case "never":
             reshard_after_forward = False
-        elif reshard_after_forward_policy == "default":
-            if pp_enabled:
-                # For PP, do not reshard after forward to avoid per-microbatch
-                # all-gathers, which can be expensive and non-overlapped
-                reshard_after_forward = False
-            else:
-                # As an optimization, do not reshard after forward for the last
-                # transformer block since FSDP would prefetch it immediately
-                reshard_after_forward = int(layer_id) < len(model.layers) - 1
-        else:
+        case "default":
+            # For PP, by default do not reshard after forward to avoid per-microbatch
+            # all-gathers, which can be expensive and non-overlapped
+            reshard_after_forward = not pp_enabled
+        case _:
             raise ValueError(
                 f"Invalid reshard_after_forward_policy: {reshard_after_forward_policy}."
             )
 
-        # NOTE: in an MoE layer, the router and the shared experts
-        #       are sharded together with the TransformerBlock
+    if model.tok_embeddings is not None:
+        fully_shard(
+            model.tok_embeddings,
+            **fsdp_config,
+            reshard_after_forward=reshard_after_forward,
+        )
+
+    for layer_id, transformer_block in model.layers.items():
+        # NOTE: When EP is enabled, In an MoE layer, we use the following FSDP wrapping
+        # - the router and the shared experts are sharded together with the TransformerBlock
+        # - the routed experts are sharded with the remaining dp_mod_ep_mesh
         if transformer_block.moe_enabled and dp_mod_ep_mesh:
             fsdp_mod_ep_config = fsdp_config.copy()
             fsdp_mod_ep_config["mesh"] = dp_mod_ep_mesh
             fully_shard(
                 transformer_block.moe.experts,
                 **fsdp_mod_ep_config,
                 reshard_after_forward=reshard_after_forward,
+                # NOTE: When dp_mod_ep * ep > num_experts, FSDP default dim-0 sharding
+                #       causes inefficiency, so we choose to do FSDP sharding on dim-1.
+                # TODO: Even when EP is not used, we may still want to
+                #       shard the experts on non-0 dim.
+                shard_placement_fn=lambda param: Shard(1),
             )
             # NOTE: # Although the FSDP sharding of experts is done on a mesh of
             #       a different size than other parameters, the gradient division
@@ -336,7 +345,17 @@ def apply_fsdp(
             **fsdp_config,
             reshard_after_forward=reshard_after_forward,
         )
-    fully_shard(model, **fsdp_config, reshard_after_forward=not pp_enabled)
+
+    # As an optimization, do not reshard_after_forward the last layers by default
+    # since FSDP would prefetch them immediately after the forward pass
+    if model.norm is not None and model.output is not None:
+        fully_shard(
+            [model.norm, model.output],
+            **fsdp_config,
+            reshard_after_forward=reshard_after_forward_policy == "always",
+        )
+
+    fully_shard(model, **fsdp_config)
 
 
 def apply_moe_ep_tp(
@@ -362,9 +381,18 @@ def apply_moe_ep_tp(
                 ),
                 # replicate computation for the router
                 "moe.router.gate": NoParallel(),
-                # input Replicate, output Partial
-                "moe.shared_expert": TensorParallel(),
             }
+            if transformer_block.moe.shared_experts is not None:
+                # input Replicate, output Partial
+                moe_layer_plan.update(
+                    {
+                        "moe.shared_experts.w1": ColwiseParallel(),
+                        "moe.shared_experts.w2": RowwiseParallel(
+                            output_layouts=Partial()
+                        ),
+                        "moe.shared_experts.w3": ColwiseParallel(),
+                    }
+                )
             parallelize_module(
                 module=transformer_block,
                 device_mesh=tp_mesh,

torchtitan/experiments/llama4/model/args.py

@@ -85,28 +85,28 @@ def get_nparams_and_flops(
     ) -> tuple[int, float]:
         nparams_embedding = 0
         nparams_moe_router = 0
-        nparams_shared_expert = 0
+        nparams_shared_experts = 0
         nparams_experts = 0
         nparams_dense = 0
 
         for name, p in model.named_parameters():
             if "embedding" in name:
                 nparams_embedding += p.numel()
                 nparams_dense += p.numel()
-            elif "moe.shared_expert" in name:
-                nparams_shared_expert += p.numel()
+            elif "moe.shared_experts" in name:
+                nparams_shared_experts += p.numel()
             elif "moe.router" in name:
                 nparams_moe_router += p.numel()
             elif "moe.experts" in name:
                 nparams_experts += p.numel()
             else:
                 nparams_dense += p.numel()
 
-        nparams_sparse = nparams_moe_router + nparams_shared_expert + nparams_experts
+        nparams_sparse = nparams_moe_router + nparams_shared_experts + nparams_experts
         nparams = nparams_dense + nparams_sparse
         nparams_sparse_active = (
             nparams_moe_router
-            + nparams_shared_expert
+            + nparams_shared_experts
             + nparams_experts * self.moe_args.top_k // self.moe_args.num_experts
         )
 

torchtitan/experiments/llama4/scripts/convert_hf_to_dcp_with_gpus.py

@@ -57,11 +57,11 @@ def convert_to_titan_fqns(fqn: str) -> list[str]:
     elif "feed_forward.router.weight" in fqn:
         return [f"layers.{layer}.moe.router.gate.weight"]
     elif "feed_forward.shared_expert.down_proj.weight" in fqn:
-        return [f"layers.{layer}.moe.shared_expert.w2"]
+        return [f"layers.{layer}.moe.shared_experts.w2"]
     elif "feed_forward.shared_expert.gate_proj.weight" in fqn:
-        return [f"layers.{layer}.moe.shared_expert.w3"]
+        return [f"layers.{layer}.moe.shared_experts.w3"]
     elif "feed_forward.shared_expert.up_proj.weight" in fqn:
-        return [f"layers.{layer}.moe.shared_expert.w1"]
+        return [f"layers.{layer}.moe.shared_experts.w1"]
     elif "post_attention_layernorm.weight" in fqn:
         return [f"layers.{layer}.ffn_norm.weight"]
     elif "self_attn.k_proj" in fqn:
@@ -86,7 +86,7 @@ def convert_to_hf_shape(fqn: str, titan_fqns: list[str], dtensor: DTensor) -> li
     elif "shared_expert" in fqn:
         s = dtensor.shape
         # TODO: this is not right but I have to do this to load the checkpoint.
-        return torch.Size((s[2], s[1]))
+        return torch.Size((s[1], s[0]))
     return dtensor.shape
 
 
@@ -96,7 +96,7 @@ def convert_to_titan_tensors(fqn: str, full_tensor: torch.Tensor) -> torch.Tenso
     elif "shared_expert" in fqn:
         # TODO: this is not right but I have to do this to load the checkpoint.
         full_tensor = full_tensor.transpose(1, 0)
-        full_tensors = [full_tensor.unsqueeze(0)]
+        full_tensors = [full_tensor]
     else:
         full_tensors = [full_tensor]
     return full_tensors

torchtitan/models/deepseek_v3/model/args.py

@@ -126,28 +126,28 @@ def get_nparams_and_flops(self, model: nn.Module, seq_len: int) -> tuple[int, in
         """
         nparams_embedding = 0
         nparams_moe_router = 0
-        nparams_shared_expert = 0
+        nparams_shared_experts = 0
         nparams_experts = 0
         nparams_dense = 0
 
         for name, p in model.named_parameters():
             if "embedding" in name:
                 nparams_embedding += p.numel()
                 nparams_dense += p.numel()
-            elif "moe.shared_expert" in name:
-                nparams_shared_expert += p.numel()
+            elif "moe.shared_experts" in name:
+                nparams_shared_experts += p.numel()
             elif "moe.router" in name:
                 nparams_moe_router += p.numel()
             elif "moe.experts" in name:
                 nparams_experts += p.numel()
             else:
                 nparams_dense += p.numel()
 
-        nparams_sparse = nparams_moe_router + nparams_shared_expert + nparams_experts
+        nparams_sparse = nparams_moe_router + nparams_shared_experts + nparams_experts
         nparams = nparams_dense + nparams_sparse
         nparams_sparse_active = (
             nparams_moe_router
-            + nparams_shared_expert
+            + nparams_shared_experts
             + nparams_experts * self.moe_args.top_k // self.moe_args.num_experts
         )
 

torchtitan/models/deepseek_v3/model/model.py

@@ -8,52 +8,15 @@
 from typing import Tuple
 
 import torch
-import torch.nn.functional as F
 from torch import nn
 
 from torchtitan.models.attention import build_attention, init_attention_mask
-from torchtitan.models.moe import MoE
+from torchtitan.models.moe import FeedForward, MoE
 from torchtitan.protocols.train_spec import ModelProtocol
 
 from .args import DeepSeekV3ModelArgs
 
 
-class FeedForward(nn.Module):
-    """
-    FeedForward module
-
-    Args:
-        dim (int): Input dimension.
-        hidden_dim (int): Hidden dimension of the feedforward layer.
-        multiple_of (int): Value to ensure hidden dimension is a multiple of this value.
-        ffn_dim_multiplier (float | None): Custom multiplier for hidden dimension. Defaults to None.
-
-    Attributes:
-        w1 (Linear): Linear transformation for the first layer.
-        w2 (Linear): Linear transformation for the second layer.
-        w3 (Linear): Linear transformation for the third layer.
-
-    """
-
-    def __init__(
-        self,
-        dim: int,
-        hidden_dim: int,
-    ):
-        super().__init__()
-        self.w1 = nn.Linear(dim, hidden_dim, bias=False)
-        self.w2 = nn.Linear(hidden_dim, dim, bias=False)
-        self.w3 = nn.Linear(dim, hidden_dim, bias=False)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return self.w2(F.silu(self.w1(x)) * self.w3(x))
-
-    def init_weights(self, init_std: float = 0.02):
-        nn.init.trunc_normal_(self.w1.weight, mean=0.0, std=0.02)
-        for linear in (self.w2, self.w3):
-            nn.init.trunc_normal_(linear.weight, mean=0.0, std=init_std)
-
-
 # Adapted from https://github.com/DeepSeek-ai/DeepSeek-V3/blob/main/inference/model.py#L294
 def precompute_freqs_cis(args: DeepSeekV3ModelArgs) -> torch.Tensor:
     """

torchtitan/models/deepseek_v3/model/state_dict_adapter.py

@@ -44,9 +44,9 @@ def __init__(self, model_args: DeepSeekV3ModelArgs, hf_assets_path: str | None):
             "model.layers.{}.mlp.experts.{}.up_proj.weight": "layers.{}.moe.experts.w3",
             "model.layers.{}.mlp.experts.{}.down_proj.weight": "layers.{}.moe.experts.w2",
             "model.layers.{}.mlp.gate.weight": "layers.{}.moe.router.gate.weight",
-            "model.layers.{}.mlp.shared_experts.gate_proj.weight": "layers.{}.moe.shared_expert.w1",
-            "model.layers.{}.mlp.shared_experts.up_proj.weight": "layers.{}.moe.shared_expert.w3",
-            "model.layers.{}.mlp.shared_experts.down_proj.weight": "layers.{}.moe.shared_expert.w2",
+            "model.layers.{}.mlp.shared_experts.gate_proj.weight": "layers.{}.moe.shared_experts.w1",
+            "model.layers.{}.mlp.shared_experts.up_proj.weight": "layers.{}.moe.shared_experts.w3",
+            "model.layers.{}.mlp.shared_experts.down_proj.weight": "layers.{}.moe.shared_experts.w2",
             "model.layers.{}.mlp.gate.e_score_correction_bias": "layers.{}.moe.expert_bias",
             "model.norm.weight": "norm.weight",
             "lm_head.weight": "output.weight",
@@ -163,11 +163,6 @@ def to_hf(self, state_dict: dict[str, Any]) -> dict[str, Any]:
                 layer_num = re.search(r"\d+", key).group(0)
                 new_key = to_hf_map[abstract_key]
                 new_key = new_key.format(layer_num)
-
-                # torchtitan shape: (1, s[1], s[2]) -> HF shape: (s[1], s[2])
-                if "shared_expert" in key:
-                    value = value.squeeze(0)
-
                 hf_state_dict[new_key] = value
 
             else:
@@ -217,11 +212,6 @@ def from_hf(self, hf_state_dict: dict[str, Any]) -> dict[str, Any]:
                 layer_num = re.search(r"\d+", key).group(0)
                 new_key = self.from_hf_map[abstract_key]
                 new_key = new_key.format(layer_num)
-
-                # HF shape: (s[1], s[2]) -> torchtitan shape: (1, s[1], s[2])
-                if "shared_experts" in key:
-                    value = value.unsqueeze(0)
-
                 state_dict[new_key] = value
 
             else: