Restore full training state in checkpoint writer engines

[ai-yang]

FastCheckpointEngine and DecoupledCheckpointEngine can save complete ZeRO-3 training checkpoints but fail to restore them under PyTorch's default weights-only loader, which rejects DeepSpeed's own ZeroStageEnum.

Pass weights_only=False in both load() methods, matching the existing TorchCheckpointEngine behavior for trusted complete training checkpoints. The production change is limited to these two calls.

Fixes #8500.

Regression coverage

TestCheckpointWriterResume covers synchronous/asynchronous Python writers and legacy/ZIP serialization. It runs real ZeRO-3 training, saves step 1, and performs the next optimizer step to commit pending asynchronous writes. A fresh engine must restore the saved parameters and step count; its next Adam update must match uninterrupted training with zero tolerance. This checks the observable effect of restoring optimizer state, not only successful deserialization.

The asynchronous case uses a real subprocess. The only environment patch supplies the single-node launcher's CROSS_RANK / CROSS_SIZE values.

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

• New regression: 4 passed on CPU/Gloo and 4 passed on RTX 3090/CUDA/NCCL. Each pytest case executes with both one and two ranks, giving 8 distributed executions per backend.
• The four cases cover synchronous/asynchronous Python writers and legacy/ZIP serialization, including exact saved-parameter restoration, the restored training step count, and the next Adam update against uninterrupted training.
• Existing GPU regression: TestOtherOptimizerCheckpoint::test_checkpoint_fp32_optimizer: 1 passed with two ranks.
• All applicable pre-commit hooks on the three 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_other_optimizer.py::TestCheckpointWriterResume \
  --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_other_optimizer.py::TestCheckpointWriterResume \
  --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 with the complete pretrained Qwen3.5-0.8B language model: 24 decoder layers, 752,393,024 trainable language-model parameters, BF16, ZeRO-3, PyTorch AdamW, sequence length 32, microbatch 1 per GPU. Hardware/software: one/four RTX 3090 24 GiB GPUs, driver 580.173.02, PyTorch 2.12.1+cu126, Transformers 5.10.4, NCCL.

Synchronous/asynchronous writer × one/four GPUs gave four baseline deserialization failures and four successful fixed runs. All restored model parameters, the training step count, and all model parameters after the next update matched the uninterrupted reference exactly. These full-model GPU runs used ZIP serialization and excluded the vision branch. Current-base regression results are listed separately above.

Compatibility and scope

Full pickle deserialization requires self-produced or otherwise trusted checkpoints; weights_only=False must not be treated as safe for untrusted files. This restores the existing complete-checkpoint convention without adding a public API or changing the file format. An allowlist or a tensor/primitive-only format would require a separate compatibility design.

#6751 updated other load sites but did not cover these two calls. #7742 addresses asynchronous process reliability, a different failure mode. Multi-node, AIO/GDS, NVMe offload, long-run convergence, and the full DeepSpeed suite were not tested.

[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:35.263845Z	960efec	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: 960efec32a

ℹ️ 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 sign-off trailer

Commit 8f7c35c4db8f386ff9f742234159b9d44ad16b34 is a non-merge commit, but its message contains no Signed-off-by trailer, violating the repository's commit requirements; add the configured author's sign-off before merging.

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

Useful? React with 👍 / 👎.