GPU memory & the model

1 · Where the VRAM actually goes essential

Here's the moment this lesson exists to prevent. You pick a model, you rent a GPU, you type the launch command — and thirty seconds later the log spits CUDA out of memory and the process dies. You try a bigger card. It happens again on a different request. Now you're guessing at hardware with the meter running, and every guess costs real money. The engineers who never hit this aren't luckier; they did the arithmetic before they rented anything, because GPU memory isn't mysterious — it's a budget with a fixed number of line items.

Think of the card's VRAM the way you'd think of a moving van with a strict weight limit. You can't just "hope it fits" — you add up the weight of every thing you're loading and compare it to the limit stamped on the side. A GPU is exactly that: an A100 gives you 40 or 80 GB, an H100 gives you 80, and if the sum of what you load crosses that line, nothing loads at all — it's a hard wall, not a slow-down. So the real skill isn't "make it faster," it's knowing what's in the van.

And the beautiful part is that there are only four things ever competing for that space, no matter how exotic the model:

• Weights — the model's actual learned numbers. A fixed, one-time cost the moment you load it, present whether you're serving or training.
• KV cache — scratch memory kept per token during generation. Zero at load, then it grows with every user and every token of context they send.
• Activations — intermediate results the backward pass must hold onto. Tiny at inference, large during training.
• Optimizer state — the bookkeeping an optimizer like Adam keeps per parameter. Zero at inference, often the biggest item when training.

The single most useful thing to internalize before the table below: serving and training pay different bills. Serving pays for weights + KV cache and little else, which is why one card can often serve a model. Training additionally pays for activations + gradients + optimizer state, which is why the same model that serves on one GPU needs a cluster to train. The common beginner mistake is assuming "if it loads, it trains" — it doesn't, and section 4 shows the gap is roughly 7×. Read the four line items as a checklist you run on every deployment: which boxes apply to my case, and do they sum under the ceiling?

A GPU's memory (VRAM) is a hard ceiling — an A100 has 40 or 80 GB, an H100 80 GB. When people say a model "doesn't fit", they mean the sum of four things exceeds that ceiling. Learn the four line items and you can predict fit before you launch:

Inference pays for weights + KV cache (plus a little working memory). Training adds activations + gradients + optimizer state — which is why the same model needs far more VRAM to train than to serve.

2 · Precision — bytes per parameter essential

Start with the mistake almost everyone makes once. Someone glances at "Llama-3 70B," sees the number 70, and reasons "70 billion… so about 70 GB, that'll fit my 80 GB card." They're wrong by a factor of two, and the reason is a confusion so common it's worth burning into memory: a parameter count is not a byte count. "70B" tells you how many numbers the model has. It says nothing about how big each one is — and that second fact is the entire difference between "fits" and "OOM."

Here's the analogy that fixes it for good. Imagine a warehouse holding 70 billion boxes. Knowing the box count doesn't tell you how much floor space you need — you also need to know how big each box is. A parameter is a box; its precision is the box size, measured in bytes. Store each number in 32 bits (fp32) and each box is 4 bytes; in 16 bits (fp16/bf16) it's 2 bytes; in 8 bits (int8) it's 1 byte; in 4 bits (int4) it's half a byte. Same 70 billion boxes, but the floor space you rent swings by 8× depending purely on box size.

So the master formula for weight memory is refreshingly simple: bytes = params × bytes-per-param. Every capacity question in this lesson bottoms out in that one multiplication. The precision column below is therefore the most important constant on the page — not because the bit-splitting details matter day to day, but because it's the number you multiply by to turn a headline parameter count into a real gigabyte figure. Notice the pattern as you read it: each step down the precision ladder roughly halves the footprint. That halving is exactly why quantization (IC2) exists — it doesn't make the model smarter, it makes each box smaller so more of them fit on the card you already have.

Every number in the model occupies a fixed number of bytes set by its precision. This single constant drives the whole weight-memory calculation. Memorize the table:

fp16 and bf16 are both 2 bytes — they differ in how the 16 bits split between exponent and mantissa (bf16 trades precision for fp32's dynamic range, which is why it's the training default). For memory math they're identical. Quantization (IC2) is precision applied to the weights after training; here we care about the bytes it buys you.

vram_weights.pyGIB = 1024 ** 3  # bytes in a gibibyte

def vram_for_model(params_b, bytes_per_param):
    """Weight memory in GiB. params_b = billions of params."""
    total_bytes = params_b * 1e9 * bytes_per_param
    return total_bytes / GIB

BYTES = {"fp32": 4, "fp16": 2, "bf16": 2, "fp8": 1, "int8": 1, "int4": 0.5}

for size in (7, 13, 70):
    line = [f"{size:>3}B:"]
    for fmt in ("bf16", "int8", "int4"):
        gib = vram_for_model(size, BYTES[fmt])
        line.append(f"{fmt}={gib:6.1f} GiB")
    print("  ".join(line))

  7B:  bf16=  13.0 GiB  int8=   6.5 GiB  int4=   3.3 GiB
 13B:  bf16=  24.2 GiB  int8=  12.1 GiB  int4=   6.1 GiB
 70B:  bf16= 130.4 GiB  int8=  65.2 GiB  int4=  32.6 GiB

The formula is just params × bytes/param. A 70B model in bf16 is 70e9 × 2 bytes ≈ 130.4 GiB of weights alone — already over a single 80 GB card. In int4 (0.5 bytes) it's ≈ 32.6 GiB, which fits. That is the whole reason quantization exists for serving large models.

3 · The KV cache — memory that grows with tokens intermediate

This is the section that explains a bug that has confused thousands of teams. You size your GPU carefully — weights fit with room to spare, you launch, everything's green. Then traffic ramps up, or one user pastes a very long document, and suddenly the same model that fit an hour ago throws CUDA out of memory. Nothing about the weights changed. What changed is the second, moving line item on the card: the KV cache. And unlike weights, which you can look up once, the KV cache is a live cost that breathes with your traffic.

To see why it exists, picture how the model generates text: one token at a time, and to produce each new token it "attends" to every token before it. The naive way would be to re-read and re-encode the entire conversation for every single new token — quadratically wasteful. So the model does the obvious optimization: the first time it processes a token it computes that token's key and value vectors and keeps them, so later tokens can look them up instead of recomputing (this is the attention mechanism from IC3). Think of it as the model's running notebook: every token it has ever seen leaves two notes (a K and a V) on every layer, and the notebook stays open in VRAM for the whole request.

That notebook is the trap, because it grows in four dimensions at once. The formula below spells it out, but read the intuition first — the cache size is proportional to:

• × sequence length — one more token means one more page of notes, per layer. Long contexts are expensive.
• × batch size — every concurrent request keeps its own notebook. Ten users is ten times the cache.
• × number of layers — the notes are kept at every transformer layer, not once.
• × 2 — because it's a Key and a Value stored per token, not one number.

The detail that trips people, and the one worth getting right: the formula uses n_kv_heads, the number of key/value heads — not the (usually larger) number of query heads. Modern large models deliberately use grouped-query attention (GQA), where many query heads share a smaller set of KV heads, precisely to shrink this cache. Read that as a design decision aimed straight at the memory bill, not an accident. The punchline you'll prove in the lab: on a busy server it's the KV cache, not the weights, that runs you out of memory first — so capacity planning at serving time is mostly KV-cache planning.

During generation the model caches the key and value vectors for every token it has already seen, at every layer, so it never recomputes attention over the past (IC3). That cache is pure VRAM overhead on top of the weights, and it grows linearly with sequence length and batch size. The formula:

kv_bytes = 2 × n_layers × n_kv_heads × head_dim × seq_len × batch × bytes_per_elem

The leading 2 is for K and V. Note n_kv_heads, not query heads — grouped-query attention (GQA) shrinks the cache by sharing KV across query heads, which is exactly why modern large models use it.

kv_cache.pyGIB = 1024 ** 3

def kv_cache_gb(n_layers, n_kv_heads, head_dim, seq_len, batch, bytes_per_elem=2):
    """KV-cache VRAM in GiB. Factor 2 is for K and V."""
    elems = 2 * n_layers * n_kv_heads * head_dim * seq_len * batch
    return elems * bytes_per_elem / GIB

# 70B-class model with grouped-query attention (8 KV heads), fp16 cache
cfg = dict(n_layers=80, n_kv_heads=8, head_dim=128, seq_len=8192)

for batch in (1, 8, 32):
    gib = kv_cache_gb(batch=batch, **cfg)
    print(f"batch={batch:>2} @ 8k ctx -> KV cache = {gib:6.2f} GiB")

batch= 1 @ 8k ctx -> KV cache =   2.50 GiB
batch= 8 @ 8k ctx -> KV cache =  20.00 GiB
batch=32 @ 8k ctx -> KV cache =  80.00 GiB

For a 70B-class model (80 layers, 8 KV heads via GQA, head_dim 128) a single 8k-token sequence in fp16 is only ~2.5 GiB — but 32 concurrent sequences is ~80 GiB, the whole card on its own. KV cache, not weights, is usually what OOMs a busy inference server.

4 · Activations & the training memory blow-up advanced

Here's a failure that feels like a betrayal. You've been happily serving a 13B model on a single 80 GB card for weeks. Then you decide to fine-tune it — same model, same card — and it OOMs before the first step even completes. It feels wrong: how can a model that fits not fit? The answer is that training is not "inference plus a backward pass." It's a fundamentally hungrier regime, and the extra appetite comes from two consumers that are simply absent at inference time.

Walk through what training actually has to keep in memory that serving doesn't. First, activations: to compute gradients, the backward pass needs the intermediate values from the forward pass — you can't know how to nudge a weight without remembering what flowed through it. At inference you compute a layer's output, use it, and throw it away; at training you must hoard those intermediate results until the backward pass consumes them. Second, and usually larger, is the optimizer state. An optimizer like Adam doesn't just apply the raw gradient — it keeps a running mean and running variance of the gradients for every single parameter, so it can adapt the step size. That's bookkeeping the model never carries when it's only answering questions.

Now count the bytes per parameter for the classic mixed-precision Adam recipe, because the arithmetic is the whole lesson:

• an fp32 master copy of the weights Adam updates — 4 bytes,
• the gradient, one per weight each step, in fp16/bf16 — 2 bytes,
• Adam's first moment m (mean of gradients), fp32 — 4 bytes,
• Adam's second moment v (variance of gradients), fp32 — 4 bytes.

That's ~14 bytes per parameter of persistent state — and note we haven't counted a single activation yet. Compare it to serving, which pays roughly 2 bytes per parameter (just the bf16 weights). The ratio is about 7×, and it's why the lab shows a 7B model costing ~13 GiB to serve but ~91 GiB of state to train. The critical takeaway — and the thing to say out loud in an interview — is that the optimizer state, not the weights, blows the training budget. That reframing is what points you at the right fixes: LoRA/QLoRA (FT3) train a tiny adapter so there's almost no optimizer state to hold, and sharding (MS2) splits that state across cards. You attack the 12 bytes of grad+optimizer, not the 2 bytes of weights.

Inference is memory-light: weights + KV cache and you're done. Training is a different regime. Backprop must keep the forward activations around to compute gradients, and the optimizer keeps its own state per parameter. The four-way split for Adam-style training in mixed precision:

That's roughly 14 bytes/param just for weights + gradients + optimizer state in the classic mixed-precision Adam recipe (4 + 2 + 4 + 4) — before a single activation. A 7B model is thus ~91 GiB of state, which is why it does not train on one 80 GB card without sharding (MS2) or memory-saving tricks. Activations are extra on top and scale with batch × seq × hidden × layers; gradient checkpointing trades compute to shrink them.

train_memory.pyGIB = 1024 ** 3

def inference_state_gb(params_b, bytes_per_param=2):
    """Just the weights (bf16 serving)."""
    return params_b * 1e9 * bytes_per_param / GIB

def training_state_gb(params_b):
    """Mixed-precision Adam persistent state, bytes/param:
       fp32 master weights (4) + grads (2) + Adam m (4) + Adam v (4) = 14."""
    bytes_per_param = 4 + 2 + 4 + 4
    return params_b * 1e9 * bytes_per_param / GIB

for size in (7, 13):
    inf = inference_state_gb(size)
    tr = training_state_gb(size)
    print(f"{size:>2}B: serve(bf16)={inf:6.1f} GiB | "
          f"train(Adam)={tr:6.1f} GiB | {tr/inf:.1f}x more")

 7B: serve(bf16)=  13.0 GiB | train(Adam)=  91.3 GiB | 7.0x more
13B: serve(bf16)=  24.2 GiB | train(Adam)= 169.5 GiB | 7.0x more

The same 7B model: ~13 GiB to serve in bf16, ~91 GiB of persistent state to train with Adam — a 7× gap, and that's before activations. This is the single most important reason training needs clusters while inference often needs one card.

5 · Does it fit? Batch & sequence vs the ceiling professional

There are two ways to find out whether a model fits a card. The expensive way is to rent the GPU, launch, and watch — you either get lucky or you OOM in production with users watching. The professional way is to add up the four line items on paper first and only rent the hardware the arithmetic already told you would fit. This section is about turning "does it fit?" from a nerve-wracking experiment into a single line of math you can run before spending a cent.

The whole question reduces to one inequality: weights + KV cache + headroom ≤ card VRAM. Two of those terms you now know cold — weights from section 2 (fixed, precision × params) and KV cache from section 3 (grows with batch and context). The insight that makes this planning rather than accounting is which terms you can actually move. Weights are frozen once you pick a precision; you don't get to negotiate them at runtime. The KV cache, by contrast, is the dial you turn every time you set a batch size or a max context length — so batch and sequence length are your real capacity knobs, and they trade directly against how many users you can pack onto a card.

One term deserves its own emphasis because skipping it is the classic rookie error: headroom. You never plan to fill 100% of the card. CUDA needs its own context, the memory allocator fragments, attention builds temporary buffers, and the occasional long request spikes above average — all of that eats VRAM you didn't put in your spreadsheet. Budget to roughly 85–90% of the card and treat the rest as slack. A deployment sized to exactly 80.0 GiB of an 80 GB card isn't "efficient," it's a pager alert waiting for the first busy minute. The checker in the lab bakes this in with a 10% reserve, so read its "usable" number (≈72 GiB) as the line you actually plan against.

Capacity planning is one question: does weights + KV cache + headroom fit the card? Weights are fixed; KV cache is the dial you actually turn with batch size and context length. The professional move is to compute the answer, not to try-and-OOM in production.

fits_80gb.pyGIB = 1024 ** 3

def vram_for_model(params_b, bytes_per_param):
    return params_b * 1e9 * bytes_per_param / GIB

def kv_cache_gb(n_layers, n_kv_heads, head_dim, seq_len, batch, bytes_per_elem=2):
    elems = 2 * n_layers * n_kv_heads * head_dim * seq_len * batch
    return elems * bytes_per_elem / GIB

def fits_80gb(params_b, bytes_per_param, kv_gb, headroom=0.10):
    """Does weights + KV cache fit an 80 GB card with headroom to spare?"""
    card = 80.0
    usable = card * (1 - headroom)
    need = vram_for_model(params_b, bytes_per_param) + kv_gb
    return need, usable, need <= usable

# 13B bf16 serving 8 users at 4k ctx
kv = kv_cache_gb(n_layers=40, n_kv_heads=40, head_dim=128, seq_len=4096, batch=8)
need, usable, ok = fits_80gb(13, 2, kv)
print(f"13B bf16 + KV {kv:.1f} GiB -> need {need:5.1f} / {usable:.1f} usable -> "
      f"{'FITS' if ok else 'OOM'}")

# 70B bf16, weights alone
need, usable, ok = fits_80gb(70, 2, 0.0)
print(f"70B bf16 weights only     -> need {need:5.1f} / {usable:.1f} usable -> "
      f"{'FITS' if ok else 'OOM'}")

# 70B int4, weights alone
need, usable, ok = fits_80gb(70, 0.5, 0.0)
print(f"70B int4 weights only     -> need {need:5.1f} / {usable:.1f} usable -> "
      f"{'FITS' if ok else 'OOM'}")

13B bf16 + KV 25.0 GiB -> need  49.2 / 72.0 usable -> FITS
70B bf16 weights only     -> need 130.4 / 72.0 usable -> OOM
70B int4 weights only     -> need  32.6 / 72.0 usable -> FITS

Read the output as a capacity plan: a 13B model in bf16 on an 80 GB card leaves plenty of room for KV cache, so you can push batch/context high; a 70B model in bf16 doesn't fit at all and must be quantized or sharded. The knobs — precision, batch size, max sequence length — are exactly what you tune to stay under the ceiling while maximizing throughput (IC1/IC4).

6 · Debugging an OOM methodically professional

Watch what a panicking engineer does when they hit CUDA out of memory: they lower the batch size a bit, it OOMs again; they grab a bigger card, it OOMs on a different request; they toggle three flags at once from a Stack Overflow answer and pray. That's debugging by superstition, and it burns hours and GPU-hours both. The reason it fails is that they're treating the OOM as a random event instead of what it actually is — a budget that went over by a specific amount, on a specific line item, that you can name.

Because you now know there are only four consumers, an OOM stops being a mystery and becomes an accounting problem: something in {weights, KV cache, activations, optimizer state} grew past the ceiling, and your job is to figure out which one and shrink that one. The discipline is two moves — attribute, then attack — and the order matters more than any individual trick.

The attribution step is where the profiler earns its keep, and the signal is in when the memory appears. Weights are constant — resident the moment you load, unchanging after. If memory climbs as you add concurrent requests, that's the KV cache. If it spikes specifically during .backward(), that's activations and optimizer state. So don't guess: log torch.cuda.memory_allocated() and max_memory_allocated() around load and around a forward/backward pass, and let the numbers point at the culprit. Only then do you attack — and you attack the biggest line item first, with the cheapest lever that shrinks it, rather than sprinkling fixes at random. The ordered playbook below is that logic written out; the one habit to carry from it is that reaching for a bigger card is the last resort, never the first, because it's the most expensive lever and it hides the real bug instead of fixing it.

"CUDA out of memory" is not a mystery once you know the four consumers. Debug it as a budget problem, in order — attribute the memory, then attack the biggest line item:

OOM debugging methodology

1. Read the allocator, don't guess. Log torch.cuda.memory_allocated() and max_memory_allocated() around load and around a forward pass to see what's resident vs. peak.
2. Attribute it to a consumer. Weights are constant; if memory climbs with concurrency it's KV cache; if it spikes only during .backward() it's activations/optimizer.
3. Attack the biggest line item. KV cache → cap batch/context, use GQA/paged attention, quantize the KV cache. Weights → quantize (IC2) or shard (MS2). Optimizer → LoRA/QLoRA (FT3), 8-bit optimizer, or shard.
4. Reduce activations if training. Enable gradient checkpointing, lower the micro-batch, use gradient accumulation to keep the effective batch while cutting peak memory.
5. Fix fragmentation last. If allocated is well under total but you still OOM, set PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True and avoid wild batch-size swings.

7 · Measuring it for real on a GPU advanced

Here's where a careful engineer keeps themselves honest. Everything so far has been paper arithmetic — powerful, cheap, and approximate. The formulas capture the big rocks (weights, KV cache, optimizer state), but real hardware always carries extra weight the spreadsheet omits: the CUDA context itself, allocator overhead, cuDNN workspace, temporary buffers a kernel grabs mid-attention. So the mature workflow is not "trust the math" or "just measure" — it's predict, then verify: size it offline to pick the smallest plausible card, then confirm once on the real thing before you commit.

The tool that closes the loop is the CUDA memory allocator, which PyTorch exposes directly. Two calls tell you almost everything. torch.cuda.memory_allocated() reports what's resident right now — call it after loading the model and it should land within a hair of your vram_for_model prediction (if a 7B in bf16 reports ~14 GiB, your arithmetic is trustworthy). torch.cuda.max_memory_allocated() reports the peak the allocator ever reached, which is the number that actually decides whether you OOM — because a forward or backward pass transiently spikes above the resident baseline. The gap between those two numbers is your headroom requirement, measured rather than guessed.

One honest caveat about the lab below, and why it's labeled the way it is: it imports torch and calls CUDA, so it genuinely cannot run in this offline, CPU-only environment — it's shown so you know exactly what to run on a real GPU box. That's the opposite of the calculators in labs 1–4, which are pure stdlib arithmetic you can run right here to predict what this snippet will later confirm. Read the two together as the full discipline: the offline math proposes, the allocator disposes.

Everything above is arithmetic you can do offline. On real hardware you confirm it with the CUDA allocator. This snippet is real and correct but needs a GPU + PyTorch — it will not run in this offline environment; it's here so you know exactly what to measure:

measure_gpu.py# >>> needs a GPU + PyTorch: pip install torch transformers accelerate
# This is REAL and correct but will NOT run in an offline/CPU-only environment.
import torch
from transformers import AutoModelForCausalLM

def mib(n):
    return n / (1024 ** 2)

# Load weights in bf16 to halve the memory vs fp32.
model = AutoModelForCausalLM.from_pretrained(
    "mistralai/Mistral-7B-v0.3",
    torch_dtype=torch.bfloat16,
    device_map="cuda",
)

# Confirm the arithmetic from vram_weights.py against the real allocator.
print(f"allocated after load: {mib(torch.cuda.memory_allocated()):.0f} MiB")

ids = torch.randint(0, 32000, (1, 512), device="cuda")
with torch.no_grad():
    _ = model(ids)  # forward pass builds a little KV/working memory

print(f"peak allocated:       {mib(torch.cuda.max_memory_allocated()):.0f} MiB")
# ~14 GiB resident for a 7B model in bf16 -> matches vram_for_model(7, 2).

8 · Tech-lead — VRAM as a capacity-planning discipline tech-lead

Two teams ship the same model; one is calm and one is on fire. The first over-provisioned "to be safe" and is quietly burning thousands a month on GPUs running at 40% utilization. The second sized to the edge and pages on-call every time traffic spikes. Both made the same mistake from opposite directions: they skipped the arithmetic and substituted a gut feeling. A tech lead's job on GPU memory is to make that gut feeling unnecessary — to hold a VRAM budget model in hand so hardware choices, batch limits, and cost forecasts fall out of numbers rather than nerves.

The shift in altitude here is from "will this run?" to "what's the cheapest fleet that meets the SLO?" — and it's the same four consumers, now parameterized by precision, batch, and context and multiplied across replicas. The calculators in this lesson are that model in miniature: they let a lead answer "which card, how many, at what cost per token" (IC1) in a spreadsheet instead of a procurement cycle. The discipline that separates seniors is ordering the levers by cost, cheapest first, and only escalating when the cheap ones are exhausted — because every lever trades away something (accuracy, latency, or communication overhead), and reaching for the expensive one first wastes the cheap savings you left on the table.

That ordering, roughly, runs: quantize the weights (IC2) to cut the fixed cost with minimal accuracy hit; cap context and batch to bound the KV cache; use GQA, paged KV, or KV quantization to pack more concurrency onto each card; and only then shard across GPUs (MS2), because sharding adds cross-card communication and complexity that you don't want until a single card genuinely cannot hold the model. Read this section as the executive summary of the whole page: the arithmetic you learned in sections 1–7 isn't trivia, it's the cheapest risk-reduction tool you have — cheaper than an over-provisioned fleet and far cheaper than a production OOM.

A lead doesn't launch instances to find out what fits. They keep a VRAM budget model — the four consumers, parameterized by precision, batch, and context — and use it to choose hardware, set max_model_len and batch limits, and forecast cost per token (IC1). The calculators in this lesson are that model in miniature: they turn "will it fit?" from an expensive experiment into a line of arithmetic.

The strategic levers a lead reaches for, in order of cheapness: quantize weights (IC2) to cut the fixed cost; cap context and batch to bound the KV cache; GQA / paged KV / KV quantization to serve more concurrency per card; and only then shard across GPUs (MS2) when a single card genuinely can't hold the model. Sharding adds communication cost, so it's the last resort, not the first reach.

def weight_gib(params_b, bytes_per):
    total_bytes = params_b * 1e9 * bytes_per
    return total_bytes / (1024 ** 3)

for name, p in [("7B", 7), ("13B", 13), ("70B", 70)]:
    bf16 = weight_gib(p, 2)
    int8 = weight_gib(p, 1)
    int4 = weight_gib(p, 0.5)
    print(f"{name:>4}:  bf16={bf16:6.1f} GiB  int8={int8:6.1f} GiB  int4={int4:6.1f} GiB")
# 7B: bf16 ~13.0, int8 ~6.5, int4 ~3.3 -- halving bytes halves the footprint

def kv_gib(n_layers, n_kv_heads, head_dim, seq, batch, bytes_per=2):
    elems = 2 * n_layers * n_kv_heads * head_dim * seq * batch
    return elems * bytes_per / (1024 ** 3)

# a 70B-ish config: 80 layers, 8 KV heads, head_dim 128, bf16
for b in (1, 8, 32):
    gib = kv_gib(n_layers=80, n_kv_heads=8, head_dim=128, seq=8192, batch=b)
    print(f"batch={b:>2} @ 8k ctx -> KV cache = {gib:6.2f} GiB")
# batch 1 -> 2.50, batch 8 -> 20.00, batch 32 -> 80.00 : linear in batch AND seq

def serve_gib(params_b):        # bf16 weights only
    return params_b * 1e9 * 2 / (1024 ** 3)

def train_gib(params_b):        # bf16 weight + fp32 grad + fp32 m + fp32 v
    bytes_per = 2 + 4 + 4 + 4   # = 14 bytes/param (Adam, mixed precision)
    return params_b * 1e9 * bytes_per / (1024 ** 3)

for name, p in [("7B", 7), ("13B", 13)]:
    s, t = serve_gib(p), train_gib(p)
    print(f"{name}: serve(bf16)={s:6.1f} GiB | train(Adam)={t:6.1f} GiB | {t/s:.1f}x more")
# ~7x more -- optimizer state, not the weights, dominates a training budget

def usable(card_gib, headroom=0.10):
    return card_gib * (1 - headroom)

def fits(need_gib, card_gib, headroom=0.10):
    u = usable(card_gib, headroom)
    verdict = "FITS" if need_gib <= u else "OOM"
    return f"need {need_gib:6.1f} / {u:4.1f} usable -> {verdict}"

W = lambda p, bp: p * 1e9 * bp / (1024 ** 3)
print("13B bf16 + KV 25.0 GiB    ->", fits(W(13, 2) + 25.0, 80))
print("70B bf16 weights only     ->", fits(W(70, 2), 80))
print("70B int4 weights only     ->", fits(W(70, 0.5), 80))
# 13B+KV fits; 70B bf16 OOMs (130 GiB); 70B int4 fits (33 GiB)

def triage(dominant, over_by_gib):
    # dominant term from a memory profile: 'kv', 'weights', or 'activations'
    plan = []
    if dominant == "kv":
        plan += ["lower max batch / concurrent seqs",
                 "cap context length or enable paged-KV (vLLM)"]
    elif dominant == "activations":
        plan += ["enable gradient/activation checkpointing",
                 "reduce micro-batch, use grad accumulation"]
    elif dominant == "weights":
        plan += ["quantize weights (bf16 -> int8 -> int4)",
                 "shard weights across GPUs (tensor parallel / FSDP)"]
    plan.append(f"still {over_by_gib:.0f} GiB over? move to a bigger card")
    return plan

for step in triage("kv", 12):
    print("-", step)

import math

def kv_per_user_gib(n_layers=80, n_kv_heads=8, head_dim=128, seq=8192, bytes_per=2):
    return 2 * n_layers * n_kv_heads * head_dim * seq * bytes_per / (1024 ** 3)

def plan(users, card_gib=80, weight_gib=32.6, headroom=0.10, gpu_hr=2.0):
    usable = card_gib * (1 - headroom)
    per_user = kv_per_user_gib()
    users_per_replica = int((usable - weight_gib) // per_user)
    if users_per_replica < 1:
        return "weights + one user do not fit -- need a bigger card"
    replicas = math.ceil(users / users_per_replica)
    monthly = replicas * gpu_hr * 24 * 30
    return (f"KV/user={per_user:.2f} GiB, users/replica={users_per_replica}, "
            f"replicas={replicas}, ~${monthly:,.0f}/mo")

print(plan(500))
# real per-card headroom must be MEASURED (torch.cuda.max_memory_allocated) -- needs GPU