PyTorch not working smoothly with DeepSpeed ZeRO Stage-3

ConnollyLeon: <p>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.</p>
<p>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.</p>
<p>Could anyone tell us how I can remove this param ref to help reduce memory?</p>

ConnollyLeon: <p><a class="mention" href="/u/alband">@albanD</a> Could you please help me out?</p>

albanD: <p>Hi,</p>
<p>In general, you should never use <code>.data</code> <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=10" title=":slight_smile:" class="emoji" alt=":slight_smile:"><br>
Could you give more details on what you’re trying to do here? Because it is expected that the autograd saves (a lot) of things in the graph.</p>

ConnollyLeon: <p>I am trying to save memory as much as possible.<br>
The implementation actually comes from DeepSpeed. this line set torch.data to a torch.ones(1) tensor. <a href="https://github.com/microsoft/DeepSpeed/blob/master/deepspeed/runtime/zero/partition_parameters.py#L736" class="inline-onebox" rel="noopener nofollow ugc">DeepSpeed/partition_parameters.py at master · microsoft/DeepSpeed · GitHub</a></p>
<p>However, the parameters memory didn’t release correctly. And I found out that if a parameter Tensor has its corresponding intermediate results, the parameters memory won’t be released.</p>
<p>The parameter data will only be gathered when they are needed. Like gather in pre_sub_module_backward_function here:<a href="https://github.com/microsoft/DeepSpeed/blob/master/deepspeed/runtime/zero/stage3.py#L1590" class="inline-onebox" rel="noopener nofollow ugc">DeepSpeed/stage3.py at master · microsoft/DeepSpeed · GitHub</a></p>
<p>But If the parameters memory not release correctly, here gathering won’t consume any other memory. In other word, I reset the param.data to parameter Tensor, but the memory_allocated  has no change.</p>

albanD: <aside class="quote no-group" data-username="ConnollyLeon" data-post="4" data-topic="134660">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/connollyleon/48/36864_2.png" class="avatar"> ConnollyLeon:</div>
<blockquote>
<p>However, the parameters memory didn’t release correctly.</p>
</blockquote>
</aside>
<p>I wouldn’t say that. You do remove the parameter successfully. The problem is that other part of the code are also using this Tensor and so the Tensor cannot be freed. This is actually expected.</p>
<p>If you want to play tricks with what the autograd saves for backward, you have hooks for that: <a href="https://pytorch.org/docs/1.10.0/notes/autograd.html#hooks-for-saved-tensors" class="inline-onebox">Autograd mechanics — PyTorch 1.10.0 documentation</a>. There is also a tutorial for that.</p>

ConnollyLeon: <p>Thanks, I will try it out later</p>