fix test

Signed-off-by: Zhiyu Li <[email protected]>

Author: ZhiyuLi-Nvidia

nemo_automodel/components/distributed/parallelizer.py

@@ -123,21 +123,30 @@ def _resolve_fsdp_mesh_dims(
         Returns:
             Tuple of dimension names to use for the FSDP mesh.
         """
-        if dp_shard_cp_mesh_name in device_mesh.mesh_dim_names:
-            # Best case: use pre-flattened (dp_replicate, dp_shard_cp)
+        mesh_dim_names = getattr(device_mesh, "mesh_dim_names", None)
+
+        # If mesh_dim_names not available, use the default passed names (backward compatibility)
+        if mesh_dim_names is None:
+            logger.debug(f"FSDP mesh: ({dp_replicate_mesh_name}, {dp_shard_cp_mesh_name}) [no mesh_dim_names]")
+            return (dp_replicate_mesh_name, dp_shard_cp_mesh_name)
+
+        # Use flattened dimension if available
+        if dp_shard_cp_mesh_name in mesh_dim_names:
             logger.debug(f"FSDP mesh: ({dp_replicate_mesh_name}, {dp_shard_cp_mesh_name}) [flattened]")
             return (dp_replicate_mesh_name, dp_shard_cp_mesh_name)
-        elif device_mesh["cp"].size() == 1:
-            # CP unused: construct 2D mesh without cp
-            logger.debug(f"FSDP mesh: ({dp_replicate_mesh_name}, dp_shard) [cp=1, unflattened]")
+
+        # Check if cp dimension exists and get its size
+        cp_size = device_mesh["cp"].size() if "cp" in mesh_dim_names else 1
+
+        if cp_size == 1:
+            logger.debug(f"FSDP mesh: ({dp_replicate_mesh_name}, dp_shard) [cp={cp_size}, unflattened]")
             return (dp_replicate_mesh_name, "dp_shard")
         else:
-            # This should be unreachable: if cp > 1, flattening should have succeeded
             raise RuntimeError(
-                f"Cannot construct FSDP mesh: cp={device_mesh['cp'].size()} but '{dp_shard_cp_mesh_name}' dimension not found. "
+                f"Cannot construct FSDP mesh: cp={cp_size} but '{dp_shard_cp_mesh_name}' dimension not found. "
                 f"FSDP requires 1D or 2D mesh, but would need 3D: ({dp_replicate_mesh_name}, dp_shard, cp). "
                 f"Flattening (dp_shard, cp) -> '{dp_shard_cp_mesh_name}' should occur in FSDP2Manager._get_device_mesh() (fsdp2.py). "
-                f"Available: {device_mesh.mesh_dim_names}"
+                f"Available: {mesh_dim_names}"
             )
 
     def parallelize(

nemo_automodel/recipes/base_recipe.py

@@ -498,12 +498,18 @@ def _resolve_dp_dim_name(self, include_cp: bool = False) -> str:
         Returns:
             The dimension name to use for DP operations.
         """
-        if include_cp and self.device_mesh["cp"].size() > 1:
-            # Prefer flattened "dp_cp", fallback to "dp_shard"
-            return "dp_cp" if "dp_cp" in self.device_mesh.mesh_dim_names else "dp_shard"
+        mesh_dim_names = getattr(self.device_mesh, "mesh_dim_names", None)
+
+        # If mesh_dim_names not available, use default behavior
+        if mesh_dim_names is None:
+            mesh_dim_names = []
+
+        cp_size = self.device_mesh["cp"].size() if "cp" in mesh_dim_names else 1
+
+        if include_cp and cp_size > 1:
+            return "dp_cp" if "dp_cp" in mesh_dim_names else "dp_shard"
         else:
-            # Prefer flattened "dp", fallback to "dp_shard"
-            return "dp" if "dp" in self.device_mesh.mesh_dim_names else "dp_shard"
+            return "dp" if "dp" in mesh_dim_names else "dp_shard"
 
     def _get_dp_group(self, include_cp: bool = False):
         """Get the data parallel process group.