Sort tensors/weight based on their offset in object when loading checkpoints

Author: Ilya Isaev

s3torchconnector/src/s3torchconnector/dcp/s3_file_system.py

@@ -252,7 +252,7 @@ def _escape_path(string):
         return "/".join(parts)
 
 
-from torch.distributed.checkpoint.planner import SavePlan
+from torch.distributed.checkpoint.planner import SavePlan, LoadPlan
 import dataclasses
 from dataclasses import dataclass
 
@@ -359,6 +359,11 @@ def __init__(
     def validate_checkpoint_id(cls, checkpoint_id: Union[str, os.PathLike]) -> bool:
         return S3FileSystem.validate_checkpoint_id(checkpoint_id)
 
+    def prepare_local_plan(self, plan: LoadPlan) -> LoadPlan:
+        # Sort items in plan based on their offset in checkpoints shards
+        plan.items.sort(key=lambda  item: self.storage_data[item.storage_index].offset)
+        return plan
+
 
 def _path_or_str_to_str(path: Union[str, os.PathLike]) -> str:
     return path if isinstance(path, str) else str(path)