revise per disscussion round 2, add SAC and optional padding

Author: elfiegg

tests/integration_tests/models.py

@@ -64,37 +64,6 @@ def build_model_tests_list() -> list[OverrideDefinitions]:
             "deepseek_v3_pp+fsdp+tp+ep+etp",
             ngpu=8,
         ),
-        # Integration Test Cases for DeepSeek V3 with DeepEP
-        OverrideDefinitions(
-            [
-                [
-                    "--model.name deepseek_v3",
-                    "--parallelism.data_parallel_shard_degree 4",
-                    "--parallelism.expert_parallel_degree 2",
-                    "--parallelism.moe_comm_backend deep_ep",
-                ],
-            ],
-            "DeepSeek V3 FSDP+EP+DeepEP",
-            "deepseek_v3_fsdp+ep+deepep",
-            ngpu=4,
-        ),
-        OverrideDefinitions(
-            [
-                [
-                    "--model.name deepseek_v3",
-                    "--parallelism.pipeline_parallel_degree 2",
-                    "--parallelism.pipeline_parallel_schedule Interleaved1F1B",
-                    "--parallelism.data_parallel_shard_degree 2",
-                    "--parallelism.tensor_parallel_degree 2",
-                    "--parallelism.expert_parallel_degree 4",
-                    "--parallelism.expert_tensor_parallel_degree 1",
-                    "--parallelism.moe_comm_backend deep_ep",
-                ],
-            ],
-            "DeepSeek V3 PP+FSDP+TP+EP+DeepEP",
-            "deepseek_v3_pp+fsdp+tp+ep+deepep",
-            ngpu=8,
-        ),
         # Integration Test Cases for Qwen3 dense and MoE model
         OverrideDefinitions(
             [
@@ -123,23 +92,6 @@ def build_model_tests_list() -> list[OverrideDefinitions]:
             "qwen3_fsdp+tp+ep+etp",
             ngpu=4,
         ),
-        # Integration Test Cases for Qwen3 with DeepEP
-        OverrideDefinitions(
-            [
-                [
-                    "--model.name qwen3",
-                    "--model.flavor debugmodel_moe",
-                    "--parallelism.data_parallel_shard_degree 2",
-                    "--parallelism.tensor_parallel_degree 2",
-                    "--parallelism.expert_parallel_degree 2",
-                    "--parallelism.expert_tensor_parallel_degree 2",
-                    "--parallelism.moe_comm_backend deep_ep",
-                ],
-            ],
-            "Qwen3 FSDP+TP+EP+ETP+DeepEP",
-            "qwen3_fsdp+tp+ep+etp+deepep",
-            ngpu=4,
-        ),
         # Integration Test Cases for Llama 4
         OverrideDefinitions(
             [
@@ -158,24 +110,6 @@ def build_model_tests_list() -> list[OverrideDefinitions]:
             "llama4_pp+fsdp+tp+ep+compile",
             ngpu=8,
         ),
-        # Integration Test Cases for Llama 4 with DeepEP
-        OverrideDefinitions(
-            [
-                [
-                    "--model.name llama4",
-                    "--parallelism.pipeline_parallel_degree 2",
-                    "--parallelism.pipeline_parallel_schedule Interleaved1F1B",
-                    "--parallelism.data_parallel_shard_degree 2",
-                    "--parallelism.tensor_parallel_degree 2",
-                    "--parallelism.expert_parallel_degree 4",
-                    "--parallelism.expert_tensor_parallel_degree 1",
-                    "--parallelism.moe_comm_backend deep_ep",
-                ],
-            ],
-            "Llama 4 PP+FSDP+TP+EP+DeepEP",
-            "llama4_pp+fsdp+tp+ep+deepep",
-            ngpu=8,
-        ),
     ]
 
     return model_tests

torchtitan/config/job_config.py

@@ -416,17 +416,26 @@ class Parallelism:
     Note that this is still an experimental feature.
     """
 
-    moe_comm_backend: Literal["standard", "deep_ep"] = "standard"
+    expert_parallel_comm_backend: Literal["standard", "deepep"] = "standard"
     """
-    MoE expert-parallel communication backend. No effect for non-MoE models or when ep = 1.
+    Expert-parallel communication backend. No effect for non-MoE models or when ep = 1.
     
     - "standard": Uses PyTorch all-to-all collectives (default)
-    - "deep_ep": Uses DeepEP custom kernels for more efficient communication
+    - "deepep": Uses DeepEP custom kernels for more efficient communication
     
     DeepEP requires installation:
     https://github.com/deepseek-ai/DeepEP.
     """
 
+    deepep_use_alignment_padding: bool = False
+    """
+    Whether to use alignment padding for DeepEP token dispatch.
+    Only applies when expert_parallel_comm_backend="deepep".
+    
+    Recommended for large models (671B+) where the padding overhead is 
+    amortized over more compute. May cause slowdown for smaller models.
+    """
+
 
 @dataclass
 class Checkpoint:

torchtitan/distributed/__init__.py

@@ -13,15 +13,13 @@
 from torch.distributed.tensor.placement_types import Placement
 
 from torchtitan.distributed.parallel_dims import ParallelDims
-from torchtitan.distributed.expert_parallel import ExpertParallelDeepEP
-from torchtitan.distributed.deepep import MoEFlexTokenDispatcher
+from torchtitan.distributed.expert_parallel import DeepEPExpertParallel
 
 
 __all__ = [
     "ParallelDims",
     "NoParallel",
-    "MoEFlexTokenDispatcher",
-    "ExpertParallelDeepEP",
+    "DeepEPExpertParallel",
 ]
 
 

torchtitan/distributed/deepep/__init__.py

@@ -6,9 +6,14 @@
 
 """DeepEP distributed communication primitives for MoE."""
 
-from .flex_dispatcher import MoEFlexTokenDispatcher
+from .deepep import (
+    dispatch_tokens,
+    combine_tokens,
+    DispatchState,
+)
 
 __all__ = [
-    "MoEFlexTokenDispatcher",
+    "dispatch_tokens",
+    "combine_tokens",
+    "DispatchState",
 ]
-