address comments

Author: elfiegg

torchtitan/distributed/utils.py

@@ -341,21 +341,12 @@ def _get_distributed_backend(enable_cpu_backend):
         prefix = comm_config.save_traces_file_prefix
         os.makedirs(dump_dir, exist_ok=True)
         _warn_overwrite_env(TRACE_FILE, f"{dump_dir}/{prefix}")
-
-    # _ranks argument is only available in newer PyTorch versions
-    init_kwargs = {
-        "backend": _get_distributed_backend(enable_cpu_backend),
-        "timeout": timedelta(seconds=comm_config.init_timeout_seconds),
-    }
-    # Try with _ranks first (newer PyTorch), fall back without it
-    try:
-        torch.distributed.init_process_group(
-            **init_kwargs,
-            _ranks=ranks if ranks is not None else [],
-        )
-    except TypeError:
-        # Older PyTorch doesn't support _ranks
-        torch.distributed.init_process_group(**init_kwargs)
+ 
+    torch.distributed.init_process_group(
+        backend=_get_distributed_backend(enable_cpu_backend),
+        timeout=timedelta(seconds=comm_config.init_timeout_seconds),
+        _ranks=ranks if ranks is not None else [],
+    )
 
     return torch.distributed.get_world_size()
 

torchtitan/models/deepseek_v3/infra/parallelize.py

@@ -44,21 +44,8 @@
     # used to compute the scaling factor for quantization.
     torch.ops.aten.max.default,
     torch._higher_order_ops.flex_attention,
+    torch._higher_order_ops.inductor_compiled_code,
 }
-# Add optional ops if available (requires newer PyTorch)
-try:
-    _op_sac_save_list.add(torch._higher_order_ops.inductor_compiled_code)
-except AttributeError:
-    pass
-
-# Add DeepEP custom ops to SAC save list
-try:
-    import torchtitan.distributed.deepep.deepep  # noqa: F401
-    _op_sac_save_list.add(torch.ops.deepep.dispatch.default)
-    _op_sac_save_list.add(torch.ops.deepep.combine.default)
-except (ImportError, AttributeError):
-    pass
-
 
 # Adapted from llama4/infra/parallelize.py
 def parallelize_deepseekv3(
@@ -115,11 +102,17 @@ def parallelize_deepseekv3(
             job_config.parallelism.expert_parallel_comm_backend == "deepep" 
             and not parallel_dims.ep_enabled
         ):
+            use_deepep = False
             logger.warning(
                 "expert_parallel_comm_backend='deepep' has no effect when EP=1. "
                 "Using standard communication."
             )
 
+        if use_deepep:
+            import torchtitan.distributed.deepep.deepep  # noqa: F401
+            _op_sac_save_list.add(torch.ops.deepep.dispatch.default)
+            _op_sac_save_list.add(torch.ops.deepep.combine.default)
+
         # DeepEP + ETP is not supported yet
         if use_deepep and parallel_dims.etp_enabled:
             raise NotImplementedError(

torchtitan/models/deepseek_v3/model/args.py

@@ -112,7 +112,7 @@ def update_from_config(self, job_config: JobConfig, **kwargs) -> None:
         )
 
         # Configure expert parallel communication backend from config (defaults to "standard")
-        self.expert_parallel_comm_backend = job_config.parallelism.expert_parallel_comm_backend
+        self.moe_impl = job_config.parallelism.expert_parallel_comm_backend
 
     def get_nparams_and_flops(
         self, model: nn.Module, seq_len: int

torchtitan/models/deepseek_v3/model/model.py

@@ -19,7 +19,7 @@
     get_document_mask_mod,
     ScaledDotProductAttentionWrapper,
 )
-from torchtitan.models.moe import FeedForward, MoE
+from torchtitan.models.moe import FeedForward, MoE, build_moe
 from torchtitan.protocols.model import AttentionMasksType
 from torchtitan.protocols.train_spec import ModelProtocol
 
@@ -350,13 +350,11 @@ def __init__(self, layer_id: int, model_args: DeepSeekV3ModelArgs):
 
         self.moe_enabled = layer_id >= model_args.n_dense_layers
         if self.moe_enabled:
-            # Use build_moe factory to support different communication backends
-            from torchtitan.models.moe import build_moe
             self.moe = build_moe(
                 args=model_args.moe_args,
                 dim=model_args.dim,
                 hidden_dim=model_args.moe_inter_dim,
-                communication_backend=model_args.expert_parallel_comm_backend,
+                moe_impl=model_args.moe_impl,
             )
         else:
             self.feed_forward = FeedForward(model_args.dim, model_args.inter_dim)

torchtitan/models/llama4/infra/parallelize.py

@@ -55,14 +55,6 @@
     torch._higher_order_ops.inductor_compiled_code,
 }
 
-# Add DeepEP custom ops to SAC save list
-try:
-    import torchtitan.distributed.deepep.deepep  # noqa: F401
-    _op_sac_save_list.add(torch.ops.deepep.dispatch.default)
-    _op_sac_save_list.add(torch.ops.deepep.combine.default)
-except (ImportError, AttributeError):
-    pass
-
 
 def parallelize_llama(
     model: nn.Module,
@@ -117,11 +109,19 @@ def parallelize_llama(
         job_config.parallelism.expert_parallel_comm_backend == "deepep" 
         and not parallel_dims.ep_enabled
     ):
+        use_deepep = False
         logger.warning(
             "expert_parallel_comm_backend='deepep' has no effect when EP=1. "
             "Using standard communication."
         )
 
+    if use_deepep:
+        # Import deepep module to register custom ops before accessing them
+        import torchtitan.distributed.deepep  # noqa: F401 - registers torch.ops.deepep
+        _op_sac_save_list.add(torch.ops.deepep.get_dispatch_layout.default)
+        _op_sac_save_list.add(torch.ops.deepep.dispatch.default)
+        _op_sac_save_list.add(torch.ops.deepep.combine.default)
+    
     # DeepEP + ETP is not supported yet
     if use_deepep and parallel_dims.etp_enabled:
         raise NotImplementedError(

torchtitan/models/moe/moe.py

@@ -506,16 +506,16 @@ def init_weights(
                 )
 
 
-def build_moe(args: MoEArgs, dim: int, hidden_dim: int, communication_backend: str = "standard") -> nn.Module:
+def build_moe(args: MoEArgs, dim: int, hidden_dim: int, moe_impl: str = "standard") -> nn.Module:
     """Factory for MoE with different backends: 'standard' (all-to-all) or 'deepep' (DeepEP).
 
     If 'deepep' is requested but DeepEP is not installed, falls back to standard with a warning.
     """
-    if communication_backend == "deepep":
+    if moe_impl == "deepep":
         try:
-            from .moe_deepep import MoEWithDeepEP
+            from .moe_deepep import DeepEPMoE
             logger.info(f"DeepEP MoE: num_experts={args.num_experts}, top_k={args.top_k}, dim={dim}, hidden_dim={hidden_dim}")
-            return MoEWithDeepEP(moe_args=args, dim=dim, hidden_dim=hidden_dim)
+            return DeepEPMoE(moe_args=args, dim=dim, hidden_dim=hidden_dim)
         except ImportError as e:
             logger.warning(
                 f"DeepEP requested but not available: {e}. "

torchtitan/models/moe/moe_deepep.py

@@ -11,7 +11,7 @@
 from .moe import MoE, MoEArgs
 
 
-class MoEWithDeepEP(MoE):
+class DeepEPMoE(MoE):
     """
     Mixture of Experts with DeepEP communication.