PyTorch & Resource Accounting

Why Resource Accounting Comes Before Code

It's tempting to treat "how much will this cost" as an afterthought — something you check after the model works. CS336 puts it in Lecture 2, immediately after the course overview and before tokenization, because the opposite order is how research budgets get blown: a team spends weeks getting an architecture correct, then discovers the training run they designed needs 40 GPU-years, or that the activations alone won't fit in memory at the batch size they wanted. Resource accounting is arithmetic you can do on a napkin before writing a single line of model code, and it tells you up front which experiments are even possible on the hardware you have.

What Lecture 2 Actually Covers

PyTorch & einops, Briefly

The "PyTorch" half of Lecture 2's title is mostly about a habit, not a library feature: writing tensor operations so that their shapes are legible at the call site instead of buried in comments. Raw PyTorch reshaping code like x.view(b, h, s, d).transpose(1, 2) is correct but says nothing about what b, h, s, and d mean without reading surrounding context. The einops library replaces that with named-axis notation — rearrange(x, "b h s d -> b s (h d)") — so the operation documents its own shape transformation. This matters more than it looks: most of the actual bugs in from-scratch transformer implementations (in assignment1-basics and everywhere else) are shape-mismatch or shape-mixup bugs — attention heads getting concatenated in the wrong order, a sequence and batch axis getting silently swapped — and einops-style code makes that entire bug class visible at a glance rather than crashing three layers downstream with a cryptic matmul dimension error.

Counting FLOPs

A FLOP (floating-point operation) is one multiply or one add. The dominant cost in a transformer is matrix multiplication, and a matrix multiply of a (m × k) matrix by a (k × n) matrix costs 2·m·k·n FLOPs — one multiply and one add per output element, summed over the k-length dot product, hence the factor of 2. Every linear layer, every attention projection, and the attention score matrix itself are matrix multiplies, so this one formula is most of what you need.

For a full forward pass through a transformer with N non-embedding parameters processing one token, the standard shortcut used across the field is: forward pass ≈ 2·N FLOPs per token (each parameter participates in roughly one multiply-add per token it touches). Backward pass costs roughly twice the forward pass, because computing gradients requires propagating error back through the same matrices twice — once with respect to the inputs, once with respect to the weights — giving ≈ 4·N FLOPs per token. Forward plus backward together is commonly summarized as ≈ 6·N FLOPs per token, and total training compute for a run of D tokens is then approximately 6·N·D. This is the estimate given in Kaplan et al., "Scaling Laws for Neural Language Models" (2020, arXiv:2001.08361): non-embedding training compute C ≈ 6NBS (batch size × steps), i.e. roughly 6N FLOPs per training token, with the backward pass costing about twice the forward pass. It ignores the attention-score term, which grows with context length, so it slightly underestimates long-context training.

Counting Memory

Training memory is not just "the model." Four things typically compete for GPU memory at once, and mixed-precision training (the near-universal default) makes the accounting slightly more involved than "bytes per parameter times number of parameters":

The ZeRO family of optimizations (Rajbhandari et al., "ZeRO: Memory Optimizations Toward Training Trillion Parameter Models" (2019, arXiv:1910.02054), from Microsoft's DeepSpeed team) exists specifically to stop replicating those terms on every data-parallel GPU: ZeRO partitions optimizer state, gradients and eventually parameters across devices, so per-GPU memory for model state shrinks roughly in proportion to the number of GPUs.

Arithmetic Intensity & the Roofline

Knowing the FLOPs and the memory traffic separately isn't enough — what determines actual GPU utilization is their ratio. Arithmetic intensity is defined as FLOPs performed per byte moved between GPU memory (HBM) and the compute cores: AI = FLOPs / bytes_moved. A GPU has both a peak compute throughput (FLOPs/second) and a peak memory bandwidth (bytes/second); the "roofline" model from Williams, Waterman & Patterson, "Roofline: An Insightful Visual Performance Model for Multicore Architectures" (Communications of the ACM 52(4), 2009) says any real workload's achievable throughput is capped by whichever of those two limits it hits first. An operation with low arithmetic intensity (like an elementwise activation function, which moves a lot of data but does very little compute per element) is memory-bound: the GPU's compute units sit idle waiting for data. An operation with high arithmetic intensity (like a large matrix multiply, which reuses the same loaded values across many multiply-adds) is compute-bound: the GPU is doing useful work close to its peak FLOPs/second. This is precisely why kernel fusion (combining several memory-bound elementwise ops into one pass over the data) and why larger batch sizes (which raise the arithmetic intensity of matrix multiplies by reusing loaded weights across more rows) are two of the most effective performance levers in practice — they're both, at bottom, ways of moving an operation further right on the roofline before it hits the memory-bandwidth ceiling.

Hands-On with assignment1-basics

Reading assignment1-basics's tests/adapters.py shows it defines 21 functions — run_transformer_lm, run_multihead_self_attention, run_rmsnorm, run_get_batch, get_adamw_cls, and so on — for the model and training-loop pieces you build in that assignment, but none of them are FLOPs- or memory-accounting functions. In the actual course, that reasoning is a written-analysis section of the assignment's PDF handout (its exact wording and point values aren't reproduced here), applied to the same GPT-2-scale configuration this series' Part 1 introduced.

The practical exercise this section sets up: once you have real num_non_embedding_params for a model shape (Part 1's transformer_param_count function computes this), and a target token count and batch/sequence configuration for a training run, you can answer three questions before you run any training code — will it fit in memory, how long will it take on your hardware, and is the resulting workload compute-bound or memory-bound. The functions below do exactly that, and were executed directly (not just written) against a GPT-2-small-sized configuration to confirm the output is sane:

def transformer_training_flops(num_non_embedding_params, num_tokens):
    # Field-standard heuristic: forward ~= 2N FLOPs/token, backward ~= 2x forward,
    # so a full training step costs roughly 6N FLOPs per token processed.
    forward_flops_per_token = 2 * num_non_embedding_params
    backward_flops_per_token = 2 * forward_flops_per_token
    return (forward_flops_per_token + backward_flops_per_token) * num_tokens


def mixed_precision_training_memory_bytes(num_params, activation_bytes=0):
    # bf16 compute copies of params/grads, fp32 master weights + Adam's two moments.
    params_bf16 = num_params * 2
    grads_bf16 = num_params * 2
    master_weights_fp32 = num_params * 4
    adam_m_fp32 = num_params * 4
    adam_v_fp32 = num_params * 4
    optimizer_and_master = master_weights_fp32 + adam_m_fp32 + adam_v_fp32
    return {
        "params": params_bf16,
        "grads": grads_bf16,
        "optimizer_and_master_weights": optimizer_and_master,
        "activations": activation_bytes,
        "total": params_bf16 + grads_bf16 + optimizer_and_master + activation_bytes,
    }


def arithmetic_intensity(flops, bytes_moved):
    return flops / bytes_moved


N = 124_000_000  # GPT-2-small-ish non-embedding param count, from Part 1
tokens = 300_000_000_000

flops = transformer_training_flops(N, tokens)
mem = mixed_precision_training_memory_bytes(N)
ai = arithmetic_intensity(flops, mem["total"])

print(f"Training FLOPs: {flops:.3e}")
print(f"Memory (GB): { {k: round(v/1e9, 3) for k, v in mem.items()} }")
print(f"Arithmetic intensity (illustrative): {ai:.3e}")

# Executed output for this exact configuration:
# Training FLOPs: 2.232e+20
# Memory (GB): {'params': 0.248, 'grads': 0.248, 'optimizer_and_master_weights': 1.488, 'activations': 0.0, 'total': 1.984}
# Arithmetic intensity (illustrative): 1.125e+11

Two things worth noticing in that output. First, the optimizer-and-master-weights term (1.488 GB) is exactly 6× the params term (0.248 GB) — a direct consequence of the 12-bytes-vs-2-bytes ratio described above, and the single biggest reason "the model is only 250MB" does not mean training fits in 250MB of GPU memory. Second, the activation term is deliberately left at zero here: activation memory depends on batch size, sequence length, and how many intermediate tensors your implementation checkpoints versus recomputes, which is architecture- and implementation-specific in a way parameter count alone isn't — extending this function with a real activation estimate for your own run_transformer_lm implementation is a natural next exercise once that function is built.

⚠️ Confidence Notes and Gaps

🔗 Full Reference List