Fix ZeRO checkpoint DP merge on scalar optimizer state - #8396
Fix ZeRO checkpoint DP merge on scalar optimizer state#8396Kayvan-Zahiri wants to merge 1 commit into
Conversation
merge_state concatenated every tensor it met, so reshaping a ZeRO 1/2 checkpoint to a smaller dp_degree died on torch's 0-dim step counter with "zero-dimensional tensor (at position 0) cannot be concatenated". Return replicated scalars unchanged, matching the value.dim() > 0 guard in stage_1_and_2.py. merge_state_dict also iterated only dict_b, silently dropping keys held by dict_a alone, and passed [str(key)] instead of the accumulated path, so the diagnostic printed the leaf key rather than the nested one. Merge from dict_a, then add dict_b's exclusive keys, and propagate key_list. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01DQughgr95y4B9H1jfaQH8o Signed-off-by: kzahiri1 <kzahiri@dons.usfca.edu>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 35e294e035
ℹ️ 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".
| for key, value in dict_b.items(): | ||
| if key in dict_a.keys(): | ||
| merged_dict[key] = merge_state(dict_a[key], dict_b[key], [str(key)]) | ||
| for key, value in dict_a.items(): |
There was a problem hiding this comment.
This is a non-merge commit, but its message has no Signed-off-by trailer, so it violates the repository's mandatory commit policy. Recreate the commit using --signoff before merging.
AGENTS.md reference: AGENTS.md:L8-L8
Useful? React with 👍 / 👎.
|
The sign-off is already there. The commit message on
Nothing to recreate here, but happy to reorder the trailers so |
1 participant
merge_stateconcatenates every tensor it meets. Withelastic_checkpoint=False(the default) ZeRO 1/2 savesself.optimizer.state_dict()verbatim (stage_1_and_2.py:2628), so per-parameter state carries torch'sstep, a 0-dim tensor for Adam/AdamW. Reshaping that checkpoint to a smallerdp_degreesends the shards throughmerge_stateand dies:raised at
reshape_utils.py:111fromZeROCheckpoint.get_state_for_rank. Replicated scalars are now returned unchanged, mirroring thevalue.dim() > 0guard instage_1_and_2.pyand thestate_name != "step"skip in_strip_tensor_paddings.Same file:
merge_state_dictiterated onlydict_b, silently dropping keys held bydict_aalone, and passed[str(key)]rather than the accumulated path, so the mismatch diagnostic printedgroupsinstead ofopt.groups.merge_statehad no test coverage. The three added tests fail on master with exactly those symptoms.test_reshape_checkpoint.py+test_convert_checkpoint.py: 12 passed / 1 skipped before, 15 passed / 1 skipped after.pre-commit run --filespasses on both changed files.🤖 Generated with Claude Code