CUDA out of memory - nV RTX 3060 laptop GPU - time to move to Colab?

[bjpcjp]

Bug description

Consistent CUDA memory issues when running compute_grpo_loss in book section 6.10 ("putting it all together).

What operating system are you using?

Linux

Where do you run your code?

Local (laptop, desktop)

Environment

Python version: 3.12.3 (main, Mar 23 2026, 19:04:32) [GCC 13.3.0]
PyTorch version: 2.11.0+cu130
CUDA available: True
CUDA device: NVIDIA GeForce RTX 3060 Laptop GPU
CUDA version: 13.0
Device count: 1

[bjpcjp]

after using "!export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True":

[rasbt]

Could you try smaller numbers of rollouts and max_new_tokens? But yeah, the 6 GB RAM in this card would be too low otherwise and a Lightning AI or Google Colab would be more feasible options. It would perhaps be a good idea to add a table with minimum memory requirements with different settings.

[bjpcjp]

Interesting experiment - I reset the kernel & re-ran the entire workbook past compute_grpo_loss, then iterated on generating stats.

Test 1: max_new_tokens 256 --> 128. This test passed w/o error.
Test 2: max_new_tokens 128 --> 192. This returns an OOM error.
Test 3: max_new_tokens 192 --> 128. OOM still occurs.

It looks like num_rollouts has a minimum of 2 (setting to 1 returns an assertion error).

I agree, this looks like I'm destined to visit LAI or Colab.

[danielpenchev98]

Hmm, i faced the very same problem but with a bigger GPU - RTX 3090. The way I tried to fix it is by calculating the logprob of and immediately call the backwards() for each rollout separately:

    total_pg_loss = 0.0
    for i, d in enumerate(rollout_data):
        logp = sequence_logprob(model, d["token_ids"], d["prompt_len"])
        # Divide by num_rollouts so accumulated gradient == mean over rollouts
        pg_loss_i = -(advantages[i].detach() * logp) / num_rollouts
        pg_loss_i.backward()
        total_pg_loss += pg_loss_i.item()

If I'm not mistaken, in that way, we aren't keeping all rollout computation graphs in the VRAM at the same time.