deepspeedai / deepspeedai/DeepSpeed
zero_to_fp32.py drops buffers
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 43.1k
- Forks
- 5k
- Avg merge
- 4d 15h
- Merged PRs (30d)
- 112
Description
At least for deepspeed 0.3.16, zero_to_fp32.py seems to drop content from state dicts that was registered with register_buffer. It looks like the current script gets states out of the optimizer, so maybe it would make sense that any state that isn't an optimizable model parameter gets dropped? The result is that running zero_to_fp32.py on models with buffers results in a new state dict that has fewer entries and consequently can't be used for loading.
I can try and create a minimal reproduction if it would help, although it may take me some time to get to that. My hope is that the issue is pretty straightforward to reproduce.
Also tagging @stas00, as requested.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start by inspecting zero_to_fp32.py and how it reconstructs model state from optimizer state, then reproduce the issue with a model using register_buffer under DeepSpeed 0.3.16. Done means the converted state dict retains registered buffers and can be loaded without losing entries.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 56/100