Suspected memory leak during zero3 training. oom eventually after several checkpoint

Hi,

I am using zero3 with latest 0.9.2 to train 8gpus in one node, and from the stat, I see memory grow a lot after each checkpoint saving.

With checking the code, I find one suspected memory leak place,
https://github.com/microsoft/DeepSpeed/blob/master/deepspeed/runtime/engine.py#L3257
In _zero3_consolidated_16bit_state_dict, it would alloc [state_dict] local variable to collect params in rank0, and return it to parent function save_16bit_model, where it use self.checkpoint_engine.save to save it to disk.

So [state_dict] is a local variable, is it suitable to pass the param out to the parent function? would it cause some memory leak in this way?

Thx,
Lei