Repository navigation
[Dev] Preserve HDO step across DP-reshardable resume - #7629
Merged
Merged
Conversation
When DistributedOptimizer wraps a HybridDeviceOptimizer and loads a DP-reshardable checkpoint, the restored scalar `step` was written into `self.optimizer.state` and then overwritten by the local template value while the per-parameter state was copied in. Write it into the state dict that is loaded instead, and skip the per-parameter `step` in the HDO state-copy paths. Add a dist-ckpt unit test that resumes an HDO optimizer across a DP resharding and checks the step is preserved. Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011KkXwkg28vJbR2T64TETxM Signed-off-by: Hejian Sang <sanghj0923@gmail.com>
cspades
approved these changes
Sep 24, 2026
Member
|
/ok to test f685f98 |
Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011KkXwkg28vJbR2T64TETxM Signed-off-by: Hejian Sang <sanghj0923@gmail.com>
Author
Member
|
/ok to test a25b5b0 |
hxbai
approved these changes
Sep 28, 2026
This branch was previously deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
When
DistributedOptimizerwraps aHybridDeviceOptimizer(CPU optimizer offload) and loads a DP-reshardable checkpoint, the restored scalarstepis written intoself.optimizer.state, and the per-parameter state copy then overwrites it with the local template value. After resuming, Adam's bias correction restarts from the template step instead of the checkpointed one.This PR writes the restored step into the state dict that is actually loaded (
state_dict_state), and skips the per-parameterstepin the HDO state-copy paths, so the value restored fromparam_groupsis the one that sticks.Tests
New
TestDistributedOptimizer::test_hdo_dp_reshardable_preserves_stepintests/unit_tests/dist_checkpointing/test_optimizer.py: it steps an HDO optimizer 3 times, saves withdistrib_optim_sharding_type=dp_reshardable, loads into a fresh optimizer that has stepped once, and asserts the step is 3 in bothparam_groupsand the per-param state.setup_model_and_optimizerintests/unit_tests/dist_checkpointing/utils.pygains anoptimizer_config_kwargspass-through.Ran on 4× B200 with
torchrun --nproc_per_node=4,-k "hdo or dp_reshardable or HybridDevice or offload":dev@0e704edb1: 5 passed, 2 failedThe 2 failures (
test_nonreshardable_optimizer_save_load[...-dp_reshardable-None]) are identical ondevand unrelated to this change.Credit
Original fix by Nan Jiang (@nanjiangwill), developed in our internal fork (radixark#105). Upstreamed by me with the author preserved.
🤖 Generated with Claude Code
https://claude.ai/code/session_011KkXwkg28vJbR2T64TETxM