Sort checkpoint shard files by numeric rank

deepspeed/checkpoint/ds_to_universal.py

@@ -57,6 +57,7 @@
     is_autoep_zero3_partitioned_entry,
     validate_autoep_zero3_partitioned_metadata,
 )
+from deepspeed.checkpoint.utils import natural_keys
 
 
 def parse_arguments():
@@ -89,19 +90,6 @@ def parse_arguments():
     return args
 
 
-def atoi(text):
-    return int(text) if text.isdigit() else text
-
-
-def natural_keys(text):
-    '''
-    alist.sort(key=natural_keys) sorts in human order
-    http://nedbatchelder.com/blog/200712/human_sorting.html
-    (See Toothy's implementation in the comments)
-    '''
-    return [atoi(c) for c in re.split(r'(\d+)', text)]
-
-
 def _create_checkpoint_paths(base_folder, iteration, tp_degree, pp_degree):
     path_list = []
     iter_folder = f'iter_{iteration:07d}'

deepspeed/checkpoint/utils.py

@@ -4,10 +4,24 @@
 # DeepSpeed Team
 
 import os
+import re
 import torch
 from .constants import (MODEL_FILE_PREFIX, MODEL_FILE_SUFFIX, OPTIM_FILE_SUFFIX, ZERO_FILE_PREFIX)
 
 
+def atoi(text):
+    return int(text) if text.isdigit() else text
+
+
+def natural_keys(text):
+    '''
+    alist.sort(key=natural_keys) sorts in human order
+    http://nedbatchelder.com/blog/200712/human_sorting.html
+    (See Toothy's implementation in the comments)
+    '''
+    return [atoi(c) for c in re.split(r'(\d+)', text)]
+
+
 def get_model_ckpt_name_for_rank(base_folder, mp_rank_str):
     ckpt_name = os.path.join(
         base_folder,

deepspeed/runtime/engine.py

@@ -97,7 +97,7 @@
     is_autoep_zero3_partitioned_entry,
     validate_autoep_zero3_partitioned_metadata,
 )
-from deepspeed.checkpoint.utils import clone_tensors_for_torch_save
+from deepspeed.checkpoint.utils import clone_tensors_for_torch_save, natural_keys
 from deepspeed.checkpoint.ds_to_universal import dp_index_to_str
 from deepspeed.runtime.sparse_tensor import SparseTensor
 
@@ -4360,7 +4360,8 @@ def _get_all_ckpt_names(self, checkpoints_path, tag):
         import glob
 
         ckpt_files = glob.glob(ckpt_file_pattern)
-        ckpt_files.sort()
+        # Callers index this list by model-parallel rank, so it must be ordered numerically.
+        ckpt_files.sort(key=natural_keys)
         return ckpt_files
 
     def load_checkpoint(self,

deepspeed/runtime/pipe/module.py

@@ -20,7 +20,7 @@
 from .topology import PipeDataParallelTopology, PipelineParallelGrid
 from deepspeed.runtime.state_dict_factory import SDLoaderFactory
 from deepspeed.accelerator import get_accelerator
-from deepspeed.checkpoint.utils import clone_tensors_for_torch_save
+from deepspeed.checkpoint.utils import clone_tensors_for_torch_save, natural_keys
 
 
 class PipelineError(Exception):
@@ -601,7 +601,8 @@ def ckpt_layer_path_list(self, ckpt_dir, local_layer_idx):
         layer_ckpt_path = os.path.join(ckpt_dir, f'layer_{idx:02d}-')
         layer_ckpt_path += "*model_states.pt"
         ckpt_files = glob.glob(layer_ckpt_path)
-        ckpt_files.sort()
+        # Callers index this list by model-parallel rank, so it must be ordered numerically.
+        ckpt_files.sort(key=natural_keys)
         return ckpt_files
 
     def save_state_dict(self, save_dir, checkpoint_engine, exclude_frozen_params=False):

tests/unit/checkpoint/test_ckpt_file_ordering.py

@@ -0,0 +1,93 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import os
+from types import SimpleNamespace
+
+import torch
+
+from deepspeed.runtime.engine import DeepSpeedEngine
+from deepspeed.runtime.pipe.module import PipelineModule
+from deepspeed.runtime.pipe.topology import PipeModelDataParallelTopology
+from deepspeed.runtime.state_dict_factory import SDLoaderFactory
+
+# The rank field in a checkpoint file name is padded to two digits, so this is the
+# smallest interesting degree: the field overflows and lexicographic order diverges
+# from rank order.
+MP_DEGREE_OVER_PAD = 128
+MP_DEGREE_UNDER_PAD = 8
+
+
+class PipelineModuleStub:
+    """Supplies only the attributes the two checkpoint path methods read off ``self``."""
+
+    ckpt_layer_path = PipelineModule.ckpt_layer_path
+    ckpt_layer_path_list = PipelineModule.ckpt_layer_path_list
+
+    def __init__(self, topo, global_rank=0):
+        self._local_start = 0
+        self.global_rank = global_rank
+        self._grid = SimpleNamespace(_topo=topo)
+
+
+class EngineStub:
+    """Supplies only the attributes the two checkpoint name methods read off ``self``."""
+
+    _get_ckpt_name = DeepSpeedEngine._get_ckpt_name
+    _get_all_ckpt_names = DeepSpeedEngine._get_all_ckpt_names
+
+    def __init__(self, checkpoint_mp_rank=0):
+        self.checkpoint_mp_rank = checkpoint_mp_rank
+
+    def zero_optimization_partition_weights(self):
+        return False
+
+    def load_universal_checkpoint(self):
+        return False
+
+
+def write_pipeline_layer_shards(ckpt_dir, mp_degree):
+    topo = PipeModelDataParallelTopology(num_pp=1, num_mp=mp_degree, num_dp=1)
+    for rank in range(mp_degree):
+        torch.save({'rank': rank}, PipelineModuleStub(topo, rank).ckpt_layer_path(ckpt_dir, 0))
+    return topo
+
+
+def test_pipeline_layer_shards_load_by_numeric_rank(tmpdir):
+    ckpt_dir = str(tmpdir)
+    topo = write_pipeline_layer_shards(ckpt_dir, MP_DEGREE_OVER_PAD)
+
+    ckpt_list = PipelineModuleStub(topo).ckpt_layer_path_list(ckpt_dir, 0)
+    assert len(ckpt_list) == MP_DEGREE_OVER_PAD
+
+    sd_loader = SDLoaderFactory.get_sd_loader(ckpt_list, version=2.0, checkpoint_engine=None)
+    for mp_rank in range(MP_DEGREE_OVER_PAD):
+        _, sd, _ = sd_loader.load(MP_DEGREE_OVER_PAD, mp_rank, module_key=None, is_pipe_parallel=True)
+        assert sd['rank'] == mp_rank
+
+
+def test_engine_checkpoint_names_ordered_by_numeric_rank(tmpdir):
+    ckpt_dir, tag = str(tmpdir), 'global_step100'
+    os.makedirs(os.path.join(ckpt_dir, tag))
+    for rank in range(MP_DEGREE_OVER_PAD):
+        torch.save({'rank': rank}, EngineStub(rank)._get_ckpt_name(ckpt_dir, tag))
+
+    ckpt_files = EngineStub()._get_all_ckpt_names(ckpt_dir, tag)
+    assert len(ckpt_files) == MP_DEGREE_OVER_PAD
+
+    loaded = [torch.load(f, weights_only=True)['rank'] for f in ckpt_files]
+    assert loaded == list(range(MP_DEGREE_OVER_PAD))
+
+
+def test_shard_order_below_pad_width_matches_lexicographic(tmpdir):
+    # Every checkpoint written before this change has a rank field of at most two
+    # digits, where the two orderings agree, so none of them is reordered.
+    ckpt_dir = str(tmpdir)
+    topo = write_pipeline_layer_shards(ckpt_dir, MP_DEGREE_UNDER_PAD)
+
+    ckpt_list = PipelineModuleStub(topo).ckpt_layer_path_list(ckpt_dir, 0)
+    loaded = [torch.load(p, weights_only=True)['rank'] for p in ckpt_list]
+    assert loaded == list(range(MP_DEGREE_UNDER_PAD))
+    assert ckpt_list == sorted(ckpt_list)