Got CUDA OOM Error after 8 Hours of Training

username: weiqis
cooked: <p>Hi, I’m training a large CLM on 8 GPUs using FSDP with gradient checkpointing. When training a CLM, I concatenated all the tokenized examples together and then split them into blocks with exactly the same length. The training process ran smoothly for 8 hours, and then it was terminated due to a CUDA OOM error.</p>
<p>I’m trying to resume the training from the latest checkpoint. But now I get the CUDA OOM error much sooner, and it happened totally randomly. Sometimes the error happened just at the second step after resuming the training. While sometimes it happened at the 25 step, resuming from the same checkpoint. But after resuming from the latest checkpoint, the training could never run for more than 30 min.</p>
<p>My questions are:</p>
<ol>
<li>Given all my training examples have exactly the same length (and I don’t do eval during training), shouldn’t the memory consumption be identical from step to step? Why the training process can run smoothly for 8 hours and then suddenly it got CUDA OOM error?</li>
<li>After resuming from the latest checkpoint, why do I get the CUDA OOM error much sooner? And why does it happen randomly, even if I’m using the same data, same hyperparameters and same seed?</li>
</ol>
<p>I tried to explicitly call Python’s garbage collector and also clean cuda cache in each training step:</p>
<pre><code class="lang-auto">class CLMTrainer(transformers.Trainer):
    def training_step(self, model: nn.Module, inputs) -&gt; torch.Tensor:
        #gc.collect()
        #torch.cuda.empty_cache()
        res = super().training_step(model, inputs)
        gc.collect()
        torch.cuda.empty_cache()
        return res
</code></pre>
<p>But it doesn’t help.</p>
<p>My training env:</p>
<ul>
<li>torch: 2.0.0+cu118</li>
<li>cuda: 11.8</li>
<li>transformers: 4.48.1</li>
</ul>
<p>Below are the training curve and GPU memory consumption plot in the first 8 hours of training (before getting the CUDA OOM error). Everything looks fine and I didn’t see any overall increase of GPU memory.<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/8/1/81aaeae350e8dd144264161dce03ee14083ec3bc.png" data-download-href="https://discuss.pytorch.org/uploads/default/81aaeae350e8dd144264161dce03ee14083ec3bc" title="W&amp;B Chart 4_26_2023, 12_12_28 PM"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/8/1/81aaeae350e8dd144264161dce03ee14083ec3bc_2_690x362.png" alt="W&amp;B Chart 4_26_2023, 12_12_28 PM" data-base62-sha1="iv5Mx7oDigpk7UkT4bHE0mjYV9q" width="690" height="362" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/8/1/81aaeae350e8dd144264161dce03ee14083ec3bc_2_690x362.png, https://discuss.pytorch.org/uploads/default/optimized/3X/8/1/81aaeae350e8dd144264161dce03ee14083ec3bc_2_1035x543.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/8/1/81aaeae350e8dd144264161dce03ee14083ec3bc_2_1380x724.png 2x" data-dominant-color="FCFDFE"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">W&amp;B Chart 4_26_2023, 12_12_28 PM</span><span class="informations">5056×2656 487 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div><br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/d/f/df558f4269c2983dcff470d16fc3efcaf4a1c7b8.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/df558f4269c2983dcff470d16fc3efcaf4a1c7b8" title="W&amp;B Chart 4_26_2023, 12_12_45 PM"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/d/f/df558f4269c2983dcff470d16fc3efcaf4a1c7b8_2_690x362.jpeg" alt="W&amp;B Chart 4_26_2023, 12_12_45 PM" data-base62-sha1="vRHMLosaKRYUkBzm21K8IsKuUL6" width="690" height="362" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/d/f/df558f4269c2983dcff470d16fc3efcaf4a1c7b8_2_690x362.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/d/f/df558f4269c2983dcff470d16fc3efcaf4a1c7b8_2_1035x543.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/d/f/df558f4269c2983dcff470d16fc3efcaf4a1c7b8_2_1380x724.jpeg 2x" data-dominant-color="F6F3F1"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">W&amp;B Chart 4_26_2023, 12_12_45 PM</span><span class="informations">1920×1009 66.1 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>Any commands or suggestions would be appreciated. Thanks in advance!</p>

username: aditya_nanda_kishore
cooked: <p>Hi Weiqis,<br>
I am facing a similar error. I am currently training a large causal model using Trainer and FSDP. Training runs smoothly for a few steps and then runs into OOM (All the training examples have the same sequence length as well). I would like to know if you debugged this.</p>