Handout #7 — Batching and Training Memory

A broadcast can silently change your objective, and a small batch can still exhaust GPU memory. This reading connects tensor shapes, gradient buffers, and the lifetime of saved computations so you can reason about what your training loop actually does.

Use the worked examples to predict what a training run will do, then test the explanation. Try each exercise before opening its answer.

1Execution modes, batching, and memory

Slides 165–182

1.1  Three switches with three different jobs

Slides 166–167

A module's training flag, autograd's recording mode, and a parameter's gradient buffer are separate pieces of state. Changing one does not automatically change the others. A useful way to read a training loop is to ask three questions: which forward behavior is active, is a backward graph being built, and what is already in the gradient buffers?

backward() adds to leaf gradient buffers. For ordinary training, clear those buffers before computing the next update's gradients; use model.train() when resuming training after validation. Validation usually combines eval() and no_grad(). It needs no gradient clearing if it performs no backward pass.

import torch
from torch import nn

torch.manual_seed(4)
model = nn.Sequential(nn.Linear(3, 8), nn.ReLU(), nn.Dropout(0.5), nn.Linear(8, 1))
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
x, target = torch.randn(4, 3), torch.randn(4, 1)

model.train()
optimizer.zero_grad(set_to_none=True)
loss = (model(x) - target).square().mean()
loss.backward()
optimizer.step()
old_grad = model[0].weight.grad.clone()

model.eval()
assert model(x).requires_grad  # eval does not disable autograd
with torch.no_grad():
    prediction = model(x)
assert not prediction.requires_grad
torch.testing.assert_close(model[0].weight.grad, old_grad)
optimizer.zero_grad(set_to_none=True)
assert all(p.grad is None for p in model.parameters())
model.train()  # ready for the next training step

Reference: PyTorch's grad modes and evaluation mode.

1.2  Batching: give every axis a meaning

Slides 168–172

A batch applies the same parameters to multiple examples. Vectorized tensor operations let the hardware process those examples together. For language data we use $X\in\mathbb R^{B\times T\times d}$: batch, sequence position, embedding coordinate. Token IDs have shape $(B,T)$; looking up each ID's embedding introduces the final feature axis. Other axis orders work if operations follow that convention. An axis's size does not tell you its meaning.

The multiplication acts on the last axis, while the bias broadcasts over examples and positions. Its leading cost is $BTdh$ multiply-accumulates. Larger batches can improve hardware utilization, but their saved activations also consume more memory.

import torch

torch.manual_seed(0)
B, T, d, h = 2, 3, 4, 5
X = torch.randn(B, T, d)
W = torch.randn(h, d)
b = torch.randn(h)
Y = X @ W.T + b
loop = torch.stack([torch.stack([W @ token + b for token in seq]) for seq in X])
assert Y.shape == (B, T, h)
torch.testing.assert_close(Y, loop)
print("Vectorized and loop results agree:", Y.shape)

1.3  Broadcasting: compatible shapes can still mean the wrong thing

Slides 173–178

For elementwise arithmetic, compare axes from right to left: paired sizes must match or one must be 1; absent leading axes act like size 1. Thus a feature bias of shape $(d,)$ can be added to $(B,T,d)$ without writing a loop or allocating repeated copies of that bias. The output still needs storage. Broadcasting knows sizes, not whether an axis means “example” or “feature.”

The following regression bug is especially dangerous: subtraction succeeds, but compares every prediction with every target. Averaging afterward hides the extra axis.

import torch

predictions = torch.tensor([[10.], [20.], [30.]])  # (3, 1)
targets = torch.tensor([11., 22., 33.])             # (3,)
wrong = predictions - targets                     # (3, 3)
assert wrong.shape == (3, 3)
assert torch.isclose(wrong.square().mean(), torch.tensor(454. / 3.))

# Choose one consistent convention: one scalar per example.
matched = predictions.squeeze(-1)                 # (3,), including when B=1
assert matched.shape == targets.shape
errors = matched - targets
assert errors.tolist() == [-1., -2., -3.]
assert torch.isclose(errors.square().mean(), torch.tensor(14. / 3.))
# Alternatively: predictions - targets[:, None] gives shape (3, 1).

The wrong objective is $B^{-2}\sum_{i,j}(\hat y_i-y_j)^2$, whereas the intended one is $B^{-1}\sum_i(\hat y_i-y_i)^2$. The wrong objective pushes every prediction toward the batch's mean target, rather than its paired label. Assert prediction/target shapes before a pointwise loss; use squeeze(-1) rather than dropping every size-1 axis with squeeze().

Arithmetic, comparisons, and operations such as torch.where use elementwise broadcasting. Matrix multiplication and advanced indexing use broadcasting on different axes:

import torch

# matmul contracts the inner matrix dimension (4), broadcasts leading dimensions.
a = torch.randn(32, 10, 3, 4)
b = torch.randn(10, 4, 5)
assert (a @ b).shape == (32, 10, 3, 5)

# Advanced indexing broadcasts the index tensors to form an index grid.
grid = torch.arange(12).reshape(3, 4)
rows = torch.tensor([[0], [2]])   # (2, 1)
cols = torch.tensor([[1, 3]])     # (1, 2)
assert grid[rows, cols].tolist() == [[1, 3], [9, 11]]

Reference: PyTorch broadcasting semantics.

1.4  What occupies GPU memory, and when?

Slides 179–182

In a fresh inference process, memory holds model parameters, live intermediate activations, inputs/outputs, and temporary workspace. Without graph recording, intermediates can be released once no operation or reference needs them. During training, backward needs saved tensors from the forward pass, so many more remain live. Merely switching an existing training process to eval() does not delete its optimizer state or old gradients.

Training memory includes parameters, parameter gradients, optimizer state, saved activations, and temporary workspace. For a simple float32 Adam accounting, parameters, gradients, and two moment buffers use roughly $16P$ bytes for $P$ parameters, before activations and temporary buffers. Mixed precision and implementation choices change this estimate.

Peak memory varies during a step. Optimizer state is often allocated lazily at the first step; backward typically releases saved graph tensors, but parameter .grad buffers persist until cleared. zero_grad(set_to_none=True) clears those references. Validation should use model.eval() and torch.no_grad(): the first changes layer behavior, the second disables graph recording. Holding graph-connected losses in a list can retain graphs; logging loss.item() avoids that reference.

Read the figure from the bottom up. The green base is present throughout. The first persistent yellow band marks optimizer-state allocation; after that, each cycle starts from a higher baseline. Within a cycle, saved forward activations rise before backward consumes them and builds gradients. The tallest total stack matters for OOM, even if that peak lasts only briefly.

The allocator may keep freed storage reserved for reuse, so a GPU monitor need not show an immediate drop. An out-of-memory error depends on the peak live allocation and available memory, not an average over the step. Its location is informative: an error on the first Adam step may be optimizer-state allocation, while one late in forward often points to saved activations.

1.5  Responding to OOM: smaller live batches, same objective

Start by reducing batch size or sequence length when activation memory dominates. For evaluation, use no_grad() to avoid saving a backward graph and eval() for evaluation behavior. Do not place a training forward pass in no_grad() just to save memory: that removes the path needed to compute its gradients.

Gradient accumulation performs backward after each small microbatch, releasing its saved tensors before processing the next. Keep the weights fixed until all microbatches in the update are processed. For an additive per-example objective,

If microbatches have equal sizes, divide each mean loss by $K$. If their sizes differ, weight by example count. For a token-averaged language-model loss, count valid target tokens instead. The following checks the gradients against one full batch, including an uneven final microbatch.

import copy
import torch
from torch import nn

torch.manual_seed(5)
full = nn.Linear(3, 2).double()
small = copy.deepcopy(full)
x = torch.randn(5, 3, dtype=torch.float64)
y = torch.tensor([0, 1, 0, 1, 1])
full.zero_grad(set_to_none=True)
nn.functional.cross_entropy(full(x), y).backward()

optimizer = torch.optim.SGD(small.parameters(), lr=0.1)
optimizer.zero_grad(set_to_none=True)
for xb, yb in zip(x.split(2), y.split(2)):  # sizes 2, 2, 1
    loss = nn.functional.cross_entropy(small(xb), yb, reduction="sum") / len(y)
    loss.backward()                       # add this microbatch's contribution
for p_full, p_small in zip(full.parameters(), small.parameters()):
    torch.testing.assert_close(p_full.grad, p_small.grad)
optimizer.step()                          # one update for all five examples
optimizer.zero_grad(set_to_none=True)

This equality assumes examples do not interact through batch-dependent operations. BatchNorm statistics change with microbatching; dropout masks and floating-point order can also produce numerical differences. Accumulation saves activation memory, but does not shrink the model, its gradients, or its optimizer state. Summing graph-connected losses and calling backward only at the end retains all those graphs and defeats the main memory benefit.

Activation checkpointing instead saves selected boundary tensors and recomputes intermediates during backward. It trades extra computation for fewer saved activations and can be combined with accumulation. Mixed precision reduces storage for some tensors; distributing model state or offloading optimizer state addresses other memory costs. These methods have their own numerical and performance tradeoffs and are later-course topics.

2Optional application: the language-model loss

Optional extension beyond slides 165–182, applying batching and evaluation to language models.

Handout #6 §3.1 gave us a training loop with a loss in it. For a language model that loss is cross-entropy, and Handout #2 §4.4 computed it by hand. You will never do that again — but you will compute it in PyTorch constantly, and there are three ways to get it wrong that are worth inoculating against before you write your own loop.

2.1  It is one line — and it is already the loss

Handout #1 §1.5 established that F.cross_entropy(logits, target) is $-\log p(\text{target})$. Handout #2 §4.1 established that perplexity is $2^H$ where $H$ is that quantity averaged. Put them together:

API sketch: assumes a model, inputs, and targets from your training loop.

import torch
import torch.nn.functional as F

logits  = model(inputs)                     # (T, V) -- one row of scores per position
targets = ...                               # (T,)   -- the word that actually came next

loss = F.cross_entropy(logits, targets)     # mean negative log-probability
ppl  = torch.exp(loss)                      # perplexity

# Verify the identity by hand:
logp   = F.log_softmax(logits, dim=-1).gather(1, targets[:, None]).squeeze(1)
assert torch.allclose(loss, -logp.mean())   # cross_entropy IS mean -log p

2.2  Trap: the off-by-one

A causal model predicts token $t$ from tokens $< t$, so the inputs and the targets are the same sequence, shifted by one. This is the objective-constructor widget from Handout #1 §2.1, in two lines:

x = torch.tensor([[5, 9, 2, 7, 4]])         # one sequence of 5 tokens

inputs  = x[:, :-1]                         # [[5, 9, 2, 7]]  -- drop the last
targets = x[:,  1:]                         # [[9, 2, 7, 4]]  -- drop the first

logits = model(inputs)                      # (B, T-1, V)
loss = F.cross_entropy(
    logits.reshape(-1, logits.size(-1)),    # (B*(T-1), V) -- cross_entropy wants 2-D
    targets.reshape(-1),                    # (B*(T-1),)
)

A length-5 sequence yields 4 predictions, not 5 — there is nothing to predict the first token from, and nothing after the last to predict. Forget the shift and you are training the model to copy its input, which produces a beautifully low loss and a completely useless model. Note also the reshape: cross_entropy wants (N, V) and (N,), so the batch and time axes get flattened together. Perplexity does not care, because it averages over all positions anyway.

2.3  Trap: padding, and the flattering number it gives you

Real batches contain sequences of different lengths, padded to a common length. Padding positions can become easy to predict. If their losses are lower than those of real targets, including them drags reported perplexity down. If their losses are higher, it pushes perplexity up. Either way, padding is not part of the intended evaluation target set.

import torch
import torch.nn.functional as F

PAD_ID, V = 0, 16
# Already aligned target positions: 23 real targets and 17 padding targets.
flat_targets = torch.tensor([1] * 23 + [PAD_ID] * 17)
# Construct probabilities so real-target perplexity is 12.7,
# while padding is predicted with probability 0.999.
correct_p = torch.tensor([1 / 12.7] * 23 + [0.999] * 17, dtype=torch.float64)
probs = ((1 - correct_p) / (V - 1))[:, None].repeat(1, V)
probs[torch.arange(40), flat_targets] = correct_p
flat_logits = probs.log()  # log probabilities are valid logits
masked = F.cross_entropy(flat_logits, flat_targets, ignore_index=PAD_ID)
unmasked = F.cross_entropy(flat_logits, flat_targets)
assert abs(torch.exp(masked).item() - 12.7) < 1e-10
assert 4.3 < torch.exp(unmasked).item() < 4.4
print(torch.exp(masked).item(), torch.exp(unmasked).item())  # 12.7, about 4.31

On that batch — 42% padding, and a model that has learned to predict padding almost perfectly — the unmasked figure understates perplexity by about two thirds. (The direction and size of the bias depend on the losses on padding versus real targets.) And nothing errors: you get a number, it looks good, and it is wrong. Exclude padding targets with ignore_index or an explicit mask, and sanity-check the denominator. Continue the example above:

total = F.cross_entropy(flat_logits, flat_targets,
                        ignore_index=PAD_ID, reduction="sum")
n_real = (flat_targets != PAD_ID).sum()          # 23, not 40

assert torch.allclose(total / n_real, masked)    # summed real-target loss / 23

This is the same averaged-versus-summed distinction as Handout #1 Exercise 1.5b, with a subtler failure mode: there the risk was dividing by the wrong number, here it is counting the wrong tokens. Whenever you see a perplexity, ask what $n$ was.

2.4  Finding the token that cost you

Because $H$ is an average of per-token surprisals, you can always break it back apart — useful when a model is doing well overall and failing somewhere specific:

Continue the unpadded 2-D logits, targets, and loss from §2.1; use the same evaluated positions in both averages.

import math

logp = F.log_softmax(logits, dim=-1).gather(1, targets[:, None]).squeeze(1)
bits = -logp / math.log(2)                  # surprisal per token, in bits

worst = bits.argmax()
print(f"token {worst.item()} cost {bits[worst]:.2f} bits")

# and the average of those bars is exactly the loss:
assert torch.allclose(2 ** bits.mean(), torch.exp(loss))

That last assertion is Handout #2 §4.1's two definitions meeting in code: the mean surprisal in bits, exponentiated base 2, equals the mean loss in nats, exponentiated base $e$. If it ever fails, check the log bases, reductions, and evaluated positions.

2.5  Application questions