Fixing CUDA Out of Memory During VLA Training

Read the error line before you touch anything

CUDA out of memory is not one error. PyTorch's allocator prints a five-part sentence and each part points at a different fix, so reading it is the difference between one config change and an afternoon of bisection. The wording comes from CUDACachingAllocator.cpp. Here it is, from a fine-tuning run that died on an 80 GB card.

torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 2.34 GiB.
GPU 0 has a total capacity of 79.15 GiB of which 812.00 MiB is free.
Process 21874 has 78.35 GiB memory in use. Of the allocated memory
71.02 GiB is allocated by PyTorch, and 5.61 GiB is reserved by PyTorch
but unallocated. If reserved but unallocated memory is large try setting
PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.

First check: is it the GPU at all?

Robot pipelines decode video on the CPU, in worker processes. A good share of reported out-of-memory failures on a LeRobot dataset never touched VRAM at all.

Where the memory goes

Hugging Face's GPU memory anatomy page gives the per-parameter accounting for mixed-precision AdamW. It tells you which failures batch size can fix.

Two consequences. A full fine-tune of a 3 B vision-language-action model does not fit on 24 GB at any batch size, which is why GR00T N1.7 and Pi0.5 are 80 GB-tier here while SmolVLA and ACT are not. And only the last row responds to batch size, so a run that dies at batch 1 needs a structural change.

Published figures agree. openpi states above 22.5 GB for a LoRA fine-tune on an RTX 4090 and above 70 GB for a full one on an A100 80 GB or H100. NVIDIA's guide states a 40 GB minimum and puts the default recipe, which tunes the projector and diffusion head rather than the language backbone, under about 35 GB. The gap between 35 and 70 is entirely which modules carry gradients.

Pull the levers in this order

Order matters because the early levers are free and the late ones are not. Steps 1 to 5 cost time or nothing. From step 6 you are trading model quality, so do not start there just because the flag is easy to type.

1. 1
   Free the card you already have

   A dead notebook kernel or a crashed run will sit on 20 GB. On a rented pod this is the commonest reason a run OOMs at step 0 after working yesterday.

   bash

   nvidia-smi --query-compute-apps=pid,process_name,used_memory --format=csv
   kill -9 <pid>   # only processes you recognise as yours

2. 2
   Check the dtype default

   lerobot's Pi0.5 config ships dtype: str = "float32". Training a 3 B model in fp32 doubles weights and activations for no benefit on Ampere or newer. Check the default before you set anything: lerobot's GR00T config already ships use_bf16 = True, so there is nothing to turn on there, while the accelerator's own mixed_precision still defaults to "no".

   bash

   lerobot-train --policy.type=pi05 --policy.dtype=bfloat16 ...   # default is "float32"
   lerobot-train --accelerator.mixed_precision=bf16 ...            # default is "no"
   
   # lerobot's groot policy is already use_bf16=True by default:
   # params stay FP32 and compute runs under BF16 autocast.

3. 3
   Halve the batch size

   This tells you whether you are activation-bound or parameter-bound. If 32 dies and 16 runs, activations dominate. If batch 1 dies too, go to step 6.

   bash

   python gr00t/experiment/launch_finetune.py --global-batch-size 16 ...
   lerobot-train --batch_size=16 ...

4. 4
   Accumulate to restore the effective batch

   Buys back the optimisation behaviour you gave away, at the cost of wall-clock time. GR00T's guide works an example: 4 GPUs, 8 accumulation steps, per-GPU batch 8 gives an effective global batch of 256.

   bash

   # GR00T: global_batch_size is PRE-accumulation, and the class warns you
   python gr00t/experiment/launch_finetune.py \
     --global-batch-size 16 --gradient-accumulation-steps 4 ...
   
   # lerobot main: it lives under the accelerator config
   lerobot-train --batch_size=16 --accelerator.gradient_accumulation.steps=4 ...

5. 5
   Turn on gradient checkpointing

   Drop the cached activations and recompute them in the backward pass. Chen et al. 2016 took a 1000-layer residual network from 48 GB to 7 GB for about 30 percent extra running time; Hugging Face quotes about 20 percent for transformer fine-tunes.

   bash

   lerobot-train --policy.type=pi05 --policy.gradient_checkpointing=true ...
   
   # raw PyTorch: use_reentrant has no default any more
   from torch.utils.checkpoint import checkpoint
   out = checkpoint(block, x, use_reentrant=False)

6. 6
   Freeze what you do not need to train

   Here the trade starts. Both stacks freeze most of the backbone already, and un-freezing pushes GR00T past 80 GB. lerobot names the price: the frozen-VLM variant is less memory, at some cost in success rate.

   bash

   # GR00T defaults: tune_llm=False, tune_visual=False,
   # tune_projector=True, tune_diffusion_model=True.
   # Adding --tune-llm or --tune-visual is what breaks a 40 GB card.
   
   lerobot-train --policy.type=pi05 \
     --policy.freeze_vision_encoder=true --policy.train_expert_only=true ...

7. 7
   Switch to LoRA, if your version has it

   Adapters cut optimizer state and gradients to the adapter parameters. Two version traps: in lerobot main PEFT is a top-level config, not a policy flag, so it is --peft.* and it needs the peft extra installed; and on the Isaac-GR00T side the N1.5 script had --lora-rank while the N1.7 FinetuneConfig has no LoRA fields at all.

   bash

   # lerobot main: PeftConfig is top-level (method_type defaults to LORA, r to 16).
   # Passing any --peft.* field is what turns PEFT on; there is no --policy.use_peft.
   lerobot-train --peft.method_type=LORA --peft.r=16 --peft.lora_alpha=32 ...
   
   # Isaac-GR00T, N1.5 branch only
   python scripts/gr00t_finetune.py --lora-rank 64 --lora-alpha 128 ...

8. 8
   Shrink the inputs, knowing the cost

   Fewer cameras, smaller images and shorter chunks cut activations hard and change the task the model solves. GR00T N1.7 gates resolution behind two flags set together.

   bash

   python gr00t/experiment/launch_finetune.py \
     --shortest-image-edge 224 --crop-fraction 0.9 ...
   
   lerobot-train --policy.chunk_size=25 --policy.n_action_steps=25 ...

9. 9
   Tune the allocator, which costs nothing

   If the error showed several GB reserved-but-unallocated, this is your fix and none of the steps above were needed. In PyTorch 2.13 the variable is PYTORCH_ALLOC_CONF.

   bash

   export PYTORCH_ALLOC_CONF=expandable_segments:True
   
   # only if that is unavailable, and the docs call this a last resort:
   export PYTORCH_ALLOC_CONF=max_split_size_mb:512,garbage_collection_threshold:0.8

10. 10
    Rent the right card

    Nine levers in, you have spent real time and given away accuracy. Price that against the right GPU for a few hours: 4 to 12 USD on the 80 GB tier here, 1 to 3 USD on 24 GB.

    bash

    uv run torchrun --nproc_per_node=8 --master_port=29500 \
      gr00t/experiment/launch_finetune.py \
      --base-model-path nvidia/GR00T-N1.7-3B \
      --num-gpus 8 --global-batch-size 64 ...

Which levers silently change the run

Some flags are pure infrastructure. Others alter what the model learns, and the effect surfaces weeks later as a policy that behaves differently from the one you validated. Comparing two checkpoints means holding every row marked yes constant.

The BatchNorm folklore does not apply here

There is a folk rule that lowering batch size breaks BatchNorm. Check before acting on it. lerobot's ACT builds its ResNet-18 backbone with FrozenBatchNorm2d, so the statistics are baked in, and GR00T, Pi0.5 and SmolVLA build on transformer backbones that normalise within each sample rather than across the batch. Across all five trainable policies here the batch-size risk is gradient noise, not normalisation.

That risk is real enough. The default in FinetuneConfig on main is a global batch of 64 at a learning rate of 1e-4, and those numbers were chosen together. Cut to batch 8 and keep 1e-4 and you have changed the optimisation problem, not just the footprint.

Freezing and LoRA: the honest trade

Isaac-GR00T N1.7 specifics

Read from NVIDIA/Isaac-GR00T on main on 2026-08-24. The entry point moved: N1.5 used scripts/gr00t_finetune.py, N1.7 uses gr00t/experiment/launch_finetune.py, a tyro CLI over FinetuneConfig. Tutorials written for N1.5 have half the flag names wrong.

Two details when hunting bytes. launch_finetune.py hard-codes load_bf16 = False and backbone_trainable_params_fp32 = True, so the profile is not simply everything in bf16. And the CLI exposes no seed: the Isaac-GR00T README says to expect 5 to 6 percent variance from non-deterministic augmentation, which is the floor for judging whether a memory change hurt your policy.

lerobot specifics: Pi0.5, SmolVLA, ACT

lerobot moved its runtime onto Hugging Face Accelerate, relocating several memory flags into an accelerator config that carries mixed precision, gradient accumulation, FSDP2 and DDP. Two sub-configs are unfinished, and that is the trap.

# Pi0.5 on lerobot, memory-conscious. The docs describe the un-frozen
# version of this command as sized for a single 80 GB GPU.
lerobot-train \
  --dataset.repo_id=$HF_USER/my_so100_dataset \
  --policy.type=pi05 \
  --policy.pretrained_path=lerobot/pi05_base \
  --policy.dtype=bfloat16 \
  --policy.gradient_checkpointing=true \
  --policy.freeze_vision_encoder=true \
  --policy.train_expert_only=true \
  --policy.device=cuda \
  --batch_size=16 \
  --accelerator.gradient_accumulation.steps=4 \
  --num_workers=4 --steps=30000 --seed=1000

SmolVLA is the easy case: it ships freeze_vision_encoder = True and train_expert_only = True, and uses only the first 16 layers of the SmolVLM2 backbone. The Hugging Face guide suggests batch 64. If it still will not fit on 24 GB, look at resize_imgs_with_padding, which defaults to 512 by 512.

ACT rarely runs out of memory: at roughly 80 M parameters the whole 18-bytes-per-parameter budget is about 1.4 GB. When it does, the cause is action chunking and camera count. lerobot's ACT defaults are chunk_size 100 and n_action_steps 100 over a 512-dimensional transformer, Chunk size sets the number of decoder queries, and each camera runs the shared ResNet-18 once and flattens its feature map into extra encoder tokens, so cameras and chunk length grow the activations independently.

Stop guessing: take a memory snapshot

If you have pulled four levers and still do not know where the memory went, measure it. PyTorch records every allocation with its Python stack and renders it as a browser timeline.

import torch

torch.cuda.memory._record_memory_history(max_entries=100_000)
try:
    for step, batch in enumerate(loader):
        train_one_step(batch)
        if step == 5:      # peak is reached early
            break
finally:
    torch.cuda.memory._dump_snapshot('oom.pickle')
    torch.cuda.memory._record_memory_history(enabled=None)

print('allocated', torch.cuda.memory_allocated() / 2**30)
print('reserved ', torch.cuda.memory_reserved() / 2**30)
print('peak     ', torch.cuda.max_memory_allocated() / 2**30)

Read it for shape, not totals. A flat block that appears once and never moves is weights and optimizer state. A sawtooth rising through the forward and collapsing in the backward is activations. A staircase growing every step is a leak, usually a list of loss tensors that still carry their graph. No flag here fixes that; loss.detach() does.

Allocator settings that cost nothing

Two ways to get a run that fits

When the answer is simply a bigger card

Engineers underrate this because it feels like giving up. Price it instead. Three hours fighting a 40 GB card, landing on a frozen vision encoder plus LoRA plus batch 4, gets you a policy worse than the reference recipe. The alternative was 8 more dollars. The levers above earn their keep when you are constrained by hardware you own. For a single policy on an SO-100, the cheapest fix is the right tier: 80 GB for GR00T or Pi0.5, or SmolVLA and ACT on 24 GB.

One caveat covers the subject. Fitting the run is not the same as getting a policy that works on the arm. A model that only fits because you froze the vision tower and halved the resolution can train to a low loss and still fail in your kitchen. That is low loss with no useful behaviour. Compare the constrained recipe against the reference one on rollouts, and use the benchmark arena for what the unconstrained versions reach.

Compact version: read the error, rule out host RAM, check your dtype, then work down the list and log every lever from the 'changes the policy' column. That log explains, weeks later, why the training run that fit on the card produced a different policy from the guide's. To skip the tuning, the GR00T N1.7 on SO-100 guide and the SmolVLA on SO-100 guide list what each run uses, and pricing shows what the right card costs. Still choosing a model? Start with the VLA overview and the SO-100 guide.