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)