Pick a depth. Each prompt opens in your AI pre-loaded with the lesson. Click a row to preview the prompt.
GPU memory during training splits into four categories that scale differently: parameters (fixed, one copy per model), gradients (same size as parameters, one per trainable weight), optimizer state (Adam keeps two extra fp32 buffers per parameter — momentum and variance — which is why Adam alone can triple your memory bill versus plain SGD), and activations (the intermediate tensors saved for backward, which scale with batch size and sequence length, not model size). Most first-time OOM debugging is guesswork — halving batch size and hoping — because the engineer never built the budget table that would have told them which term was actually dominant. Once you can compute this breakdown from a parameter count and a training config, an OOM stops being a mystery: you look at the table, see which term is largest, and cut that one first.
The demo estimates the four memory categories for a representative model size under mixed-precision Adam training, and prints which term dominates — a table you'll reuse verbatim in Module 5 when you start trading memory for compute.
mixed_precision=False and compare the total; explain why plain fp32 training isn't just '2x bigger weights' but changes the master-weights term too.batch and observe which single term grows — confirm it's activations_gb and not the parameter-related terms.def memory_budget_gb(num_params_b, batch, seq_len, hidden, layers, mixed_precision=True):
params = num_params_b * 1e9
bytes_per_param = 2 if mixed_precision else 4 # bf16 weights (+ fp32 master copy counted below)
param_mem = params * bytes_per_param
grad_mem = params * bytes_per_param
master_mem = params * 4 if mixed_precision else 0 # fp32 master weights
adam_mem = params * 4 * 2 # momentum + variance, always fp32
# rough activation estimate: ~34 * layers * batch * seq * hidden bytes (fp16), a common rule of thumb
act_mem = 34 * layers * batch * seq_len * hidden * 2
total = param_mem + grad_mem + master_mem + adam_mem + act_mem
return {
"params_gb": param_mem / 1e9,
"grads_gb": grad_mem / 1e9,
"master_weights_gb": master_mem / 1e9,
"optimizer_state_gb": adam_mem / 1e9,
"activations_gb": act_mem / 1e9,
"total_gb": total / 1e9,
}
budget = memory_budget_gb(num_params_b=7, batch=4, seq_len=2048, hidden=4096, layers=32)
for k, v in budget.items():
print(f"{k:20s} {v:8.1f} GB")python3 main.py