zero_to_fp32.py drops buffers

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.