I notice that when using zero stage-3 in deepspeed to train model without any recomputing (checkpointing) methods, the parameters of the model cannot correctly released. I have already release param.data with my ZeRO, which set param.data to torch.Tensor([1]). However, the memory consumption doesn’t decrease, which means param.data still remains in the memory.
I think the problem occurs in autograd module, which may possibly create a weak ref to “param.data” when there are some corresponding intermediate results ( a.k.a. intermediate activations) in GPU memory.
Could anyone tell us how I can remove this param ref to help reduce memory?