How to manipulate model weights within a ZeRO 3 training process?

What I want to do is pretty straightforward - average the current model weights after a specified training steps with previous saved checkpoints (or torch model, whichever it works better with deepspeed). One problem I've found so far is directly accessing the model state_dict can lead to tensors of size 0, which is also mentioned in deepspeed official doc caused by model sharding.
I wonder how can I achieve this in an elegant way? I was told that even if I managed to gather all the tensors and do the average operation, I also need to cope with the fp32 model parameters maintained in the optimizer in additional to replacing the model parameters in fp16?
Any insights on this issue would be greatly appreciated. One important note is that I do need to realize this in the ZeRO 3 setting.

Start with the ZeRO 3 documentation on model sharding and the behavior of model state_dict tensors. Trace how saved checkpoints, gathered parameters, and the optimizer's fp32 parameters are handled during training. Done means identifying and documenting an elegant supported way to average weights at specified steps while keeping model and optimizer parameters consistent.