CUDA OoM Error for a model written using Lightnig

username: gerbersht
cooked: <p>I’ve been having this issue for a while when I try doubling the embed dim from 64 to 128 or increasing the batch size from 1 I keep getting this error. This is a language model with 1.2 million parameters. It works with 300k parameters with the same data but if I use a larger dataset with 200k samples it crashes. The current dataset has 4k samples. The samples are sequences of 50 proteins stacked on top of each other. I’d appreciate if someone could help out.</p>
<p>The error code is:</p>
<p><code>torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 150.00 MiB (GPU 0; 31.74 GiB total capacity; 21.32 GiB already allocated; 9.86 GiB free; 21.58 GiB allowed; 21.37 GiB reserved in total by PyTorch)</code></p>
<p>Here are the arguments:<br>
<code>nb_blocks=6, embed_dim=128, nb_heads=4, nb_epochs=20, warmup_steps=300, learning_rate=0.0001, check_val_every=1000, batch_size=1</code></p>
<p>I’m on a cluster setup using NVIDIA V100 PCIe 32 GB GPUs.</p>
<p>The Lightning settings are:</p>
<pre><code class="lang-auto">slurm_args = {
            "accelerator": "gpu",
            "devices": int(os.environ["SLURM_GPUS_ON_NODE"]),
            "num_nodes": int(os.environ["SLURM_NNODES"]),
            "strategy": "ddp",
            "precision": 16,
        }

...

    accelerator = "cuda" if torch.cuda.is_available() else "cpu"
    trainer_args = {
        "max_epochs": args.nb_epochs,
        "log_every_n_steps": LOGGING_STEPS,
        "val_check_interval": VAL_CHECK_STEPS,
        "logger": wandb_logger,
        "callbacks": callbacks,
        "accelerator": accelerator,
        **slurm_args,
    }
</code></pre>
<p>Things I have tried:</p>
<ol>
<li>Setting precision to 16 (down from 32): halved the memory overflow from 300MiB to 150MiB (At least some progress but didn’t solve the problem)</li>
<li><code>torch.cuda.set_per_process_memory_fraction(0.7)</code>: freed up space but this space didn’t get allocated to be used for the memory overflow, and I can’t figure out how to allocate this spare memory to be used for the above problem</li>
<li>Manually overriding allocations and moving weights, data, and the model around to CPU. Didn’t do anything as pretty much all of the initialization already gets initialized to the CPU.</li>
<li>Set <code>accelerator</code> in <code>slurm_args</code> to <code>cpu</code>. Solved the memory issue but each iteration went from taking 1.2s to 25s. In other words, it’s so slow it’s useless</li>
<li>Changed the strategy to <code>FSDPStrategy(cpu_offload=True)</code> from <code>ddp</code> and it didn’t change a single thing</li>
<li>Tried setting <code>PYTORCH_CUDA_ALLOC_CONF</code>, didn’t do anything</li>
<li>Set <code>accumulate_grad_batches</code> in the arguments with different values, no change</li>
</ol>

username: yf225
cooked: <p>Does it fail on the first iteration? If so, then the batch size + model size is too large for the GPU.</p>
<p>If it fails on subsequent iterations, you might want to try <a href="https://pytorch.org/docs/stable/torch_cuda_memory.html" class="inline-onebox" rel="noopener nofollow ugc">Understanding CUDA Memory Usage — PyTorch 2.5 documentation</a>, to see if there is gradual increase of peak memory over the course of training, and see what is the source of the increase.</p>

username: gerbersht
cooked: <p>It did fail on the first iteration and yes the model size was large (but that was the point) with a batch size of 1. I fixed it by checkpointing the attention blocks. So I can now increase the model size to about ~200M, which is more than enough.</p>