Update on "[MoE][2/n]Move EP setup from trainer to config registry and add model_registry params"

acisseJZhong · acisseJZhong · commit 515ff1f8fd62 · 2026-04-14T17:25:24.000-07:00
## Summary

Remove the `apply_ep()` call from `Trainer.__post_init__` and move EP-aware token dispatcher configuration to model config time.

Previously, model configs always created a `LocalTokenDispatcher.Config` (EP=1 default), then `Trainer.__post_init__` would immediately replace it with the correct EP dispatcher via `apply_ep()`. This two-phase setup split EP concerns between the model config layer and the trainer, making it harder to reason about what dispatcher a model actually uses.

Now:
- `model_registry()` accepts `attn_backend` and `moe_comm_backend` params across all model families (llama3, llama4, deepseek_v3, gpt_oss, qwen3), so config registries can build the right dispatcher config upfront
- `apply_ep()` and `find_pad_multiple` imports are removed from `trainer.py`
- `ep_degree` is moved from `LocalTokenDispatcher.Config` (base) to `AllToAllTokenDispatcher.Config` and `DeepEPTokenDispatcher.Config` where it's actually used
- Stale comment referencing removed `ExpertSequenceParallel` class is updated

## Test Plan

Existing unit tests: `python -m pytest tests/unit_tests/test_expert_parallel.py -x`


[ghstack-poisoned]
diff --git a/torchtitan/models/common/config_utils.py b/torchtitan/models/common/config_utils.py
@@ -242,34 +242,6 @@ def make_token_dispatcher_config(
         )
 
 
-def apply_ep(
-    layers: list,
-    *,
-    ep_size: int,
-    sp_size: int = 1,
-    moe_comm_backend: str = "standard",
-    hybridep_non_blocking_expert_capacity_factor: float | None = None,
-    pad_multiple: int | None = None,
-) -> None:
-    """Replace token dispatchers in MoE layers for expert parallelism.
-
-    Mutates layer configs in-place: for each MoE layer, replaces the
-    token_dispatcher with the appropriate config based on EP settings.
-    """
-    for layer_cfg in layers:
-        if layer_cfg.moe is not None:
-            td = layer_cfg.moe.experts.token_dispatcher
-            layer_cfg.moe.experts.token_dispatcher = make_token_dispatcher_config(
-                num_experts=td.num_experts,
-                top_k=td.top_k,
-                score_before_experts=td.score_before_experts,
-                ep_size=ep_size,
-                sp_size=sp_size,
-                moe_comm_backend=moe_comm_backend,
-                hybridep_non_blocking_expert_capacity_factor=hybridep_non_blocking_expert_capacity_factor,
-                pad_multiple=pad_multiple,
-            )
-
 
 def make_experts_config(
     *,
diff --git a/torchtitan/models/deepseek_v3/__init__.py b/torchtitan/models/deepseek_v3/__init__.py
@@ -10,6 +10,7 @@
 
 import torch.nn as nn
 from torchtitan.components.loss import build_cross_entropy_loss
+from torchtitan.config import ParallelismConfig
 from torchtitan.components.optimizer import register_moe_load_balancing_hook
 from torchtitan.distributed.pipeline_parallel import pipeline_llm
 from torchtitan.models.common import Embedding, Linear, RMSNorm, RoPE, TransformerBlock
@@ -170,7 +171,8 @@ def _build_dsv3_layers(
     router_route_norm: bool = False,
     score_before_experts: bool = False,
     attn_backend: str = "sdpa",
-    moe_comm_backend: str = "standard",
+    parallelism: ParallelismConfig | None = None,
+    pad_multiple: int | None = None,
 ) -> list[TransformerBlock.Config]:
     """Build the list of per-layer TransformerBlock configs.
 
@@ -180,6 +182,8 @@ def _build_dsv3_layers(
     Router and expert inits are constructed per-layer so depth-scaled
     initializers are correct for each layer's position.
     """
+    if parallelism is None:
+        parallelism = ParallelismConfig()
     layers = []
     for layer_id in range(n_layers):
         attn_cfg = _make_dsv3_attn_config(
@@ -225,7 +229,10 @@ def _build_dsv3_layers(
                     top_k=router_top_k,
                     param_init=_depth_experts_init(layer_id),
                     score_before_experts=score_before_experts,
-                    moe_comm_backend=moe_comm_backend,
+                    moe_comm_backend=parallelism.expert_parallel_comm_backend,
+                    ep_size=parallelism.expert_parallel_degree,
+                    pad_multiple=pad_multiple,
+                    hybridep_non_blocking_expert_capacity_factor=parallelism.hybridep_non_blocking_expert_capacity_factor,
                 ),
                 shared_experts=make_ffn_config(
                     dim=dim,
@@ -251,7 +258,8 @@ def _build_dsv3_layers(
 
 def _debugmodel(
     attn_backend: str = "sdpa",
-    moe_comm_backend: str = "standard",
+    parallelism: ParallelismConfig | None = None,
+    pad_multiple: int | None = None,
     **kwargs,
 ) -> DeepSeekV3Model.Config:
     dim = 256
@@ -284,7 +292,8 @@ def _debugmodel(
         router_score_func="softmax",
         score_before_experts=False,
         attn_backend=attn_backend,
-        moe_comm_backend=moe_comm_backend,
+        parallelism=parallelism,
+        pad_multiple=pad_multiple,
     )
     return DeepSeekV3Model.Config(
         vocab_size=vocab_size,
@@ -315,7 +324,8 @@ def _debugmodel(
 
 def _16b(
     attn_backend: str = "flex",
-    moe_comm_backend: str = "standard",
+    parallelism: ParallelismConfig | None = None,
+    pad_multiple: int | None = None,
     **kwargs,
 ) -> DeepSeekV3Model.Config:
     dim = 2048
@@ -348,7 +358,8 @@ def _16b(
         router_score_func="softmax",
         score_before_experts=False,
         attn_backend=attn_backend,
-        moe_comm_backend=moe_comm_backend,
+        parallelism=parallelism,
+        pad_multiple=pad_multiple,
     )
     return DeepSeekV3Model.Config(
         vocab_size=vocab_size,
@@ -379,7 +390,8 @@ def _16b(
 
 def _236b(
     attn_backend: str = "flex",
-    moe_comm_backend: str = "standard",
+    parallelism: ParallelismConfig | None = None,
+    pad_multiple: int | None = None,
     **kwargs,
 ) -> DeepSeekV3Model.Config:
     dim = 5120
@@ -416,7 +428,8 @@ def _236b(
         router_route_scale=16.0,
         score_before_experts=False,
         attn_backend=attn_backend,
-        moe_comm_backend=moe_comm_backend,
+        parallelism=parallelism,
+        pad_multiple=pad_multiple,
     )
     return DeepSeekV3Model.Config(
         vocab_size=vocab_size,
@@ -447,7 +460,8 @@ def _236b(
 
 def _671b(
     attn_backend: str = "flex",
-    moe_comm_backend: str = "standard",
+    parallelism: ParallelismConfig | None = None,
+    pad_multiple: int | None = None,
     **kwargs,
 ) -> DeepSeekV3Model.Config:
     dim = 7168
@@ -485,7 +499,8 @@ def _671b(
         router_route_norm=True,
         score_before_experts=False,
         attn_backend=attn_backend,
-        moe_comm_backend=moe_comm_backend,
+        parallelism=parallelism,
+        pad_multiple=pad_multiple,
     )
     return DeepSeekV3Model.Config(
         vocab_size=vocab_size,
@@ -525,11 +540,13 @@ def _671b(
 def model_registry(
     flavor: str,
     attn_backend: str = "sdpa",
-    moe_comm_backend: str = "standard",
+    parallelism: ParallelismConfig | None = None,
+    pad_multiple: int | None = None,
 ) -> ModelSpec:
     config = deepseekv3_configs[flavor](
         attn_backend=attn_backend,
-        moe_comm_backend=moe_comm_backend,
+        parallelism=parallelism,
+        pad_multiple=pad_multiple,
     )
     return ModelSpec(
         name="deepseek_v3",
diff --git a/torchtitan/models/deepseek_v3/config_registry.py b/torchtitan/models/deepseek_v3/config_registry.py
@@ -8,6 +8,7 @@
 from torchtitan.components.lr_scheduler import LRSchedulersContainer
 from torchtitan.components.metrics import MetricsProcessor
 from torchtitan.components.optimizer import OptimizersContainer
+from torchtitan.components.quantization import find_pad_multiple
 from torchtitan.components.quantization.float8 import (
     Float8GroupedMMConverter,
     Float8LinearConverter,
@@ -64,9 +65,14 @@ def deepseek_v3_debugmodel_flex_attn() -> Trainer.Config:
 
 
 def deepseek_v3_16b() -> Trainer.Config:
+    parallelism = ParallelismConfig(
+        pipeline_parallel_schedule="Interleaved1F1B",
+        expert_parallel_degree=8,
+        expert_tensor_parallel_degree=1,
+    )
     return Trainer.Config(
         hf_assets_path="./assets/hf/deepseek-moe-16b-base",
-        model_spec=model_registry("16B"),
+        model_spec=model_registry("16B", parallelism=parallelism),
         dataloader=HuggingFaceTextDataLoader.Config(
             dataset="c4",
         ),
@@ -81,11 +87,7 @@ def deepseek_v3_16b() -> Trainer.Config:
             seq_len=4096,
             steps=1000,
         ),
-        parallelism=ParallelismConfig(
-            pipeline_parallel_schedule="Interleaved1F1B",
-            expert_parallel_degree=8,
-            expert_tensor_parallel_degree=1,
-        ),
+        parallelism=parallelism,
         checkpoint=CheckpointManager.Config(interval=10),
         activation_checkpoint=ActivationCheckpointConfig(
             mode="selective",
@@ -95,9 +97,15 @@ def deepseek_v3_16b() -> Trainer.Config:
 
 
 def deepseek_v3_671b() -> Trainer.Config:
+    converters = [
+        Float8LinearConverter.Config(filter_fqns=["output", "router.gate"]),
+        Float8GroupedMMConverter.Config(fqns=["experts"]),
+    ]
     return Trainer.Config(
         hf_assets_path="./assets/hf/DeepSeek-V3.1-Base",
-        model_spec=model_registry("671B"),
+        model_spec=model_registry(
+            "671B", pad_multiple=find_pad_multiple(converters)
+        ),
         dataloader=HuggingFaceTextDataLoader.Config(
             dataset="c4",
         ),
@@ -123,10 +131,5 @@ def deepseek_v3_671b() -> Trainer.Config:
             mode="selective",
         ),
         compile=CompileConfig(enable=True, components=["loss"]),
-        model_converters=ModelConvertersContainer.Config(
-            converters=[
-                Float8LinearConverter.Config(filter_fqns=["output", "router.gate"]),
-                Float8GroupedMMConverter.Config(fqns=["experts"]),
-            ],
-        ),
+        model_converters=ModelConvertersContainer.Config(converters=converters),
     )
diff --git a/torchtitan/models/llama4/__init__.py b/torchtitan/models/llama4/__init__.py
@@ -9,6 +9,7 @@
 
 import torch.nn as nn
 from torchtitan.components.loss import build_cross_entropy_loss
+from torchtitan.config import ParallelismConfig
 from torchtitan.components.optimizer import register_moe_load_balancing_hook
 from torchtitan.distributed.pipeline_parallel import pipeline_llm
 from torchtitan.models.common import (
@@ -84,14 +85,17 @@ def _build_llama4_layers(
     interleave_moe_layer_step: int = 1,
     fixed_attn_block_size: int = 8192,
     attn_backend: str = "flex",
-    moe_comm_backend: str = "standard",
     shared_experts_hidden_dim: int | None = None,
+    parallelism: ParallelismConfig | None = None,
+    pad_multiple: int | None = None,
 ) -> list[TransformerBlock.Config]:
     """Build per-layer configs for a Llama4 model.
 
     Handles iRoPE (NoPE on every N layers) and MoE interleaving. For each
     layer, depth-scaled inits are computed using the layer index.
     """
+    if parallelism is None:
+        parallelism = ParallelismConfig()
     inner_attention, mask_type = get_attention_config(attn_backend)
     if every_n_layers_nope <= 1:
         raise ValueError("every_n_layers_nope must be greater than 1")
@@ -131,7 +135,10 @@ def _build_llama4_layers(
                 num_experts=num_experts,
                 top_k=router.top_k,
                 param_init=_depth_experts_init(layer_id),
-                moe_comm_backend=moe_comm_backend,
+                moe_comm_backend=parallelism.expert_parallel_comm_backend,
+                ep_size=parallelism.expert_parallel_degree,
+                pad_multiple=pad_multiple,
+                hybridep_non_blocking_expert_capacity_factor=parallelism.hybridep_non_blocking_expert_capacity_factor,
             )
             shared_experts = make_ffn_config(
                 dim=dim,
@@ -176,7 +183,8 @@ def _build_llama4_layers(
 
 def _debugmodel(
     attn_backend: str = "flex",
-    moe_comm_backend: str = "standard",
+    parallelism: ParallelismConfig | None = None,
+    pad_multiple: int | None = None,
     **kwargs,
 ) -> Llama4Model.Config:
     dim = 256
@@ -207,7 +215,8 @@ def _debugmodel(
             interleave_moe_layer_step=2,
             fixed_attn_block_size=256,
             attn_backend=attn_backend,
-            moe_comm_backend=moe_comm_backend,
+            parallelism=parallelism,
+            pad_multiple=pad_multiple,
         ),
         rope=RoPE.Config(
             dim=dim // n_heads,
@@ -223,7 +232,8 @@ def _debugmodel(
 
 def _17bx16e(
     attn_backend: str = "flex",
-    moe_comm_backend: str = "standard",
+    parallelism: ParallelismConfig | None = None,
+    pad_multiple: int | None = None,
     **kwargs,
 ) -> Llama4Model.Config:
     dim = 5120
@@ -265,7 +275,8 @@ def _17bx16e(
             every_n_layers_nope=4,
             interleave_moe_layer_step=1,
             attn_backend=attn_backend,
-            moe_comm_backend=moe_comm_backend,
+            parallelism=parallelism,
+            pad_multiple=pad_multiple,
         ),
         rope=RoPE.Config(
             dim=dim // n_heads,
@@ -281,7 +292,8 @@ def _17bx16e(
 
 def _17bx128e(
     attn_backend: str = "flex",
-    moe_comm_backend: str = "standard",
+    parallelism: ParallelismConfig | None = None,
+    pad_multiple: int | None = None,
     **kwargs,
 ) -> Llama4Model.Config:
     dim = 5120
@@ -323,7 +335,8 @@ def _17bx128e(
             every_n_layers_nope=4,
             interleave_moe_layer_step=1,
             attn_backend=attn_backend,
-            moe_comm_backend=moe_comm_backend,
+            parallelism=parallelism,
+            pad_multiple=pad_multiple,
         ),
         rope=RoPE.Config(
             dim=dim // n_heads,
@@ -345,11 +358,13 @@ def _17bx128e(
 def model_registry(
     flavor: str,
     attn_backend: str = "sdpa",
-    moe_comm_backend: str = "standard",
+    parallelism: ParallelismConfig | None = None,
+    pad_multiple: int | None = None,
 ) -> ModelSpec:
     config = llama4_configs[flavor](
         attn_backend=attn_backend,
-        moe_comm_backend=moe_comm_backend,
+        parallelism=parallelism,
+        pad_multiple=pad_multiple,
     )
     return ModelSpec(
         name="llama4",
diff --git a/torchtitan/models/llama4/config_registry.py b/torchtitan/models/llama4/config_registry.py
@@ -8,6 +8,7 @@
 from torchtitan.components.lr_scheduler import LRSchedulersContainer
 from torchtitan.components.metrics import MetricsProcessor
 from torchtitan.components.optimizer import OptimizersContainer
+from torchtitan.components.quantization import find_pad_multiple
 from torchtitan.components.quantization.float8 import Float8GroupedMMConverter
 from torchtitan.config import (
     ActivationCheckpointConfig,
@@ -57,10 +58,13 @@ def llama4_debugmodel() -> Trainer.Config:
 
 
 def llama4_debugmodel_fp8() -> Trainer.Config:
+    converters = [Float8GroupedMMConverter.Config(fqns=["experts"])]
     return Trainer.Config(
         hf_assets_path="./tests/assets/tokenizer",
         metrics=MetricsProcessor.Config(log_freq=1),
-        model_spec=model_registry("debugmodel"),
+        model_spec=model_registry(
+            "debugmodel", pad_multiple=find_pad_multiple(converters)
+        ),
         dataloader=HuggingFaceTextDataLoader.Config(
             dataset="c4_test",
         ),
@@ -88,11 +92,7 @@ def llama4_debugmodel_fp8() -> Trainer.Config:
             mode="selective",
         ),
         compile=CompileConfig(enable=True, components=["model", "loss"]),
-        model_converters=ModelConvertersContainer.Config(
-            converters=[
-                Float8GroupedMMConverter.Config(fqns=["experts"]),
-            ],
-        ),
+        model_converters=ModelConvertersContainer.Config(converters=converters),
     )
 
 
diff --git a/torchtitan/models/qwen3/__init__.py b/torchtitan/models/qwen3/__init__.py
