Sort checkpoint shard files by numeric rank

[ebarkhordar]

get_rank_repr pads the rank field in a shard file name to two digits, and the 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 and the order stops matching the index: model_100 sorts between model_10 and model_11, so rank 11 loads rank 100's shard. Nothing raises and the shapes still match; the weights are just wrong. That is the failure this issue predicted.

A globbed shard list has to be ordered by numeric rank, because MegatronSDLoader indexes it positionally and get_merge_state_dicts/get_split_state_dict slice contiguous rank ranges out of it. Two call sites feed that loader, PipelineModule.ckpt_layer_path_list and DeepSpeedEngine._get_all_ckpt_names, both added in the same commit and both affected. They now sort with natural_keys, which get_checkpoint_files in zero_to_fp32.py and _get_checkpoint_files in ds_to_universal.py already use for globbed checkpoint files; it moves to checkpoint/utils.py so ds_to_universal.py can drop its copy, while zero_to_fp32.py keeps its own because it ships standalone.

Widening the pad would rename shards on disk and break existing checkpoints. Sorting on the read side leaves the format alone: the two orders agree whenever every rank field is two digits or fewer, so nothing written so far is reordered. I left out the pad-overflow assert suggested in the issue, since the read path no longer depends on the pad width.

This only bites at a tensor or model parallel degree of 100 or more, which is rare and is probably why it has sat since 2021. The failure is silent, though, and the fix is a sort key.

Verified in a CPU container:

• Three tests in tests/unit/checkpoint/test_ckpt_file_ordering.py. Reverting the module.py sort key on its own fails the pipeline test (assert 100 == 11), and reverting the engine.py one on its own fails the engine test. The third pins the sub-100 order and is green either way, so I reordered the list in place to confirm it can fail.
• unit/checkpoint/ and unit/runtime/pipe/: the same 10 failures and 154 errors as unmodified master, plus the three new passes. Those are the GPU and multi-process tests this box cannot run.
• pre-commit on the changed files: clean.

I have not run a real 100-way parallel job. The repro drives the actual naming, listing and loader index over 128 shards in a single process.

Fixes #1381

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: bc7b134f58

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

Add the required Signed-off-by trailer

This is a non-merge commit, but its commit message has no Signed-off-by trailer, so it violates the repository's mandatory commit requirement and may fail the corresponding CI/DCO check. Recreate the commit using --signoff with the configured Git identity.

AGENTS.md reference: AGENTS.md:L8-L8

Useful? React with 👍 / 👎.

[ebarkhordar]

The trailer is there. The commit message ends Signed-off-by: Ehsan Barkhordar <realbarkhordar@gmail.com>, and the DCO check on this same commit is green.