Fix ZeRO-3 checkpoint tag selection during weight reconstruction

[ai-yang]

When a non-MoE ZeRO-3 model loads an explicit checkpoint tag with load_optimizer_states=False, FP32 reconstruction currently resolves latest again. As a result, a request for earlier can return its metadata while silently installing later weights, or fail if there is no latest file.

Pass the already resolved tag to get_fp32_state_dict_from_zero_checkpoint() so the reconstructed parameters and metadata come from the same checkpoint. The production change is one call argument.

Fixes #8499.

Regression coverage

TestZeROCheckpointTag trains and saves distinct parameter snapshots, then checks restored parameters against the selected snapshot with zero tolerance. It covers explicit older tags with/without latest, load_module_only=True/False, default latest selection, and full optimizer-state restoration as a control. Tests use real engine training and checkpoint files, with no loader mocks.

Current-base checks on 71d316d608a56af2fcc27b84b14cf854b7052eff plus this fix (2026-09-13):

• New regression: 6 passed on CPU/Gloo and 6 passed on RTX 3090/CUDA/NCCL. Each pytest case executes with both one and two ranks, giving 12 distributed executions per backend.
• Coverage includes explicit older tags with/without latest, both load_module_only values, omitted-tag selection, and a full optimizer-state restore control.
• Existing GPU regressions: TestZeROCheckpoint::test_not_load_optimizer_state[3-False-Adam] and TestZeROCheckpoint::test_load_module_only[3]: 2 passed, each with two ranks.
• All applicable pre-commit hooks on the two modified files and git diff --check passed.

Run the new regression from the source checkout with its test dependencies installed:

export PYTHONPATH="$PWD:$PWD/tests"
export OMP_NUM_THREADS=1 LOCAL_SIZE=2 PYTEST_DISABLE_PLUGIN_AUTOLOAD=1
DS_ACCELERATOR=cpu python -m pytest -p pytest_forked --forked \
  tests/unit/checkpoint/test_zero_optimizer.py::TestZeROCheckpointTag \
  --torch_ver=2.12.1+cu126 --cuda_ver=12.6
DS_ACCELERATOR=cuda CUDA_VISIBLE_DEVICES=0,1 python -m pytest -p pytest_forked --forked \
  tests/unit/checkpoint/test_zero_optimizer.py::TestZeROCheckpointTag \
  --torch_ver=2.12.1+cu126 --cuda_ver=12.6

The explicit version options describe the tested PyTorch installation; adjust them for another environment.

Full-model integration evidence

The same production fix was previously tested on base 29d0abbc21f11da806ca14970fa8b3ddb757a152 using the complete pretrained Qwen3.5-0.8B language model, all 24 decoder layers and all 752,393,024 language-model parameters trainable. Configuration: BF16, ZeRO-3, PyTorch AdamW, sequence length 32, microbatch 1 per GPU; one and four RTX 3090 24 GiB GPUs, driver 580.173.02, PyTorch 2.12.1+cu126, Transformers 5.10.4, NCCL.

Both baseline runs reproduced the wrong-tag load. With the fix, every restored model parameter matched the selected checkpoint exactly, and a subsequent training step updated parameters successfully. This was a short integration check, not a convergence or throughput benchmark; vision, multi-node training, and NVMe offload were outside its scope. Current-base regression results are listed separately above.

Related changes

#3116 and #4089 address adjacent checkpoint behavior but do not pass this engine argument. Open PR #8378 modifies the same call for frozen-parameter dtype handling; it addresses a different contract, and both arguments may need to be retained when rebasing. No public API or checkpoint format changes are introduced here.

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-09-13T05:56:33.585982Z	1fbcf66	PR opened

ℹ️ About Codex in GitHub

Your team has set up Codex to 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" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 1fbcf669b5

ℹ️ About Codex in GitHub

Your team has set up Codex to 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 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Add the required Signed-off-by trailer

This non-merge commit has no Signed-off-by trailer, so it does not satisfy the repository's commit requirements. Please add the trailer using the configured Git name and email before merging.

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

Useful? React with 👍 / 👎.

[chatgpt-codex-connector]

Normalize the checkpoint tag before reconstruction

When a checkpoint was saved with a numeric tag and the caller later uses the same numeric value in load_checkpoint, this ZeRO-3 path now passes the integer directly to get_fp32_state_dict_from_zero_checkpoint, whose os.path.join(checkpoint_dir, tag) raises TypeError. This previously worked with the default save_latest=True because reconstruction re-read the stringified tag from latest, and save_checkpoint explicitly supports such values by applying str(tag); normalize the resolved load tag similarly before forwarding it.

Useful? React with 👍 / 👎.