[QUESTION] How to manually do states partition and all-gather communication during training

Hi everyone! I am running deepspeed with multi-GPUs. I want to store and use some extra information EXTRA_STATES related to the model states during training, which is quite huge and takes up a lot of memory. The thing I wanna optimize is to eliminate redundant memory. Specifically, due to the exactly same EXTRA_STATES retained in each of the multiple processes, there are quite lots of memory are wasted, and what I am thinking is whether I can do some partitioning and all-gather of the EXTRA_STATES.

I think perhaps this is similar to what's happening to the optimizer states and what ZeRO stage-1 does. So I am wondering if there is backend communication interface in DeepSpeed to manually handle this EXTRA_STATES partitioning and all-gather, or if there are some tricks by which I can take advantages of the existing DeepSpeedZeroOptimizer.

Start by examining the DeepSpeedZeroOptimizer and the backend communication interface mentioned in the question. Determine whether these provide a supported way to partition and all-gather EXTRA_STATES during multi-GPU training, and document the applicable approach or clarify that the capability is unavailable.