VLA Fine-Tuning VRAM & Max Batch Calculator

VLA Fine-Tuning VRAM Principles & Engineering Guide

When fine-tuning large Vision-Language-Action (VLA) models such as SmolVLA or OpenVLA on a single consumer or workstation GPU, the primary memory bottleneck is not just the model weights, but optimizer states and activation memory scaling with batch size.

1. Memory Breakdown Formula (ZeRO Mixed Precision Adam)

Total training GPU memory consists of a fixed term (weights + gradients + optimizer states), batch-dependent activation memory, and framework/CUDA context overhead:

M_train ≈ (P_frozen × b_w + P_train × 16) + M_act(B) + M_overhead M_fixed = (P_frozen × b_w + P_train × 16) / 10^9 [GB]

* ZeRO 논문(arXiv:1910.02054) 기준: 혼합 정밀도 Adam은 FP16 가중치(2B) + FP16 그래디언트(2B) + FP32 마스터 가중치/모멘텀/분산(12B) = 파라미터당 16바이트를 점유합니다.

2. Deriving Maximum Batch Size (B_max) from 2 Short Probes

Because activation memory depends heavily on camera count, resolution, and action chunk length, theoretical modeling is prone to error. Running short probes at batch sizes 2 and 4 to measure peak VRAM allows reliable linear extrapolation:

a = (M(B2) − M(B1)) ÷ (B2 − B1) [GB / 배치] M_0 = M(B1) − a × B1 [GB, 베이스 오버헤드 절편] B_max = ⌊ (Safety_Margin × V − M_0) ÷ a ⌋

3. Python Snippet to Count Trainable Parameters

In the LeRobot environment, inspect parameter counts directly using requires_grad to compute the exact fixed term:

from lerobot.policies.smolvla.modeling_smolvla import SmolVLAPolicy

policy = SmolVLAPolicy.from_pretrained("lerobot/smolvla_base")
total = sum(p.numel() for p in policy.parameters())
train = sum(p.numel() for p in policy.parameters() if p.requires_grad)

print(f"Total: {total/1e6:.0f}M, Trainable: {train/1e6:.0f}M")
print(f"Fixed Term: {((total-train)*2 + train*16)/1e9:.2f} GB (frozen BF16)")

4. Frequently Asked Questions (FAQ)

Can SmolVLA be fine-tuned on a single 24GB GPU (RTX 3090/4090)?

Yes. With the default setup (frozen vision encoder, training only the ~100M action expert), the fixed memory requirement is only about 2.3 GB. This leaves ample room on a 24GB card for batch sizes of 4 to 8. Full fine-tuning increases the fixed term to ~7.2 GB, requiring smaller batch sizes.

Why does mixed-precision Adam consume 16 bytes per trainable parameter?

According to the ZeRO paper (arXiv:1910.02054), standard mixed-precision Adam requires 2 bytes for FP16 weights, 2 bytes for FP16 gradients, 4 bytes for the FP32 master weight copy, 4 bytes for first momentum, and 4 bytes for second variance—totaling 16 bytes per parameter.

How does training VRAM differ from inference VRAM?

Inference requires no optimizer states or backward pass activation graphs, needing only model weights (2 GB per billion parameters in FP16) and KV cache. Training adds 16 bytes per trainable parameter plus extensive backpropagation activation buffers, demanding substantially more VRAM.

Primary sources

Data & formulas verified: 2026-10-23 · Operator: Huni ([email protected])