From bc7b134f58699b01c8e0844c8b953378e74925d9 Mon Sep 17 00:00:00 2001 From: Ehsan Barkhordar Date: Tue, 8 Sep 2026 02:42:25 +0000 Subject: [PATCH] Sort checkpoint shard files by numeric rank Shard file names pad the rank field to two digits, and both readers that glob those names sort them as plain strings before indexing the result by model-parallel rank. Past a parallel degree of 100 the pad overflows, model_100 sorts between model_10 and model_11, and the index no longer identifies the shard it names. Sort with natural_keys instead, which get_checkpoint_files in zero_to_fp32.py and _get_checkpoint_files in ds_to_universal.py already use for globbed checkpoint files. The two orders agree whenever every rank field is two digits or fewer, so existing checkpoints keep loading and the on-disk name format is untouched. Signed-off-by: Ehsan Barkhordar --- deepspeed/checkpoint/ds_to_universal.py | 14 +-- deepspeed/checkpoint/utils.py | 14 +++ deepspeed/runtime/engine.py | 5 +- deepspeed/runtime/pipe/module.py | 5 +- .../checkpoint/test_ckpt_file_ordering.py | 93 +++++++++++++++++++ 5 files changed, 114 insertions(+), 17 deletions(-) create mode 100644 tests/unit/checkpoint/test_ckpt_file_ordering.py diff --git a/deepspeed/checkpoint/ds_to_universal.py b/deepspeed/checkpoint/ds_to_universal.py index b62a88ce7252..ec6d03733d04 100755 --- a/deepspeed/checkpoint/ds_to_universal.py +++ b/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}' diff --git a/deepspeed/checkpoint/utils.py b/deepspeed/checkpoint/utils.py index 5964da00728e..d98a6eeb002b 100644 --- a/deepspeed/checkpoint/utils.py +++ b/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, diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py index 86918bd71c5a..e0acd2981d92 100644 --- a/deepspeed/runtime/engine.py +++ b/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, diff --git a/deepspeed/runtime/pipe/module.py b/deepspeed/runtime/pipe/module.py index f41305bac27a..bcc0ace90625 100644 --- a/deepspeed/runtime/pipe/module.py +++ b/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): diff --git a/tests/unit/checkpoint/test_ckpt_file_ordering.py b/tests/unit/checkpoint/test_ckpt_file_ordering.py new file mode 100644 index 000000000000..2974432efbf5 --- /dev/null +++ b/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)