Why VRAM, not FLOPS, limits your batch size

The spec sheet lies by omission

GPU marketing leads with TFLOPS, but almost nobody exhausts compute first. In practice, the number that ends your scaling story is memory: capacity caps what fits, and bandwidth caps how fast you can touch it. Understanding where the bytes go — precisely, not vibes — is the highest-leverage systems knowledge an ML engineer can have, because it turns OOM errors from surprises into arithmetic.

Where the bytes go in training

Train a model with AdamW in standard mixed precision and each parameter costs roughly 16 bytes: 2 for BF16 weights, 2 for BF16 gradients, 4 for the FP32 master copy, and 8 for the two FP32 optimizer moments. An 8B-parameter model therefore wants ~128 GB before a single activation is stored — which is why "small" 8B full fine-tunes don't fit on one 80 GB H100 without sharding or optimizer tricks.

Activations come on top and scale with batch size × sequence length × hidden size × layers. They are the only major term you control at runtime, which is why the batch size knob hits the memory wall long before the GPU runs out of arithmetic. Activation checkpointing trades ~30% more compute for a large activation discount — a good deal precisely because compute is the resource you have spare.

Where the bytes go in inference

Inference swaps optimizer states for the KV cache. Every generated token must attend to every cached key and value, and the cache grows linearly with context:

# KV cache bytes per token = 2 (K and V) * layers * kv_heads * head_dim * dtype_bytes
def kv_bytes_per_token(layers, kv_heads, head_dim, dtype_bytes=2):
    return 2 * layers * kv_heads * head_dim * dtype_bytes

kv_bytes_per_token(32, 8, 128)   # Llama 3 8B  -> 131,072 B = 128 KB/token
kv_bytes_per_token(80, 8, 128)   # Llama 3 70B -> 327,680 B = 320 KB/token

At 8k context, one Llama 3 8B sequence carries a 1 GB cache; a 70B sequence carries 2.5 GB. On an H100 serving 8B in BF16, ~16 GB goes to weights and the remaining ~60 GB of usable memory supports roughly 55-60 concurrent 8k-context requests. That number — not TFLOPS — is your throughput ceiling, because batch size is what keeps the GPU's compute units fed during decode.

Note the grouped-query attention in that formula: 8 KV heads instead of 32 query heads cuts the cache 4× versus classic multi-head attention. Architecture choices are memory choices.

Bandwidth: the second memory wall

Even when everything fits, decode is bandwidth-bound: generating one token requires streaming the entire working set — weights plus cache — through the memory system. An H100's 3.35 TB/s reading 16 GB of weights bounds single-stream decode near 200 tok/s no matter how many TFLOPS sit idle. Batching is the escape hatch: the same weight-read is amortized across every sequence in the batch, which is why aggregate throughput scales so well until capacity runs out. Capacity limits the batch; the batch limits how much bandwidth you amortize; FLOPS watch from the sidelines.

What to actually do

• Training: shard state with FSDP/ZeRO, checkpoint activations selectively, and consider 8-bit optimizers before reaching for more nodes.
• Inference: quantize weights (4-bit takes Llama 3 8B from 16 GB to ~6 GB — all of it recovered as KV cache), use paged attention, and cap max-model-len at what traffic really needs. FP8 KV cache halves the per-token cost again.
• Hardware selection: buy memory for the job. An L40S at $0.69/hr with 48 GB is the right home for a quantized 8B model; an H100 at $2.49/hr earns its price on 70B-class serving; the H200's 141 GB at $3.69/hr exists for the workloads where capacity, not speed, is the binding constraint — long context, big batches, models that almost-but-don't-quite fit.

Next time a job OOMs, don't shrink the batch and shrug. Do the arithmetic above, find which term is eating the budget, and attack that term. The fix is usually cheaper than the bigger GPU — and when it isn't, at least you'll know exactly why you're paying for one.