PyTorch and autograd fluency

Why it matters

In research code, wrong gradients rarely crash; they quietly train a worse model. A broadcast that turns a per-example loss into a matrix, a tensor modified in place after it was saved, or an evaluation loop that keeps building graphs can each cost days. Interviewers for research-engineering and research-ML roles therefore ask you to predict shapes, explain what backward() does, and implement or check a gradient. The relevant PyTorch notes are short; this topic is about turning them into reflexes.

Core ideas

Tensors and views. A tensor is storage plus shape, strides and an offset. transpose, view, slicing and expand return views that share storage; clone and contiguous copy. Many surprises, such as an in-place write that changes another tensor or view() failing after a transpose, come from forgetting which operations alias.

Broadcasting. Shapes are aligned from the right; each pair of dimensions must be equal, or one of them must be 1 or missing. Predictions of shape (N, 1) minus targets of shape (N,) silently produce an (N, N) result whose mean is a meaningless loss. Assert shapes at module boundaries and name dimensions in comments.

The graph. When an operation involves a tensor with requires_grad set, autograd records a node (visible as grad_fn) and saves whatever the backward formula needs. Leaves are tensors you created, usually parameters; by default only leaves receive .grad, and gradients accumulate across backward calls, which is why training loops zero them and why gradient accumulation over micro-batches works. backward() walks the graph in reverse applying vector-Jacobian products and frees saved tensors unless retain_graph is set. Saved activations are the main source of training memory.

In-place operations. Every tensor has a version counter that in-place operations increment. If a tensor saved for backward changed before backward runs, autograd raises an error instead of using stale values. Writing through .data bypasses that check and can produce wrong gradients with no error at all.

Turning tracking off. torch.no_grad() stops recording, for evaluation and for parameter updates; inference_mode() is stricter and faster, but its outputs cannot later join autograd; detach() returns a view cut from the graph; model.eval() changes layer behaviour (dropout, batch-norm statistics) and does not disable gradients. Mixing these up causes "validation uses dropout" and "evaluation leaks memory" bugs.

Custom gradients. Subclass torch.autograd.Function, compute the output in forward, save what backward needs with save_for_backward, and return one gradient per input from backward (None for non-differentiable inputs). Do this only when built-in operations cannot express the computation or when you want a cheaper or more stable backward, for example recomputing an intermediate instead of storing it.

Gradient checking. Compare analytic gradients with central differences, (f(x + h) − f(x − h)) / 2h. In float32, rounding error swamps the comparison, so check in float64 on small inputs; gradcheck automates this by comparing full Jacobians. Choose test points away from kinks such as ReLU at zero.

x = torch.randn(4, 3, dtype=torch.float64, requires_grad=True)
torch.autograd.gradcheck(StableLogSumExp.apply, (x,))

Determinism. Seed Python, NumPy and torch; seed data-loader workers; disable cuDNN autotuning; and request deterministic algorithms, which raises an error for operations without a deterministic implementation. Results can still differ across PyTorch versions, hardware, and CPU versus GPU, and some fused attention backward passes are nondeterministic by default. Determinism is a debugging aid; statistical claims still need several seeds.

Common misconceptions

• "model.eval() turns off gradients." It only switches layer modes; wrap evaluation in no_grad or inference_mode as well.
• "Each backward call overwrites .grad." Gradients accumulate; forgetting to zero them silently changes the effective update.
• "If gradcheck passes in float32, the backward is right." Float32 finite differences are too noisy to be informative; check in float64, at several points, away from kinks.
• "Seeding makes runs identical everywhere." Seeds fix random streams, but nondeterministic kernels, other hardware or library versions, and data-loader ordering can still change results.
• "Editing weights through .data or detach() is harmless." Both hide the operation from autograd; changing values that were saved for backward this way can corrupt gradients without any error.

How to approach it

Build a small autograd engine first, then extend it to tensors so that broadcasting forces you to sum gradients over expanded dimensions. Reimplement two operations (log-sum-exp, softmax cross-entropy) as custom Functions and gradcheck them. In interviews, narrate shapes aloud and say where the graph is created and freed; senior candidates are expected to discuss activation memory, determinism costs and how they would test numerics in reduced precision.