Repository navigation
ZeRO-3: take the AutoTP shard from universal checkpoints converted without TP - #9069
Achyuthan-S wants to merge 1 commit into
Conversation
…thout TP ds_to_universal writes each state of a ZeRO-3 checkpoint saved without TP as a bare tensor, while every other path writes a dict. The ZeRO-3 loader resolved the AutoTP shard only for a dict, so loading such a checkpoint into an AutoTP run partitioned the whole weight and failed. _resolve_autotp_partition reads only the target parameter's metadata and the loaded tensor, so let the target decide. Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
There was a problem hiding this comment.
Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.
An organization owner can choose the repository organization as the billing source for reviews requested by members with a paid Copilot license. Learn about Copilot code review billing.
Codex Review Summary: ✅ CompletedThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
|
A real-model check of the same change: Qwen3-0.6B, bf16, 4× A100 (driver 595.71.05, torch 2.12.1+cu130, NCCL 2.29.7), DP4 → TP2×DP2 → DP4 under ZeRO-3, 20 steps per layout. Each change saves, runs
|
gpu-test-labeler
Bot
added
the
skip-gpu-test
A ZeRO-3 checkpoint saved without TP can't be resumed in an AutoTP + ZeRO-3 run through the universal checkpoint.
Why
ds_to_universalwrites each state of a ZeRO-3 checkpoint saved without TP as a bare tensor (merge_zero3_slices); every other path writes a dict with the tensor underparam.DeepSpeedZeroOptimizer_Stage3.load_hp_checkpoint_statetakes the AutoTP shard only when the file holds a dict. With a bare tensor it partitions the whole weight as if it were this rank's TP shard, and the copy into the flat buffer fails:_resolve_autotp_partitioncuts the shard using the target parameter's AutoTP metadata and the loaded tensor, and never reads the dict, so the check isn't needed.What
The target parameter alone decides whether to slice (
if param is not None:). Checkpoints converted before this keep loading, since the loader accepts both formats as before.Whether every conversion path should write the same format is left for later, as discussed with @delock: the loader would still have to accept bare tensors from checkpoints that were already converted.
No open PR touches this check (searched for
load_hp_checkpoint_stateand ZeRO-3 universal AutoTP; #8848 changes the converter, notstage3.py).Test
TestZero3UniversalCheckpointWithoutTpintest_autotp_uc_checkpoint.pytrains 4 steps under ZeRO-3 without TP, converts, and loads into ZeRO-3 at TP1 and at AutoTP TP2. It compares the fp32 weights, both Adam moments and the step counts with the saved state bit for bit, since conversion and restore only copy. On master the TP2 case fails with the error above; the TP1 case passes before and after._affine_resume_enginegains azero_stageargument (default 1) so the test reuses the existing resume model and helpers.Validation
pytest --forked tests/unit/checkpoint/test_autotp_uc_checkpoint.py tests/unit/checkpoint/test_universal_checkpoint.py tests/unit/checkpoint/test_autoep_affine_universal.py: 217 passed with the fix. With only thestage3.pychange reverted, the new TP2 case fails with the error above and the TP1 case passes.DS_ACCELERATOR=cpu LOCAL_SIZE=4): the same result for the new test, before and after. The FusedAdam tests in these files can't build the op on macOS, so they ran on GPU only.ds_to_universaland loads in the new layout, then compares every fp32 weight and both Adam moments with the saved state bit for bit. On master the TP2×DP2 load fails withstart (187907840) + length (1048576) exceeds dimension size (187924480); with the fix both changes restore all 930 tensors bit for bit.Checked and not changed
universal_checkpoint.load_hp_checkpoint_state) readsckpt_dict[PARAM]unconditionally, but stage 1/2 conversion always writes the dict, and a ZeRO-3 universal checkpoint doesn't reach that loader today: a ZeRO-1 engine looks formp_rank_00_model_states.pt, while the converted ZeRO-3 folder haszero_pp_rank_0_mp_rank_00_model_states.pt, so the load stops instate_dict_factory.check_ckpt_list.Validated: ZeRO-3 without offload; fp32 (the unit test) and bf16 (Qwen3-0.6B); TP1 → TP1, TP1 → TP2, and DP4 → TP2×DP2 → DP4. Not validated: CPU or NVMe offload, fp16, pipeline parallelism.