RuntimeError: CUDA OOM after training on a few batches

reubenwong97: <p>Hi everyone,</p>
<p>I am trying to fine-tune a T5 based model with reinforcement learning. I am currently using the T5 model to paraphrase a relatively short query before feeding it into my environment and retrieving my reward, and then utilize this reward to compute my losses for training.</p>
<p>I have scoured the forums and have tried to incorporate best practices like deleting variables once out of scope and storing my losses with <code>loss.item()</code>. Even more strangely, I am able to train for a couple of batching with ever increasing memory usage as tracked by <code>nvidia-smi</code>. I have even reduced by batch size to 1 but to no avail. I was able to reach the same number of batches trained with a higher batch size of 4 before meeting failure.</p>
<p>I am performing the training on an RTX 2080 ti with 11gb of RAM.</p>
<p>Does anyone know of potential causes of such behaviour? Thank you for taking the time to read this post.</p>
<p>Edit (include code):<br>
<strong>Main Training Loop</strong></p>
<pre><code class="lang-auto">for epoch in range(epochs):
        acc_loss = 0.0
        counter = 0
        for i, batch in enumerate(train_dataloader):
            queries = batch[
                "question"
            ]  # pass into step if using relative reward
            contexts = batch["context"]
            answers = batch["answer"]

            del batch
            output = agent.forward(queries)
            del queries

            reformulations = agent.decode_batch_sequence(
                output.sequences.detach().cpu(), skip_special_tokens=True
            )

            rewards = world_model.step(reformulations, contexts, answers)

            del contexts, answers, reformulations

            rewards = torch.tensor(rewards, dtype=torch.float16).to(
                agent.device
            )
            gen_probs, probs = agent.probabilities(output)
            loss = agent.policy_loss(rewards, gen_probs, probs)

            agent.optim.zero_grad()
            loss.backward()
            agent.optim.step()

            loss = loss.detach().item()
            wandb.log({"policy_loss": loss})

            acc_loss += float(loss)
            del loss
            counter += 1

        epoch_loss = acc_loss / counter
        wandb.log({"epoch_loss": epoch_loss})

        if epoch_loss &lt; best_loss:
            best_loss = epoch_loss
            agent.transformer.save_pretrained(args.save_path)

</code></pre>
<p><strong>Error</strong>:</p>
<pre><code class="lang-auto">RuntimeError: CUDA out of memory. Tried to allocate 2.00 MiB already allocated; 2.44 MiB free; 9.68 GiB reserved in total by PyTorch)
</code></pre>
<p><strong>Further Updates</strong>:<br>
After using the <code>debug_memory</code> code snippet found <a href="https://forum.pyro.ai/t/a-clever-trick-to-debug-tensor-memory/556" rel="noopener nofollow ugc">here</a>, it seems like my tensors are not being cleared after the optim.step() call, but rather freed when the computation of the loss for the next batch starts.</p>

ptrblck: <p>The intermediate forward activations should have been already freed during the <code>backward()</code> call, however the gradients would be set to zero in <code>optimizer.zero_grad()</code> and not removed. To save more memory, pass <code>set_to_none=True</code> to <code>zero_grad</code>, which would lower the memory footprint a bit more.<br>
Also, based on your description it seems that the OOM issue is raised in a specific iteration regardless of the batch size, which could indicate a larger memory requirement for this iteration. Are you using variable input shapes and could this particular batch be larger?</p>

reubenwong97: <p>Hi <a class="mention" href="/u/ptrblck">@ptrblck</a> ,</p>
<p>Thank you for taking the time to read my post and reply! You are right that a particular batch was larger than others and it caused the program to crash there. To resolve that I curbed the max length of my generated sequences.</p>
<p>However, when I tracked the tensors in my program with the <code>debug_memory</code> snippet, it does seem like the intermediate tensors are still kept until I compute the loss for the next batch, which I believe is strange?</p>
<p>But with regards to the memory issue, you are entirely spot on! Thank you for your help!</p>