Fix HSDP with ShardDegree < 8

composer/utils/checkpoint.py

@@ -289,6 +289,7 @@ def read_data(self, plan: LoadPlan, planner: LoadPlanner):
         if self.device_mesh is not None and self.device_mesh.ndim == 2:
             # Broadcast file to all replicas
             replicate_process_group = self.device_mesh.get_group(0)
+            shard_process_group = self.device_mesh.get_group(1)
             shard_size = self.device_mesh.size(1)
             rank_in_first_replica = dist.get_global_rank() % shard_size
             sender = dist.get_global_rank() == rank_in_first_replica
@@ -304,7 +305,9 @@ def read_data(self, plan: LoadPlan, planner: LoadPlanner):
 
             # Send each file to the appropriate rank
             for file_name in file_list:
-                if dist.get_local_rank() == 0:  # Only 1 rank per node needs to transfer file
+                if dist.get_local_rank() == 0 or (
+                    dist.get_global_rank(shard_process_group) == 0  # pyright: ignore[reportGeneralTypeIssues]
+                ):  # Only 1 rank per node needs to transfer file
                     full_path = os.path.join(self.destination_path, file_name)
                     log.debug(f'Transferring {full_path=}')
                     file_object = [None]
@@ -318,7 +321,7 @@ def read_data(self, plan: LoadPlan, planner: LoadPlanner):
                     )
                     received_file_object = file_object[0]
                     assert received_file_object is not None
-                    if receiver and not os.path.exists(full_path):
+                    if receiver and not os.path.exists(full_path) and dist.get_local_rank() == 0:
                         with open(full_path, 'wb') as f:
                             f.write(received_file_object['content'])
 

composer/utils/dist.py

@@ -225,13 +225,23 @@ def get_world_size() -> int:
     )
 
 
-def get_global_rank() -> int:
-    """Returns the global rank of the current process, which is on ``[0; WORLD_SIZE - 1]``.
+def get_global_rank(group: Optional[dist.ProcessGroup] = None) -> int:
+    """Returns the global rank of the current process in the input PG, which is on ``[0; group.WORLD_SIZE - 1]``.
+
+    Args:
+        group (ProcessGroup, optional): The process group. If ``None``, it will return env_var ``RANK``
 
     Returns:
-        int: The global rank.
+        int: The global rank in input process group.
     """
-    return _get_distributed_config_var(env_var='RANK', human_name='global rank', default=0, fetch_fn_name='get_rank')
+    if group is None:
+        return _get_distributed_config_var(
+            env_var='RANK',
+            human_name='global rank',
+            default=0,
+            fetch_fn_name='get_rank',
+        )
+    return dist.get_rank(group)
 
 
 def get_local_world_size() -> int: