lint code by bash ./tools/autoformat.sh

Author: wplf

megatron/core/distributed/fsdp/mcore_fsdp_adapter.py

@@ -251,7 +251,9 @@ def _init_dist_index(self, pg_collection):
 
         if enable_hsdp:
             if expt_dp_group is not None:
-                expt_mesh = _get_hsdp_tp_mesh(outer_fsdp_group, expt_dp_group, expt_tp_group, ep_size=ep_group.size())
+                expt_mesh = _get_hsdp_tp_mesh(
+                    outer_fsdp_group, expt_dp_group, expt_tp_group, ep_size=ep_group.size()
+                )
                 expt_device_mesh = DeviceMesh.from_group(
                     [outer_fsdp_group, expt_dp_group, expt_tp_group],
                     device_type="cuda",
@@ -274,7 +276,7 @@ def _init_dist_index(self, pg_collection):
                 tp_dim="tp",
                 hybrid_fsdp_group=hybrid_fsdp_group,
                 hybrid_fsdp_expt_group=hybrid_fsdp_expt_group,
-                expt_device_mesh=expt_device_mesh
+                expt_device_mesh=expt_device_mesh,
             )
         else:
             if ep_group is not None:
@@ -385,7 +387,10 @@ def _get_hsdp_tp_mesh(outer_fsdp_dp_group, dp_cp_group, tp_group, ep_size=1):
     assert (
         len(dp_tp_meshes) == 1
     ), f"[Megatron-FSDP] Current rank {rank} is not unique in the mesh ranks {mesh.tolist()}."
-    assert len(dp_tp_meshes[0].reshape(-1).tolist()) == outer_fsdp_dp_group.size() * dp_cp_group.size() * tp_group.size(), (
+    assert (
+        len(dp_tp_meshes[0].reshape(-1).tolist())
+        == outer_fsdp_dp_group.size() * dp_cp_group.size() * tp_group.size()
+    ), (
         f"[Megatron-FSDP] DP-TP mesh size {len(dp_tp_meshes[0].reshape(-1).tolist())} "
         f"does not match expected size {outer_fsdp_dp_group.size() * dp_cp_group.size() * tp_group.size()}."
     )

megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py

@@ -1869,7 +1869,9 @@ def _init_each_parameter_group_buffers(self, meta_device_init_fp8_params):
                 hsdp_buf_dp_group = self.dist_index.get_fsdp_group(
                     is_expert_parallel=group.is_expert_param
                 )
-                main_buf_extra_kwargs["dp_rank"] = self.dist_index.get_logical_hybrid_fsdp_rank(is_expert_parallel=group.is_expert_param)
+                main_buf_extra_kwargs["dp_rank"] = self.dist_index.get_logical_hybrid_fsdp_rank(
+                    is_expert_parallel=group.is_expert_param
+                )
             else:
                 main_buf_dp_group = self.dist_index.get_fsdp_group(
                     is_expert_parallel=group.is_expert_param

megatron/core/distributed/fsdp/src/megatron_fsdp/utils.py

@@ -929,7 +929,11 @@ def get_logical_hybrid_fsdp_rank(self, is_expert_parallel: bool = False):
             self.hsdp_outer_dp_shard
         ), "get_logical_hybrid_fsdp_rank is only valid when full-shard hybrid FSDP is enabled."
 
-        _hybrid_fsdp_group_name = "_hybrid_fsdp_group_ranks" if not is_expert_parallel else "_hybrid_fsdp_expt_group_ranks"
+        _hybrid_fsdp_group_name = (
+            "_hybrid_fsdp_group_ranks"
+            if not is_expert_parallel
+            else "_hybrid_fsdp_expt_group_ranks"
+        )
 
         if not hasattr(self, _hybrid_fsdp_group_name):
             dp_world_size = self.get_dp_group(is_expert_parallel).size()
@@ -944,7 +948,9 @@ def get_logical_hybrid_fsdp_rank(self, is_expert_parallel: bool = False):
             setattr(self, _hybrid_fsdp_group_name, mesh.tolist())
 
         # Find the index for the current rank in the hybrid group
-        return getattr(self, _hybrid_fsdp_group_name).index(self.get_dp_group(is_expert_parallel).rank())
+        return getattr(self, _hybrid_fsdp_group_name).index(
+            self.get_dp_group(is_expert_parallel).rank()
+        )
 
 
 class GlobalMemoryBuffer: