Fix load for non-HSDP device mesh (mosaicml#2997)

* fix * lint * lint * fix
mvpatel2000 · Feb 13, 2024 · d2b9cd3 · d2b9cd3
1 parent a88c7fe
commit d2b9cd3
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/composer/utils/checkpoint.py b/composer/utils/checkpoint.py
@@ -213,7 +213,9 @@ def __init__(self, source_path: str, destination_path: str, object_store: Union[
         super().__init__(destination_path)
 
     def read_data(self, plan: LoadPlan, planner: LoadPlanner):
-        first_replica = self.device_mesh is None or self.device_mesh.get_local_rank(mesh_dim=0) == 0
+        # Download files if not using HSDP or if on first replica with HSDP enabled
+        first_replica = self.device_mesh is None or self.device_mesh.ndim == 1 or (
+            self.device_mesh.ndim >= 2 and self.device_mesh.get_local_rank(mesh_dim=0) == 0)
 
         # 1. Download to the destination all files this rank needs if on first replica
         if first_replica: