[EP] add support for ETP=1 (#1555)

This is a followup of original EP support
#1324

### PR summary
[TBA] description + figure

### numerics verification
setup
- optimizer Adam
- steps 100, warmup_steps 20
- seed 42

comparison set
- FSDP 2
- FSDP 2, CP 2, TP 2, EP 8, ETP 1
- FSDP 2 (EP 2), PP 2, TP 2 (ETP 2)

<img width="1316" height="392" alt="image"
src="https://github.com/user-attachments/assets/2aa63714-a31e-4152-b904-cc889f8434a1"
/>

Author: tianyu-l

scripts/estimate/estimation.py

@@ -52,6 +52,7 @@ def estimate_memory(job_config: JobConfig):
         tp=parallelism_config.tensor_parallel_degree,
         pp=parallelism_config.pipeline_parallel_degree,
         ep=parallelism_config.expert_parallel_degree,
+        etp=parallelism_config.expert_tensor_parallel_degree,
         world_size=world_size,
     )
     # ParallelDims.build_mesh has to happen outside of the FakeTensorMode

scripts/generate/test_generate.py

@@ -125,6 +125,7 @@ def test_generate(
             tp=world_size,
             pp=1,
             ep=1,
+            etp=1,
             world_size=world_size,
         )
         world_mesh = parallel_dims.world_mesh

tests/unit_tests/test_model_converter.py

@@ -22,6 +22,7 @@ def build_parallel_dims(job_config, world_size):
         tp=parallelism_config.tensor_parallel_degree,
         pp=parallelism_config.pipeline_parallel_degree,
         ep=parallelism_config.expert_parallel_degree,
+        etp=parallelism_config.expert_tensor_parallel_degree,
         world_size=world_size,
     )
     return parallel_dims

torchtitan/config/job_config.py

@@ -374,9 +374,27 @@ class Parallelism:
 
     expert_parallel_degree: int = 1
     """
-    Expert parallelism degree. 1 means disabled.
-    Currently, only "dp2ep" is supported, with the following constraints:
-    context_parallel_degree <= expert_parallel_degree <= data_parallel_shard_degree * context_parallel_degree
+    Expert parallelism degree. 1 means disabled. No effect for non-MoE models.
+    Currently, it is supported with the following constraints:
+    - when etp = tp:
+      - cp <= ep <= dp_shard * cp
+      - ep % cp == 0
+      - dp_shard * cp % ep == 0
+    - when etp = 1:
+      - cp * tp <= ep <= dp_shard * cp * tp
+      - ep % (cp * tp) == 0
+      - dp_shard * cp * tp % ep == 0
+    Note that this is still an experimental feature. Some contrains will be
+    relaxed soon when we have more flexible DeviceMesh support.
+    """
+
+    expert_tensor_parallel_degree: int = 1
+    """
+    Expert tensor parallelism degree. 1 means disabled. No effect for non-MoE models, or when ep = 1.
+    With this option, the tensor parallel degree on routed experts can be different from that on other params.
+    Currently, we only support either
+    - [partial dp -> ep] etp = tp
+    - [partial dp + all tp -> ep] etp = 1
     Note that this is still an experimental feature.
     """
 

torchtitan/distributed/expert_parallel.py

@@ -363,3 +363,52 @@ def wrapper(
         return out
 
     return wrapper
+
+
+# This class is to support Sequence Parallel for ETP=1
+# when EP borrows from all TP and part of DP
+class ReordererSequenceParallel(ParallelStyle):
+    def _prepare_inputput_fn(self, mod, inputs, device_mesh):
+        top_scores, selected_experts_indices = inputs
+
+        top_scores = DTensor.from_local(top_scores, device_mesh, (Replicate(),))
+        selected_experts_indices = DTensor.from_local(
+            selected_experts_indices, device_mesh, (Replicate(),)
+        )
+
+        # TODO: If needed, we can pad tokens in case bs*slen is not divisible by TP degree
+        # if top_scores.shape[0] % device_mesh.size() != 0:
+        #     num_tokens = top_scores.shape[0]
+        #     tp_size = device_mesh.size()
+        #     n_pad = (num_tokens // tp_size + 1) * tp_size - num_tokens
+        #     selected_experts_indices = F.pad(selected_experts_indices, [0, 0, 0, n_pad])
+        #     top_scores = F.pad(top_scores, [0, 0, 0, n_pad])
+        assert top_scores.shape[0] % device_mesh.size() == 0
+
+        # split on the bs*slen dimension
+        top_scores = top_scores.redistribute(device_mesh, (Shard(0),)).to_local()
+        selected_experts_indices = selected_experts_indices.redistribute(
+            device_mesh, (Shard(0),)
+        ).to_local()
+
+        return top_scores, selected_experts_indices
+
+    def _prepare_output_fn(self, mod, outputs, device_mesh):
+        top_scores, token_indices_experts_sorted, num_tokens_per_expert = outputs
+
+        # NOTE: As we shard routed tokens along bs*slen dim across the TP ranks,
+        #       the MoE gather and scatter still require global token indices.
+        num_tokens = top_scores.shape[0]
+        local_rank = device_mesh.get_local_rank()
+        token_indices_experts_sorted += num_tokens // device_mesh.size() * local_rank
+
+        return top_scores, token_indices_experts_sorted, num_tokens_per_expert
+
+    def _apply(self, module: nn.Module, device_mesh: DeviceMesh) -> nn.Module:
+        return distribute_module(
+            module,
+            device_mesh,
+            partition_fn=None,
+            input_fn=self._prepare_inputput_fn,
+            output_fn=self._prepare_output_fn,
+        )

torchtitan/distributed/parallel_dims.py

@@ -23,6 +23,7 @@ class ParallelDims:
     tp: int
     pp: int
     ep: int
+    etp: int
     world_size: int
 
     _world_mesh: DeviceMesh = None
@@ -31,18 +32,19 @@ def __post_init__(self):
         self._validate()
 
     def _validate(self):
-        dp_replicate, dp_shard, cp, tp, pp, ep = (
+        dp_replicate, dp_shard, cp, tp, pp, ep, etp = (
             self.dp_replicate,
             self.dp_shard,
             self.cp,
             self.tp,
             self.pp,
             self.ep,
+            self.etp,
         )
-        for d in (dp_replicate, cp, tp, pp, ep):
+        for d in (dp_replicate, cp, tp, pp, ep, etp):
             assert d >= 1, "Parallelism degree should be >= 1, except for dp_shard"
 
-        assert dp_shard == -1 or dp_shard >= 1, " dp_shard must -1 or >=1."
+        assert dp_shard == -1 or dp_shard >= 1, "dp_shard must -1 or >=1."
         if dp_shard < 0:
             self.dp_shard = dp_shard = self.world_size // (dp_replicate * cp * tp * pp)
         assert dp_shard >= 1
@@ -53,8 +55,13 @@ def _validate(self):
         )
 
         if ep > 1:
-            # EP would borrow all cp and some dp_shard degree
-            assert ep % cp == 0 and (dp_shard * cp) % ep == 0
+            assert etp == tp or etp == 1, "Currently we only support ETP=TP or ETP=1"
+            if etp == tp:
+                # EP would borrow all cp and some dp_shard degree
+                assert ep % cp == 0 and (dp_shard * cp) % ep == 0
+            elif etp == 1:
+                # EP would borrow all cp and tp and some dp_shard degree
+                assert ep % (cp * tp) == 0 and (dp_shard * cp * tp) % ep == 0
 
     def build_mesh(self) -> DeviceMesh:
         # TODO: Current implementation of ParallelDims for dp2ep Expert Parallel
@@ -68,9 +75,15 @@ def build_mesh(self) -> DeviceMesh:
     def _build_mesh_with_ep(self) -> DeviceMesh:
         # With ep, dp_shard and ep are derived submeshes:
         # dp_shard = dp_shard_mod_ep * dp_shard_in_ep
-        # ep = dp_shard_in_ep * cp
-        dp_shard_mod_ep = self.dp_shard * self.cp // self.ep
-        dp_shard_in_ep = self.ep // self.cp
+        if self.etp == self.tp:
+            # ep = dp_shard_in_ep * cp
+            dp_shard_mod_ep = self.dp_shard * self.cp // self.ep
+            dp_shard_in_ep = self.ep // self.cp
+        else:
+            assert self.etp == 1
+            # ep = dp_shard_in_ep * cp * tp
+            dp_shard_mod_ep = self.dp_shard * self.cp * self.tp // self.ep
+            dp_shard_in_ep = self.ep // (self.cp * self.tp)
 
         dims = []
         names = []
@@ -121,6 +134,8 @@ def _build_mesh_with_ep(self) -> DeviceMesh:
             dp_shard_cp_mesh_dim_names.append("cp")
             dp_cp_mesh_dim_names.append("cp")
             ep_mesh_dim_names.append("cp")
+        if self.etp == 1 and self.tp_enabled:
+            ep_mesh_dim_names.append("tp")
 
         mesh[tuple(dp_mesh_dim_names)]._flatten(mesh_dim_name="dp")
         mesh[tuple(dp_shard_cp_mesh_dim_names)]._flatten(mesh_dim_name="dp_shard_cp")
@@ -218,6 +233,10 @@ def pp_enabled(self):
     def ep_enabled(self):
         return self.ep > 1
 
+    @property
+    def etp_enabled(self):
+        return self.etp > 1
+
     @property
     def fsdp_gradient_divide_factor(self) -> int:
         # This is needed for FSDP-sharded experts when Expert Parallel is enabled.

torchtitan/experiments/forge/engine.py

@@ -80,6 +80,7 @@ def __init__(self, job_config: ForgeJobConfig):
             tp=parallelism_config.tensor_parallel_degree,
             pp=parallelism_config.pipeline_parallel_degree,
             ep=parallelism_config.expert_parallel_degree,
+            etp=parallelism_config.expert_tensor_parallel_degree,
             world_size=world_size,
         )
 

torchtitan/experiments/llama4/infra/parallelize.py

@@ -25,6 +25,7 @@
     ExpertParallel,
     ExpertTensorParallel,
     NoParallel,
+    ReordererSequenceParallel,
     TensorParallel,
 )
 
@@ -87,17 +88,19 @@ def parallelize_llama(
             enable_async_tp=job_config.parallelism.enable_async_tensor_parallel,
         )
 
-    # TODO: shall we support tensorwise float8 comms for MoE 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
+                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":
@@ -344,6 +347,7 @@ def apply_moe_ep_tp(
     tp_mesh: DeviceMesh | None,
     ep_mesh: DeviceMesh | None,
     ep_tp_mesh: DeviceMesh | None,
+    etp_enabled: bool,
 ):
     for transformer_block in model.layers.values():
         if not transformer_block.moe_enabled:
@@ -365,13 +369,17 @@ def apply_moe_ep_tp(
                 # input Replicate, output Partial
                 "moe.shared_expert": TensorParallel(),
             }
+            if not etp_enabled:
+                # If TP is borrowed for EP, then split the tokens across TP ranks so that
+                # the reorderer, the all-to-all comms, and routed experts computation
+                # are effectively running Sequence Parallel (split along the folded bs*slen dim)
+                moe_layer_plan.update({"moe.reorderer": ReordererSequenceParallel()})
             parallelize_module(
                 module=transformer_block,
                 device_mesh=tp_mesh,
                 parallelize_plan=moe_layer_plan,
             )
 
-        # if ep_mesh is not None:
         experts_mesh, experts_plan = None, None
         if ep_mesh is None:
             experts_mesh = tp_mesh
@@ -381,9 +389,13 @@ def apply_moe_ep_tp(
             experts_mesh = ep_mesh
             # input / output sharding on the batch / tokens dim
             experts_plan = ExpertParallel()
-        else:
+        elif etp_enabled:
             experts_mesh = ep_tp_mesh
             experts_plan = ExpertTensorParallel(tp_mesh=tp_mesh, ep_mesh=ep_mesh)
+        else:
+            experts_mesh = ep_mesh
+            experts_plan = ExpertParallel()
+
         parallelize_module(
             module=transformer_block.moe.experts,
             device_mesh=experts_mesh,

torchtitan/experiments/llama4/train_configs/debug_model.toml

@@ -53,6 +53,7 @@ enable_async_tensor_parallel = false
 pipeline_parallel_degree = 1
 context_parallel_degree = 1
 expert_parallel_degree = 1
+expert_tensor_parallel_degree = 1
 
 [checkpoint]
 enable_checkpoint = false

torchtitan/experiments/llama4/train_configs/llama4_17bx128e.toml

@@ -46,6 +46,8 @@ pipeline_parallel_degree = 4
 # pipeline_parallel_schedule = "interleaved1f1b"
 # pipeline_parallel_microbatches = 2
 context_parallel_degree = 1
+expert_parallel_degree = 1
+expert_tensor_parallel_degree = 8
 
 [checkpoint]
 enable_checkpoint = false