fix comments

Author: fegin

torchtitan/experiments/gpt_oss/infra/parallelize.py

@@ -121,18 +121,18 @@ def parallelize_gptoss(
     dp_mesh: DeviceMesh | None = None
     if parallel_dims.fsdp_enabled or parallel_dims.ep_enabled:
         # apply FSDP or HSDP, potentially with Context Parallel
-        names = (
+        dp_mesh_names = (
             ["dp_replicate", "fsdp"] if parallel_dims.dp_replicate_enabled else ["fsdp"]
         )
-        dp_mesh = parallel_dims.get_mesh(names)
+        dp_mesh = parallel_dims.get_mesh(dp_mesh_names)
 
         # the mesh dim names of which the MoE params are sharded on via FSDP/HSDP
-        edp_mesh = None
-        if parallel_dims.ep_enabled:
-            if parallel_dims.dp_replicate_enabled:
-                edp_mesh = parallel_dims.get_mesh(["dp_replicate", "efsdp"])
-            else:
-                edp_mesh = parallel_dims.get_mesh("efsdp")
+        edp_mesh_names = (
+            ["dp_replicate", "efsdp"]
+            if parallel_dims.dp_replicate_enabled
+            else ["efsdp"]
+        )
+        edp_mesh = parallel_dims.get_mesh(edp_mesh_names)
 
         apply_fsdp(
             model,

torchtitan/experiments/simple_fsdp/deepseek_v3/parallelize.py

@@ -130,21 +130,22 @@ def parallelize_deepseekv3(
             dp_mode = "fully_shard"
 
         dp_mesh = parallel_dims.get_mesh(dp_mesh_dim_names)
-        # the mesh dim names of which the MoE params are sharded on via FSDP/HSDP
 
-        if parallel_dims.dp_replicate_enabled:
-            dp_mesh_dim_names = ["dp_replicate", "efsdp"]
-        else:
-            dp_mesh_dim_names = ["efsdp"]
-        edp_mesh = parallel_dims.get_mesh(dp_mesh_dim_names)
+        # the mesh dim names of which the MoE params are sharded on via FSDP/HSDP
+        edp_mesh_names = (
+            ["dp_replicate", "efsdp"]
+            if parallel_dims.dp_replicate_enabled
+            else ["efsdp"]
+        )
+        edp_mesh = parallel_dims.get_mesh(edp_mesh_names)
 
         for _, transformer_block in model.layers.items():
             if transformer_block.moe_enabled and parallel_dims.ep_enabled:
                 experts_shard_dim = 0
                 assert edp_mesh is not None
                 assert hasattr(transformer_block, "moe")
                 if (
-                    edp_mesh.size() * parallel_dims.ep
+                    edp_mesh["efsdp"].size() * parallel_dims.ep
                     > transformer_block.moe.experts.num_experts
                 ):
                     experts_shard_dim = 1

torchtitan/models/deepseek_v3/infra/parallelize.py

@@ -124,18 +124,18 @@ def parallelize_deepseekv3(
     dp_mesh: DeviceMesh | None = None
     if parallel_dims.fsdp_enabled or parallel_dims.ep_enabled:
         # apply FSDP or HSDP, potentially with Context Parallel
-        names = (
+        dp_mesh_names = (
             ["dp_replicate", "fsdp"] if parallel_dims.dp_replicate_enabled else ["fsdp"]
         )
-        dp_mesh = parallel_dims.get_mesh(names)
+        dp_mesh = parallel_dims.get_mesh(dp_mesh_names)
 
         # the mesh dim names of which the MoE params are sharded on via FSDP/HSDP
-        names = (
+        edp_mesh_names = (
             ["dp_replicate", "efsdp"]
             if parallel_dims.dp_replicate_enabled
             else ["efsdp"]
         )
-        edp_mesh = parallel_dims.get_mesh(names)
+        edp_mesh = parallel_dims.get_mesh(edp_mesh_names)
 
         apply_fsdp(
             model,

torchtitan/models/llama4/infra/parallelize.py

@@ -127,18 +127,18 @@ def parallelize_llama(
 
     if parallel_dims.fsdp_enabled or parallel_dims.ep_enabled:
         # dp_mesh is the mesh for FSDP/HSDP
-        names = (
+        dp_mesh_names = (
             ["dp_replicate", "fsdp"] if parallel_dims.dp_replicate_enabled else ["fsdp"]
         )
-        dp_mesh = parallel_dims.get_mesh(names)
+        dp_mesh = parallel_dims.get_mesh(dp_mesh_names)
 
         # the mesh dim names of which the MoE params are sharded on via FSDP/HSDP
-        edp_mesh = None
-        if parallel_dims.ep_enabled:
-            if parallel_dims.dp_replicate_enabled:
-                edp_mesh = parallel_dims.get_mesh(["dp_replicate", "efsdp"])
-            else:
-                edp_mesh = parallel_dims.get_mesh("efsdp")
+        edp_mesh_names = (
+            ["dp_replicate", "efsdp"]
+            if parallel_dims.dp_replicate_enabled
+            else ["efsdp"]
+        )
+        edp_mesh = parallel_dims.get_mesh(edp_mesh_names)
 
         apply_fsdp(
             model,
@@ -337,7 +337,10 @@ def apply_fsdp(
             _experts_shard_placement_fn = None
             assert edp_mesh is not None
             assert hasattr(transformer_block, "moe")
-            if edp_mesh.size() * ep_degree > transformer_block.moe.experts.num_experts:
+            if (
+                edp_mesh["efsdp"].size() * ep_degree
+                > transformer_block.moe.experts.num_experts
+            ):
                 _experts_shard_placement_fn = lambda param: Shard(1)
 
             fully_shard(

torchtitan/models/qwen3/infra/parallelize.py

@@ -116,18 +116,18 @@ def parallelize_qwen3(
 
     if parallel_dims.fsdp_enabled:
         # apply FSDP or HSDP, potentially with Context Parallel
-        names = (
+        dp_mesh_names = (
             ["dp_replicate", "fsdp"] if parallel_dims.dp_replicate_enabled else ["fsdp"]
         )
-        dp_mesh = parallel_dims.get_mesh(names)
+        dp_mesh = parallel_dims.get_mesh(dp_mesh_names)
 
         # the mesh dim names of which the MoE params are sharded on via FSDP/HSDP
-        names = (
+        edp_mesh_names = (
             ["dp_replicate", "efsdp"]
             if parallel_dims.dp_replicate_enabled
             else ["efsdp"]
         )
-        edp_mesh = parallel_dims.get_mesh(names)
+        edp_mesh = parallel_dims.get_mesh(edp_mesh_names)
 
         apply_fsdp(
             model,