GPU Memory Math for VLA Training

The ten minutes that save you a wasted rental

There are two ways to find out whether a vision-language-action model fits on the GPU you are about to rent. One is to rent it, launch the run, and read the CUDA out-of-memory traceback four minutes later. The other is arithmetic, which takes about ten minutes and works offline.

This page is that arithmetic: the four things that occupy VRAM during fine-tuning, the bytes-per-parameter constants from the published literature, the one bucket you genuinely cannot predict, and how to measure the rest in six lines of Python. The worked examples use the five policies you can actually train on this platform: GR00T N1.7, GR00T N1.5, Pi0.5, SmolVLA and ACT.

Four buckets, and only three of them are predictable

The first three buckets are the ones people quote, and together they are usually called the model state. For most VLA fine-tuning runs they are the smaller half of the problem. The fourth bucket is what decides your batch size, and the fifth is why a careful calculation still comes out optimistic.

Bytes per parameter, by training regime

The canonical accounting is in section 3.1 of the ZeRO paper. For a model of size psi, mixed-precision Adam holds an fp16 copy of the parameters and of the gradients at 2 bytes each, plus optimizer state made of an fp32 copy of the parameters, the momentum and the variance at 4 bytes each. The paper calls the optimizer multiplier K, states that mixed-precision Adam has K equal to 12, and totals it as 16 bytes per parameter. Plain fp32 Adam lands on the same 16 by a different route.

Read the first two rows again, because that pair is the most common misconception in this subject. Moving from fp32 to bf16 mixed precision does not reduce model state at all; it shuffles bytes around inside the 16. The Hugging Face GPU training guide is blunt about the fp16 case: because an fp16 and an fp32 copy of the model are both resident, a small-batch run can end up using more memory rather than less.

Frozen parameters cost 2 bytes, trainable ones cost 16

This is the lever that makes 3B-parameter VLAs trainable on one card at all. Both GR00T and SmolVLA freeze most of themselves by default. NVIDIA's fine-tuning config sets tune_llm and tune_visual to False while tune_projector and tune_diffusion_model default to True, so a default run trains the projector and the diffusion action head and leaves the Cosmos-Reason2-2B backbone alone. The SmolVLA config in lerobot does the same with different names: freeze_vision_encoder and train_expert_only are both True, and the SmolVLA paper puts roughly 100 million of its 450 million parameters in the action expert.

So the model-state floor for a fine-tune is not 16 bytes times the total. It is 2 or 4 bytes times the frozen part plus 16 bytes times the trainable part. Here is that sum for the five trainable policies. Parameter counts come from the model cards and papers; the trainable fractions come from each project's own defaults, except the GR00T trainable count, which NVIDIA does not publish and which is taken from this platform's catalog.

That table already explains the GPU tiers on this platform. ACT and SmolVLA have floors of one to three gigabytes and sit comfortably on a 24 GB card. GR00T N1.7 and Pi0.5 are on the A100 80 GB or H100 80 GB tier, and the Pi0.5 row shows why: lerobot's pi05 configuration defaults to dtype float32 with nothing frozen, so the model state alone is around 48 GB before a single image is loaded.

Activations: the bucket that sets your batch size

Activations are every intermediate tensor the backward pass will need. Unlike model state they depend on the shape of the data, and for a transformer there is a published closed form. Equation 1 of NVIDIA's activation-recomputation paper gives the per-layer cost with no parallelism, in bytes, assuming 16-bit storage.

activations_per_layer = s * b * h * (34 + 5 * a * s / h)

  s = sequence length (tokens)
  b = microbatch size
  h = hidden dimension
  a = number of attention heads

total = activations_per_layer * L        # L = number of layers

# The 34*s*b*h term is linear in tokens.
# The 5*a*s^2*b term is quadratic in tokens: that is attention.
# With full activation recomputation the whole thing collapses to 2*s*b*h*L.

For a VLA the sequence is not a sentence. It is image tokens plus language tokens plus a state token, and the image tokens dominate. Take SmolVLA's backbone as a concrete case. Its lerobot config pads images to 512x512 and its SmolVLM2 backbone uses a patch size of 16, which is 1024 patches per image, folded down by the pixel-shuffle scale_factor of 4 to 64 tokens per camera. Text hidden size is 960 with 15 attention heads, and the policy uses only the first 16 of the backbone's 32 language-model layers. Plugging those into equation 1 at the platform default batch size of 2 gives the following.

Raising the token count by 64 percent, from 200 to 328, does not raise the memory by 64 percent: it doubles it. Another 78 percent on top, to 584 tokens, multiplies it by 2.4 again. That is the quadratic attention term taking over, and it is why resolution and camera count are more expensive than they look. Batch size, by contrast, is honestly linear: the same 200-token configuration costs 0.28 GiB at batch 2, 1.14 GiB at batch 8 and 4.54 GiB at batch 32.

Worked example: GR00T N1.7, and why the guide asks for 80 GB

NVIDIA's hardware recommendation guide states two numbers that together give you a usable model of the whole run. It sets the fine-tuning minimum at one GPU with 40 GB or more of VRAM, and it says that default fine-tuning, which tunes the projector and diffusion action head rather than the backbone, keeps peak VRAM under about 35 GB per GPU. The README's single-GPU example runs at --global-batch-size 32.

1. 1
   Start from the model-state floor

   3.0 billion parameters in bfloat16 is 6.0 GB. About 40 million of them are trainable at 16 bytes each rather than 2, which adds another 0.56 GB. Call the floor 6.6 GB.

   python

   total_params   = 3.0e9
   trainable      = 40e6
   
   frozen_bytes    = (total_params - trainable) * 2      # bf16 weights
   trainable_bytes = trainable * 16                      # bf16 + fp32 master + 2 fp32 moments
   floor_gb        = (frozen_bytes + trainable_bytes) / 1e9
   print(round(floor_gb, 2))     # 6.56

2. 2
   Subtract the floor from the reported peak

   Upstream reports a peak under about 35 GB at a global batch size of 32. Everything above the floor is activations plus workspaces plus fragmentation, so that is roughly 28 GB spread across 32 samples.

   python

   peak_gb       = 35.0        # upstream figure, batch 32
   floor_gb      = 6.56
   per_sample_gb = (peak_gb - floor_gb) / 32
   print(round(per_sample_gb, 2))   # 0.89 GB per sample

3. 3
   Turn it into a batch-size rule

   Peak is roughly 6.6 GB plus 0.9 GB per sample. Reserve 10 percent of the card and solve for the batch size. This is an estimate built from two published numbers, not a measurement, so treat the answer as an upper bound to test rather than a promise.

   python

   def max_batch(card_gb, floor=6.6, per_sample=0.9, headroom=0.10):
       usable = card_gb * (1 - headroom)
       return int((usable - floor) / per_sample)
   
   for card in (24, 40, 48, 80, 96):
       print(card, max_batch(card))
   # 24 -> 16   40 -> 32   48 -> 40   80 -> 72   96 -> 88

4. 4
   Check the answer against reality, and believe reality

   The rule says batch 16 would fit on a 24 GB card. Upstream says the minimum is 40 GB. The gap is bucket five: CUDA context, cuDNN workspaces, pinned dataloader buffers and the transient spike when a checkpoint is written. When a derived rule and a vendor minimum disagree, the vendor minimum is the one that has actually been run.

   bash

   CUDA_VISIBLE_DEVICES=0 uv run python \
       gr00t/experiment/launch_finetune.py \
       --base-model-path nvidia/GR00T-N1.7-3B \
       --dataset-path demo_data/cube_to_bowl_5 \
       --embodiment-tag NEW_EMBODIMENT \
       --modality-config-path examples/SO100/so100_config.py \
       --num-gpus 1 \
       --output-dir /tmp/test_finetune \
       --max-steps 2000 \
       --global-batch-size 32 \
       --dataloader-num-workers 4

The same arithmetic explains the flag that everyone eventually reaches for and regrets. Enabling --tune-llm makes essentially all 3 billion parameters trainable, which takes the model state from 6.6 GB to 3.0 times 16, or 48 GB. That leaves about 24 GB of an 80 GB card for activations, and it is exactly why the upstream guide recommends 80 GB or more per GPU once --tune-llm or --tune-visual is on. The line between the 40 GB tier and the 80 GB tier is one boolean.

Worked example: the 24 GB tier

SmolVLA and ACT are the two policies you can train on an RTX 4090 or any other 24 GB card, and the arithmetic shows real headroom rather than a squeeze. The ACT paper is the sanity check here: its authors trained the model in about five hours on a single 11 GB RTX 2080 Ti. Both policies are trained through lerobot, whose training step defaults live in one dataclass: batch_size 8, steps 100000, num_workers 4, save_freq 20000 and seed 1000.

There is a version detail here that matters more than it should, and it is worth dating because upstream moves. On the lerobot 0.5.1 line that this platform pins, the policy config carries a use_amp field, but the training script never reads it: it builds its Accelerator without a mixed_precision setting, so training runs in fp32 and stores activations at roughly twice the bf16 cost. On lerobot main, which carries the version string 0.6.2 while 0.6.1 is the newest release on PyPI, mixed precision is a first-class accelerator setting that defaults to "no" and has to be asked for.

# single GPU, bf16 mixed precision, small batch
lerobot-train \
  --policy.type=smolvla \
  --dataset.repo_id=<user>/<dataset> \
  --batch_size=2 \
  --steps=20000 \
  --policy.device=cuda \
  --accelerator.mixed_precision=bf16

# eight GPUs, sharded
torchrun --nproc-per-node=8 $(which lerobot-train) \
    --dataset.repo_id=... --policy.type=act \
    --parallelism.dp_shard=8 --accelerator.mixed_precision=bf16

The same pinning explains the training form. The gradient accumulation field does not take effect for Pi0.5 or SmolVLA, because lerobot 0.5.1 has no such flag at all; accumulation only arrived with the accelerator config on the 0.6 line. It does apply for both GR00T variants, which use NVIDIA's trainer rather than lerobot's. The training documentation and the out-of-memory failure mode page both record which knobs are live for which trainer.

Stop guessing: measure the peak

Every number above is a floor or an estimate. The peak is a measurement, and PyTorch gives it to you for free. torch.cuda.max_memory_allocated returns the high-water mark of tensor memory since the last reset, and torch.cuda.reset_peak_memory_stats clears it. Wrap a handful of training steps and you have the real answer in under a minute of GPU time, which on the 24 GB tier costs about a cent.

import torch

GB = 1024 ** 3

def probe(batch_size, steps=5):
    torch.cuda.empty_cache()
    torch.cuda.reset_peak_memory_stats()
    for _ in range(steps):
        batch = make_batch(batch_size)      # your dataloader
        loss = policy(batch)[0]
        loss.backward()
        optimizer.step()
        optimizer.zero_grad(set_to_none=True)
    return {
        "peak_alloc_gb":    torch.cuda.max_memory_allocated() / GB,
        "peak_reserved_gb": torch.cuda.max_memory_reserved() / GB,
    }

for b in (1, 2, 4, 8, 16):
    print(b, probe(b))

Watch the gap between allocated and reserved. Allocated is what your tensors hold; reserved is what the caching allocator has taken from the driver and not returned. A large and growing gap is fragmentation, and the PyTorch CUDA notes describe the two options for it: expandable_segments, which is marked experimental and defaults to False, and max_split_size_mb, which stops the native allocator from splitting blocks larger than the given size. In PyTorch 2.13, the current stable release, the variable is named PYTORCH_ALLOC_CONF and PYTORCH_CUDA_ALLOC_CONF is kept as a backward-compatible alias.

1. 1
   Sweep the batch size on the smallest card you own

   Run the probe above at batch 1, 2, 4, 8 and 16. Two of those points give you the slope and intercept of a straight line, and the line predicts the rest.

   python

   # peak(b) ~= floor + per_sample * b
   # solve from two measured points
   b1, p1 = 2, 9.4
   b2, p2 = 8, 14.8
   per_sample = (p2 - p1) / (b2 - b1)
   floor      = p1 - per_sample * b1
   print(f"peak(b) = {floor:.1f} + {per_sample:.2f} * b")

2. 2
   Set the allocator variable if reserved runs far ahead of allocated

   Fragmentation is the failure mode where the run that worked yesterday dies at step 4000 today. Both options are documented in the PyTorch CUDA memory-management notes.

   bash

   export PYTORCH_ALLOC_CONF=expandable_segments:True
   # older PyTorch, or if you prefer the alias:
   export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True,max_split_size_mb:128

3. 3
   If you still need the shape of the problem, take a snapshot

   The snapshot records stack traces for every allocation. Drop the pickle into the visualiser at pytorch.org/memory_viz, which the docs describe as a JavaScript application that runs locally and uploads nothing.

   python

   torch.cuda.memory._record_memory_history()
   # ... run a few training steps ...
   torch.cuda.memory._dump_snapshot("my_snapshot.pickle")

The knobs, ranked by what they actually buy

Not all memory savings are equal, and several popular ones do very little for a vision policy. This ordering is by how much VRAM each option removes per unit of pain, with the sourced cost of each.

Traps that eat a day

• Gradient checkpointing is not currently reachable from lerobot's CLI. --accelerator.activation_checkpointing exists on main, but the config validation raises an error calling it a placeholder that is not wired yet. --accelerator.compile is in the same state. Plan around the batch size instead.
• Reserved memory climbing while allocated stays flat is fragmentation, not a leak. Reach for expandable_segments before you reach for a bigger card.
• The checkpoint write is a spike. GR00T's save_only_model skips optimizer, scheduler and RNG state, at the price of being unable to resume from that checkpoint.
• torch.cuda.empty_cache() releases only unused cached blocks and cannot free memory your tensors still hold. Hugging Face puts the cost of using it during training at roughly 10 percent slower.
• A LeRobot v3.0 dataset crashes the GR00T loader outright and has to be converted down to v2.1 first. That traceback has nothing to do with memory; see the dataset rejected page.
• GR00T's launch_finetune.py is a tyro CLI that exposes no seed, and NVIDIA notes 5 to 6 percent variance between runs from non-deterministic image augmentations. Two runs at the same batch size can peak slightly differently. lerobot's default seed is 1000.

Two ways to reach the same answer

Where the arithmetic stops helping

Fitting is a low bar. The memory math tells you nothing about whether 50 episodes are enough, whether your normalisation statistics are right, or whether the policy will generalise past the exact table height you recorded at. It is a gate, not a predictor.

• It says nothing about throughput. A run that fits at batch 1 with 16 accumulation steps can be far slower in wall-clock than one that fits at batch 32, at identical effective batch size.
• It says nothing about inference. Serving needs the weights plus a working set, and NVIDIA quotes a 16 GB minimum for GR00T N1.7 inference against 40 GB for fine-tuning. A policy that took an 80 GB card to train may serve on a 4090.
• It does not predict data-loading stalls. If your GPU sits at 30 percent utilisation, memory is not your problem and a bigger batch will not fix it.
• It cannot see inference latency, which is what actually decides whether the trained policy is usable on the arm.

That last point is the limit this platform does not engineer away. The control loop runs at 20 to 485 ms per action step depending on the model, and adding public-internet round trips turns a working policy into a hesitant one. Remote inference is fine for slow pick-and-place and is not fine for fast reactive motion, no matter how the training run went. Compare the numbers yourself on the GR00T versus Pi0.5 comparison or across all 85 models on the arena leaderboard, and read the VLA overview if you are still choosing an architecture.