CPU memory usage leak because of calling backward

amogh112: <p>Hi<br>
<strong>Code</strong> : <a href="https://gist.github.com/amogh112/deecd6d388c43cc58c33ed9c659b4b4d" class="inline-onebox" rel="noopener nofollow ugc">train_dataleak.py · GitHub</a><br>
I observed that during training, things were fine until 5th epoch when the CPU usage suddenly shot up(see image for RAM usage).</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/0/1/01171cb6f0ba7f3bb65a3ee8d887d6557ee557be.png" alt="image" data-base62-sha1="9DZH0lz5ooXewFsPWPMTNBBlng" width="310" height="236"></p>
<p>I made sure that loss was detached before logging.<br>
I observed that this does not happen if I comment out the loss.backward call(<a href="https://gist.github.com/amogh112/deecd6d388c43cc58c33ed9c659b4b4d#file-train_dataleak-py-L181" rel="noopener nofollow ugc">Line 181</a>)<br>
Finally I was able to get around by collected garbage using <a href="https://gist.github.com/amogh112/deecd6d388c43cc58c33ed9c659b4b4d#file-train_dataleak-py-L230" rel="noopener nofollow ugc">gc.collect() after every 50 batches.</a><br>
But it still slows the epoch by a lot.</p>
<p>Some explanation for the code:<br>
<a href="https://gist.github.com/amogh112/deecd6d388c43cc58c33ed9c659b4b4d#file-train_dataleak-py-L252" rel="noopener nofollow ugc">Link to model code</a>: I added some linear layers to a base model(VLBert) that was training without any issues. I do a forward pass for the base model, take an embedding from the output and pass it through the linear layers to produce 1 output.</p>
<p>Then BCEWithLogits is then used to calculate the <a href="https://gist.github.com/amogh112/deecd6d388c43cc58c33ed9c659b4b4d#file-train_dataleak-py-L168" rel="noopener nofollow ugc">loss wrt labels for this output.</a><br>
Then backward is called.</p>
<p>The GPU(RTX 2080 TI) usage is not a problem, just the CPU usage. Printing out the tensors using gc.get_objects() as explained <a href="https://discuss.pytorch.org/t/how-to-debug-causes-of-gpu-memory-leaks/6741/3">here</a> shows there are no tensors on the CPU after every epoch.</p>
<p>What is the reason for garbage not being collected, any issues with the code?</p>

albanD: <p>Hi,</p>
<p>This is very surprising that it only happens at the 5th epoch. Do you do anything specific at that time? because the print you shows seems to indicate a very different behavior.</p>
<p>Also are you using multiprocessing here? Does removing that helps?</p>
<blockquote>
<p>What is the reason for garbage not being collected, any issues with the code?</p>
</blockquote>
<p>The python garbage collector does not run very often. And in particular, if you don’t create objects fast, it might actually never run.<br>
It is only useful to remove reference cycles so should not be needed in most cases.<br>
Did you managed to get any insight of which objects get collected when you run this? Could it be some of the logging logic that creates these reference cycles?</p>

amogh112: <p>Thanks for your response Alban.<br>
I do not do anything specific at this time. I am using 8 GPUs with apex which uses multiprocessing I believe:</p>
<blockquote>
<p>from apex.parallel import DistributedDataParallel as DDP<br>
model = DDP(model, delay_allreduce=True)<br>
optim = AdamW(model.parameters(), lr=args.learning_rate, correct_bias=False, eps=args.adam_epsilon)<br>
model, optim = amp.initialize(model, optim, opt_level=args.opt_level, loss_scale=args.loss_scale,<br>
keep_batchnorm_fp32=args.keep_batchnorm_fp32, verbosity=0)</p>
</blockquote>
<p>Since it takes a lot more time to run 5 epochs without using the 8 GPUs, I haven’t run it that long to see if its resolved without using DDP.<br>
Not sure if relevant and a source of reference cycles, I take a pre-trained base model and load its saved weights, and then initialise the weights for added layers.</p>
<blockquote>
<p>Did you managed to get any insight of which objects get collected when you run this? Could it be some of the logging logic that creates these reference cycles?</p>
</blockquote>
<p>In the <a href="https://gist.github.com/amogh112/deecd6d388c43cc58c33ed9c659b4b4d" rel="noopener nofollow ugc">code I shared</a>, I do not see any source of reference cycles.  I am simply logging losses with Tensorboard.</p>
<p>SInce removing .backward() does not lead to a leakage, I assumed that there are a few tensors on CPU that are in the computation graph. However upon printing out the tensors on all devices at different stages using <a href="https://gist.github.com/amogh112/b7146ebdd30d1d1e62ed4d5d44c71858" rel="noopener nofollow ugc">this function</a>, I found that there are some CPU tensors while batch is being processed in training (as received from dataloader before being moved to cuda in new variables eg - <code>imgs = data["imgs"].cuda()</code>). But printing out the tensors on CPU before the next training epoch, I did not find anything on CPU. Also number of tensors on CPU at start of each batch is the same.</p>

albanD: <p>I is hard to say what the issue is here as there are way too many things that come into play.<br>
I think you want to try and simplify your code and reduce your model to have something that reproduces this fast and without using so many complex features.</p>