learninfra · Linux · Networking · Kubernetes · System Design · AI Infrastructure · Exam blueprints · Drills

GPU Memory Math

Will it fit, and on how many GPUs? Weights by precision, the KV cache, the 16 bytes per parameter of training with Adam, and how ZeRO, FSDP, tensor parallelism and activation checkpointing divide the bill.

An interactive AI Infrastructure lesson: 23 steps, about 30 minutes, on a live simulation in your browser.

You run one DGX H100: eight H100s with 80 GB of HBM each. Two requests land on your desk this week. The support team wants to serve Llama 3.1 70B as an internal assistant, and the data team wants to fine-tune Llama 3.1 8B on a year of tickets. Both ask the same thing: will it fit, and on how many GPUs?

Start small. Llama 3.1 8B has 8.03 billion parameters. Served in bf16, each one is 2 bytes: 8.03 B × 2 = 16 GB of weights. Add a 1.5 GB workspace, 1 GB of CUDA context and 0.54 GB of KV cache for one 4,096-token conversation, and the plan comes to 19 GB of 80.

What you will learn

  1. Will it fit?

    • The question you are always asked: Weights in GB = parameters in billions × bytes per parameter. An 8B model in 16-bit is 16 GB before anything else.
    • The 70B model on one GPU: A model fits only when every weight is resident on the GPUs serving it: inference reads every parameter for every token.

  2. Bytes per parameter

    • Precision is a memory dial: Bytes per parameter: fp32 4, bf16 2, fp8 1, 4-bit 0.5. Quantising a model is the cheapest way to fit it on fewer GPUs.
    • 405B on the whole server: Tensor parallelism divides the weights by tp, but the sum still has to fit: total weights must be less than the HBM of all the GPUs together, with room left over.

  3. The KV cache

    • The KV cache grows with every token: KV cache = 2 × layers × KV heads × head dim × bytes per token, times context length, times concurrent sequences. Weights are fixed; the KV cache is what grows with users.
    • Break it: the same model without GQA: The KV cache scales with KV heads, not attention heads. Grouped-query attention divides it by the group size: 64 heads sharing 8 KV heads is an 8× saving.
    • Break it: weights fit, users do not: A serving plan fits only when weights plus a useful KV cache fit. Weights at 90 % of HBM is a server with no room for users.
    • What four GPUs buy you: KV capacity in tokens = (HBM budget − weights − overhead) × GPUs / KV bytes per token. Divide it by the context length to get concurrent users.
    • Drill: the 8B model's KV cache: An 8B Llama holds 0.13 MB per token: a 128k-token context is 17 GB of cache, larger than the model's own weights.

  4. Training: 16 bytes per parameter

    • Training needs 16 bytes per parameter: Full training with Adam in mixed precision costs 16 bytes per parameter (2 weights + 2 gradients + 12 optimizer) plus activations: about eight times the memory of serving in bf16.
    • Eight GPUs, plain data parallel: Plain data parallelism copies the whole training state onto every GPU. It multiplies speed, never capacity.
    • Break it: submit it anyway: An OOM is arithmetic you did not do. It fails fast when the static states do not fit and late when activations peak on the longest batch.

  5. Dividing the bill

    • ZeRO: stop keeping eight copies: ZeRO-1 divides the optimizer by the number of data-parallel GPUs, ZeRO-2 also the gradients, ZeRO-3 / FSDP also the weights. Activations are not divided.
    • The fine-tune runs: ZeRO-3 trades memory for communication: 1.5 times the data-parallel traffic, cheap on NVLink, expensive on a slow network.
    • Tensor and pipeline parallel split more: Parameters per GPU = total / (tp × pp); ZeRO then divides states by dp. Only tensor and sequence parallelism divide the activations.
    • Longer sequences: Activation memory is linear in tokens per micro-batch: double the sequence length or the micro-batch and you double it. It is the part of the bill that changes from run to run.
    • Activation checkpointing buys memory: Activation checkpointing trades a third more compute for almost all of the activation memory: about 25 % fewer tokens per second.

  6. Sizing a real request

    • Fine-tuning 70B: which server?: Capacity per GPU sets the minimum number of GPUs; the per-GPU plan does not care whose GPU it is. The fewest GPUs = total bytes / usable memory per GPU, rounded up.
    • How many H100s, honestly: Round a plan up until it uses at most about 90 % of HBM. A plan with 0.5 % headroom is an OOM that has not happened yet.
    • Serving 405B: three answers: To serve a model that does not fit, lower the bytes per parameter, raise the memory per GPU, or add GPUs. Each trades a different cost.
    • Drill: 70B training state: 70 billion parameters × 16 bytes is about 1.1 TB: the reason full fine-tunes of large models start at dozens of GPUs, and why LoRA exists.

  7. Recap & playground

    • Cheat sheet
    • Playground