lint

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

Author: ZhiyuLi-Nvidia

nemo_automodel/components/distributed/fsdp2.py

@@ -259,9 +259,9 @@ def _try_flatten_submesh(submesh_dims: list, flattened_name: str) -> None:
                 logger.debug(f"Skipping {flattened_name} flatten (dimensions are all size 1)")
 
         # Create flattened submeshes for data parallelism combinations
-        _try_flatten_submesh(dp_mesh_dim_names, "dp")                    # dp_replicate + dp_shard
+        _try_flatten_submesh(dp_mesh_dim_names, "dp")  # dp_replicate + dp_shard
         _try_flatten_submesh(dp_shard_cp_mesh_dim_names, "dp_shard_cp")  # dp_shard + cp
-        _try_flatten_submesh(dp_cp_mesh_dim_names, "dp_cp")              # dp_replicate + dp_shard + cp
+        _try_flatten_submesh(dp_cp_mesh_dim_names, "dp_cp")  # dp_replicate + dp_shard + cp
 
         return self.device_mesh
 

nemo_automodel/components/distributed/parallelizer.py

@@ -158,9 +158,7 @@ def parallelize(
         tp_mesh = device_mesh[tp_mesh_name]
 
         # Determine the appropriate FSDP mesh dimensions (must be 1D or 2D)
-        dp_mesh_dim_names = self._resolve_fsdp_mesh_dims(
-            device_mesh, dp_replicate_mesh_name, dp_shard_cp_mesh_name
-        )
+        dp_mesh_dim_names = self._resolve_fsdp_mesh_dims(device_mesh, dp_replicate_mesh_name, dp_shard_cp_mesh_name)
         dp_mesh = device_mesh[dp_mesh_dim_names]
 
         # Extract layers from the model for parallelization