misc

Author: fegin

torchtitan/distributed/parallel_dims.py

@@ -121,7 +121,6 @@ def unflatten_mesh(
         )
 
         batch = self.dp_replicate * self.dp_shard
-        loss = self.dp_replicate * self.dp_shard * self.cp
         fsdp = self.dp_shard * self.cp
         efsdp = fsdp * self.tp // (self.etp * self.ep)
 
@@ -145,12 +144,12 @@ def unflatten_mesh(
             (self.pp, self.dp_replicate, efsdp, self.ep, self.etp),
         )
 
-        # We have created all the required 1D meshes. This part is to create the
-        # all the 2D meshes. We pre-created 2D meshes and error out if the users
-        # try to access a 2D mesh that is not pre-created.
-        hsdp_mesh = dense_mesh["dp_replicate", "fsdp"]
-        ehsdp_mesh = sparse_mesh["dp_replicate", "efsdp"]
-        ep_etp_mesh = sparse_mesh["ep", "etp"]
+        self._global_meshes = {
+            "dataloading": dataloading_mesh,
+            "loss": loss_mesh,
+            "dense": dense_mesh,
+            "sparse": sparse_mesh,
+        }
 
         self._meshes = {
             "pp": dataloading_mesh["pp"],
@@ -163,9 +162,6 @@ def unflatten_mesh(
             "ep": sparse_mesh["ep"],
             "efsdp": sparse_mesh["efsdp"],
             "etp": sparse_mesh["etp"],
-            "dp_replicate_fsdp": hsdp_mesh,
-            "dp_replicate_efsdp": ehsdp_mesh,
-            "ep_etp": ep_etp_mesh,
         }
 
         # Validate mesh sizes
@@ -191,19 +187,10 @@ def _validate_meshes(self):
             "ep": self.ep,
             "efsdp": self.dp_shard * self.cp * self.tp // (self.etp * self.ep),
             "etp": self.etp,
-            "dp_replicate_fsdp": (self.dp_replicate, self.dp_shard * self.cp),
-            "dp_replicate_efsdp": (
-                self.dp_replicate,
-                self.dp_shard * self.cp * self.tp // (self.etp * self.ep),
-            ),
-            "ep_etp": (self.ep, self.etp),
         }
 
         for mesh_name, expected_size in expected_sizes.items():
-            if isinstance(expected_size, tuple):
-                actual_size = self._meshes[mesh_name].shape
-            else:
-                actual_size = self._meshes[mesh_name].size()
+            actual_size = self._meshes[mesh_name].size()
             assert actual_size == expected_size, (
                 f"Mesh '{mesh_name}' has unexpected size: "
                 f"expected {expected_size}, got {actual_size}"
@@ -232,17 +219,24 @@ def get_mesh(self, dims: str | list[str]) -> DeviceMesh | None:
         if isinstance(dims, str):
             dims = [dims]
 
-        mesh_name = "_".join(dims)
-        if mesh_name not in self._meshes:
-            raise ValueError(
-                f"Invalid mesh dim: '{mesh_name}'. "
-                f"Valid dimensions are: {list(self._meshes.keys())}"
-            )
+        for mesh_name in dims:
+            if mesh_name not in self._meshes:
+                raise ValueError(
+                    f"Invalid mesh dim: '{mesh_name}'. "
+                    f"Valid dimensions are: {list(self._meshes.keys())}"
+                )
 
         if any(not self._mesh_exist(dim, self._meshes[dim].size()) for dim in dims):
             return None
 
-        return self._meshes[mesh_name]
+        if len(dims) == 1:
+            return self._meshes[dims[0]]
+        else:
+            for global_mesh in self._global_meshes.values():
+                if not set(dims).issubset(set(global_mesh.mesh_dim_names)):
+                    continue
+                return global_mesh[tuple(dims)]
+            raise ValueError(f"Invalid mesh name combinations {dims}.")
 
     def get_all_meshes(self, one_dimensioal_only: bool = True) -> dict[str, DeviceMesh]:
         if not self._meshes:

torchtitan/experiments/gpt_oss/infra/parallelize.py

@@ -64,6 +64,10 @@ def parallelize_gptoss(
     if job_config.parallelism.context_parallel_degree > 1 and use_flex_attn:
         raise NotImplementedError("CP support for FlexAttention is still in progress.")
 
+    model_compile_enabled = (
+        job_config.compile.enable and "model" in job_config.compile.components
+    )
+
     if parallel_dims.tp_enabled:
         if (
             job_config.parallelism.enable_async_tensor_parallel
@@ -105,10 +109,6 @@ def parallelize_gptoss(
             etp_enabled=parallel_dims.etp_enabled,
         )
 
-    model_compile_enabled = (
-        job_config.compile.enable and "model" in job_config.compile.components
-    )
-
     if job_config.activation_checkpoint.mode != "none":
         apply_ac(
             model,
@@ -127,11 +127,12 @@ def parallelize_gptoss(
         dp_mesh = parallel_dims.get_mesh(names)
 
         # the mesh dim names of which the MoE params are sharded on via FSDP/HSDP
-        dp_mod_ep_mesh_dim_names = []
+        edp_mesh = None
         if parallel_dims.ep_enabled:
             if parallel_dims.dp_replicate_enabled:
-                dp_mod_ep_mesh_dim_names.append("dp_replicate")
-            dp_mod_ep_mesh_dim_names.append("dp_shard_mod_ep")
+                edp_mesh = parallel_dims.get_mesh(["dp_replicate", "efsdp"])
+            else:
+                edp_mesh = parallel_dims.get_mesh("efsdp")
 
         apply_fsdp(
             model,
@@ -142,11 +143,7 @@ def parallelize_gptoss(
             cpu_offload=job_config.training.enable_cpu_offload,
             reshard_after_forward_policy=job_config.parallelism.fsdp_reshard_after_forward,
             ep_degree=parallel_dims.ep,
-            dp_mod_ep_mesh=(
-                parallel_dims.get_mesh(dp_mod_ep_mesh_dim_names)
-                if parallel_dims.ep_enabled
-                else None
-            ),
+            edp_mesh=edp_mesh,
         )
 
         if parallel_dims.dp_replicate_enabled:

torchtitan/models/deepseek_v3/infra/parallelize.py

@@ -92,9 +92,9 @@ def parallelize_deepseekv3(
             model,
             tp_mesh=parallel_dims.get_mesh("tp") if parallel_dims.tp_enabled else None,
             ep_mesh=parallel_dims.get_mesh("ep") if parallel_dims.ep_enabled else None,
-            etp_mesh=parallel_dims.get_mesh("etp")
-            if parallel_dims.etp_enabled
-            else None,
+            etp_mesh=(
+                parallel_dims.get_mesh("etp") if parallel_dims.etp_enabled else None
+            ),
             ep_etp_mesh=(
                 parallel_dims.get_mesh("ep_etp")
                 if parallel_dims.tp_enabled
@@ -146,7 +146,7 @@ def parallelize_deepseekv3(
             cpu_offload=job_config.training.enable_cpu_offload,
             reshard_after_forward_policy=job_config.parallelism.fsdp_reshard_after_forward,
             ep_degree=parallel_dims.ep,
-            dp_mod_ep_mesh=edp_mesh,
+            edp_mesh=edp_mesh,
             gradient_divide_factor=parallel_dims.fsdp_gradient_divide_factor,
         )
 

torchtitan/models/llama4/infra/parallelize.py

@@ -133,12 +133,12 @@ def parallelize_llama(
         dp_mesh = parallel_dims.get_mesh(names)
 
         # the mesh dim names of which the MoE params are sharded on via FSDP/HSDP
-        dp_mod_ep_mesh = None
+        edp_mesh = None
         if parallel_dims.ep_enabled:
             if parallel_dims.dp_replicate_enabled:
-                dp_mod_ep_mesh = parallel_dims.get_mesh(["dp_replicate", "efsdp"])
+                edp_mesh = parallel_dims.get_mesh(["dp_replicate", "efsdp"])
             else:
-                dp_mod_ep_mesh = parallel_dims.get_mesh("efsdp")
+                edp_mesh = parallel_dims.get_mesh("efsdp")
 
         apply_fsdp(
             model,
@@ -149,7 +149,7 @@ def parallelize_llama(
             cpu_offload=job_config.training.enable_cpu_offload,
             reshard_after_forward_policy=job_config.parallelism.fsdp_reshard_after_forward,
             ep_degree=parallel_dims.ep,
-            dp_mod_ep_mesh=dp_mod_ep_mesh,
+            edp_mesh=edp_mesh,
             gradient_divide_factor=parallel_dims.fsdp_gradient_divide_factor,
         )
 
@@ -274,7 +274,7 @@ def apply_fsdp(
     cpu_offload: bool = False,
     reshard_after_forward_policy: str = "default",
     ep_degree: int = 1,
-    dp_mod_ep_mesh: DeviceMesh | None = None,
+    edp_mesh: DeviceMesh | None = None,
     gradient_divide_factor: int | None = None,
 ):
     """
@@ -323,10 +323,10 @@ def apply_fsdp(
     for layer_id, transformer_block in model.layers.items():
         # NOTE: When EP is enabled, In an MoE layer, we use the following FSDP wrapping
         # - the router and the shared experts are sharded together with the TransformerBlock
-        # - the routed experts are sharded with the remaining dp_mod_ep_mesh
+        # - the routed experts are sharded with the remaining edp_mesh
         if transformer_block.moe_enabled and ep_degree > 1:
             fsdp_mod_ep_config = fsdp_config.copy()
-            fsdp_mod_ep_config["mesh"] = dp_mod_ep_mesh
+            fsdp_mod_ep_config["mesh"] = edp_mesh
 
             # NOTE: EP alreadys shards the routed experts on dim 0 (num_experts).
             #       When dp_mod_ep * ep > num_experts, FSDP default dim-0 sharding
@@ -335,12 +335,9 @@ def apply_fsdp(
             #       on non-0 dim. For now it may not be worth the complexity to support
             #       shard_placement_fn on the outer TransformerBlock-level FSDP.
             _experts_shard_placement_fn = None
-            assert dp_mod_ep_mesh is not None
+            assert edp_mesh is not None
             assert hasattr(transformer_block, "moe")
-            if (
-                dp_mod_ep_mesh.size() * ep_degree
-                > transformer_block.moe.experts.num_experts
-            ):
+            if edp_mesh.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

@@ -138,7 +138,7 @@ def parallelize_qwen3(
             cpu_offload=job_config.training.enable_cpu_offload,
             reshard_after_forward_policy=job_config.parallelism.fsdp_reshard_after_forward,
             ep_degree=parallel_dims.ep,
-            dp_mod_ep_mesh=edp_mesh,
+            edp_mesh=edp_mesh,
             gradient_divide_factor=parallel_dims.fsdp_gradient_divide_factor,
         )