zero_to_fp32.py drops buffers

[Mindful]

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.

[stas00]

You're correct to report that buffers aren't being treated, @Mindful

I made sure I did it right in _zero3_consolidated_fp16_state_dict(self)

https://github.com/microsoft/DeepSpeed/blob/96eb5b12e328f9398d57fb577ee088cfcfb5444f/deepspeed/runtime/engine.py#L1953

it does handle buffers correctly, so at the very least you should be able to save the fp16 version of the params and buffers until this is sorted out.

---

back to the issue at hand:

The params map is saved here:

https://github.com/microsoft/DeepSpeed/blob/96eb5b12e328f9398d57fb577ee088cfcfb5444f/deepspeed/runtime/engine.py#L1925

But honestly I'm not quite sure where the buffers are saved in the optimizer checkpoint.

Probably should look somewhere around this starting point:

https://github.com/microsoft/DeepSpeed/blob/96eb5b12e328f9398d57fb577ee088cfcfb5444f/deepspeed/runtime/zero/stage3.py#L3061

Is it possible that buffers aren't partitioned by default?

I saw several situations where this was the case, but didn't have a chance to debug it.

I'm quite busy at the moment so won't be able to look at it for a while, so I'm hoping someone else will have a chance to investigate how the buffers are stored.

And, of course, DeepSpeed devs will probably be able to solve this in no time.

[stas00]

@Mindful, please try this PR: #1146

You will have to load and save the checkpoint or create a fresh one, since I added new metadata that needs to be saved there.

Please comment in the PR if you run into any issues.

Thank you!