P13 — Tensor Framework and Automatic Differentiation

Run it first. There is a companion page that builds this project's machinery as numbered, independently runnable blocks and then assembles them into one measured system: P13 hands-on — block by block (handson/h*.py). Every number on it was produced by running the code. Read it alongside the milestones below.

Large · 110 hours · split across two stages · Python, with a C/Rust kernel layer

• Phase I — Autodiff core: Medium, 66 h, Weeks 21–26 (Stage 1)
• Phase II — Execution optimisation: Small, 44 h, Weeks 51–54 (Stage 2)

Table of Contents

• The Loop, Instantiated
• Why This Project Matters
• Prerequisites
• Duration and Size
• Central Technical Questions
• Reverse Mode, Derived
• Showcase — The Knob That Decides Whether a Model Fits
• Phase I — Autodiff Core (W21–W26)
• Phase II — Execution Optimisation (W51–W54)
• Concepts To Study
• Primary-Source Readings
• Experiments
• Benchmarks and Metrics
• Correctness Tests
• Failure Tests
• Expected Difficulties
• Scope Boundaries
• Deliverables
• Exit Criteria
• Extension Ideas
• Connections
• References

The Loop, Instantiated

Why This Project Matters

In P01 you called loss.backward() several thousand times. This project is where that line stops being magic.

The specific insight, which is not obvious until you build it: automatic differentiation is not calculus, it is bookkeeping. The derivative rules for individual operations are trivial — you know them. The entire engineering content is recording the operation graph, traversing it in reverse topological order, and accumulating contributions when a value is used more than once. Realising that the hard part is data structures, not mathematics, changes how you read every ML systems paper.

Phase II delivers the second insight, which matters more for your work: a framework's overhead can exceed the arithmetic it dispatches. At small tensor sizes, PyTorch spends more time deciding what to do than doing it. That is the same lesson as P02's constant-factor result and P11's dispatch measurement, and seeing it a third time in a third domain is what turns it into intuition.

Prerequisites

• Phase I: P01 — its Transformer is the correctness target and its gradients the reference
• Phase II: Phase I; P11-II helpful (eager vs graph is AST-walk vs bytecode, one level up)
• From math.md: §Chain Rule and Jacobians (3 h) — do this before milestone 1, it is genuinely load-bearing here

Duration and Size

Central Technical Questions

1. Why reverse mode? Derive the cost of forward vs reverse for \(f: \mathbb{R}^n \to \mathbb{R}^m\) and say exactly when each wins.
2. What must be stored during the forward pass, and why is that the memory wall?
3. How does broadcasting work in reverse? The backward of a broadcast is a sum, and getting the axes right is the most common autodiff bug.
4. Why is matmul's backward two more matmuls? Derive it.
5. Where does framework time go at small vs large tensor sizes?
6. What does fusion actually save? Not FLOPs — bytes.

Reverse Mode, Derived

For \(f = f_L \circ \cdots \circ f_1: \mathbb{R}^n \to \mathbb{R}^m\), the chain rule gives the Jacobian as a product:

\[ J = J_L J_{L-1} \cdots J_1 \]

You never form these matrices; you form products with vectors. Two associativity choices:

• Forward mode evaluates right to left, propagating a Jacobian-vector product (a tangent). One pass gives you one column of \(J\) — the derivative with respect to one input. Cost: \(O(n)\) passes for the full Jacobian.
• Reverse mode evaluates left to right, propagating a vector-Jacobian product (an adjoint). One pass gives you one row of \(J\) — the derivative of one output with respect to everything. Cost: \(O(m)\) passes.

Neural network training has \(n \approx 10^7\)–\(10^{11}\) parameters and \(m = 1\) scalar loss. So reverse mode costs one backward pass, and forward mode costs \(10^7\) of them. That ratio is the entire reason deep learning is computationally possible.

Concretely, for a 10M-parameter model where one forward pass takes 10 ms:

Reverse mode is roughly \(5 \times 10^6\) times faster here. And the price is memory: reverse mode must keep every intermediate activation alive until its adjoint is consumed, so memory grows with graph depth. That is the trade — time for space — and gradient checkpointing is the knob that trades back.

Why matmul's backward is two matmuls

For \(C = AB\) with upstream gradient \(\bar{C} = \partial L/\partial C\):

\[ \bar{A} = \bar{C} B^\top, \qquad \bar{B} = A^\top \bar{C} \]

Derive this by writing \(C_{ij} = \sum_k A_{ik}B_{kj}\) and applying the chain rule elementwise; the index gymnastics collapse into those two products. The practical consequence: a backward pass costs about twice a forward pass, which is where the "training is ~3× inference per token" rule of thumb (1 forward + 2 backward) comes from — the same \(C \approx 6ND\) arithmetic used in scaling-law work.

Showcase — The Knob That Decides Whether a Model Fits

Fifteen minutes with a calculator. Reverse mode's price is memory; this is the exchange rate.

# P13 -- gradient checkpointing: the memory/compute trade, computed.
import math
def plain(n):   return n, 1.0                       # store all n activations, 1 fwd
def sqrt_ck(n):
    seg = max(1, round(math.sqrt(n)))
    return seg + seg, 1.0 + 1.0                      # store ~2*sqrt(n), recompute once
print(f"{'layers':>8}{'plain mem':>12}{'ckpt mem':>10}{'saving':>9}{'extra compute':>15}")
for n in (10, 100, 1000, 10000):
    pm,_ = plain(n); cm,cc = sqrt_ck(n)
    print(f"{n:>8}{pm:>12,}{cm:>10,}{pm/cm:>8.0f}x{'+100% fwd':>15}")
print("\\nO(sqrt(n)) memory for one extra forward pass. At 1000 layers that is a 16x")
print("memory reduction for a ~33% increase in step time (fwd+bwd = 3 units; +1 fwd).")
print("\\nThis is the knob that decides whether a model FITS. Measure your own frontier")
print("in E4 -- the constants depend on your activation sizes, not on the asymptotics.")

  layers   plain mem  ckpt mem   saving  extra compute
      10          10         6       2x      +100% fwd
     100         100        20       5x      +100% fwd
    1000       1,000        64      16x      +100% fwd
   10000      10,000       200      50x      +100% fwd
\nO(sqrt(n)) memory for one extra forward pass. At 1000 layers that is a 16x
memory reduction for a ~33% increase in step time (fwd+bwd = 3 units; +1 fwd).
\nThis is the knob that decides whether a model FITS. Measure your own frontier
in E4 -- the constants depend on your activation sizes, not on the asymptotics.

A 16× memory reduction for ~33% more time, at 1,000 layers. That is not a micro- optimisation, it is the difference between a model fitting and not fitting — and it follows from the reverse-mode derivation rather than from any implementation trick.

Phase I — Autodiff Core (W21–W26)

Medium, 66 hours, 6 weeks.

The exit test for Phase I

Take the Transformer block from P01, unchanged in structure. Run one forward and one backward pass in PyTorch and in your framework, from identical initial weights and identical input. Assert:

\[ \frac{\max_i |g_i^{\text{yours}} - g_i^{\text{torch}}|}{\max_i |g_i^{\text{torch}}| + \epsilon} < 10^{-5} \]

for every parameter tensor. This is a demanding test and it will fail several times before it passes. Each failure localises a specific bug: a wrong broadcast reduction, a missing gradient accumulation for a reused tensor, a transposed matmul backward, or an in-place operation that corrupted a saved activation.

Bisect by layer. Check the gradient at the output first, then work backwards. The first layer where the discrepancy appears contains the bug.

Phase II — Execution Optimisation (W51–W54)

Small, 44 hours, 4 weeks.

What fusion actually saves

Consider d = relu(a * b + c) on tensors of \(N\) fp32 elements.

Unfused — three kernels, each reading its inputs from and writing its output to DRAM:

Arithmetic intensity: \(3N/32N = 0.094\) FLOP/byte.

Fused — one kernel reading a, b, c and writing d:

Arithmetic intensity: \(3N/16N = 0.188\) FLOP/byte — exactly 2×, and DRAM traffic halves.

Both versions are far below any ridge point (17 FLOP/byte on a server CPU, 295 on an H100 — see tools/roofline.py), so both are firmly memory-bound, and halving the bytes should halve the time. Predict a 2× speedup; measure it; explain the gap. The gap is usually launch overhead at small \(N\) and imperfect vectorisation at large \(N\).

Note what fusion did not do: the FLOP count is identical. Fusion is a data-movement optimisation. That framing is the direct bridge to P14.

The dispatch-overhead crossover

At small tensor sizes, framework overhead dominates. The measurement from tools/bench.py on a 64×64 fp32 matmul: numpy's median is 1.8 µs, of which a substantial fraction is Python call overhead, argument parsing, and dtype dispatch rather than arithmetic. The same call at 512×512 takes 160 µs and achieves 1,679 GFLOP/s through Apple Accelerate — near the hardware's practical ceiling.

So the same operation is dispatch-dominated at one size and compute-dominated at another. Find your framework's crossover in milestone 11. It is the number that determines whether graph mode is worth building at all for your workloads.

Concepts To Study

• Computational graphs: nodes, edges, topological order, dynamic vs static
• Forward vs reverse mode, and the derivation above
• VJPs and JVPs; why you never materialise a Jacobian
• Broadcasting: NumPy semantics, and the backward rule (sum over broadcast axes, keeping dims)
• Strides and views: contiguity, why transpose is free and reshape sometimes is not
• Gradient accumulation for tensors used more than once — the bug that produces gradients that are too small by an integer factor
• In-place operations and why they break autodiff
• Numerical stability: log-sum-exp, the fused softmax + cross-entropy backward
• Memory in reverse mode; gradient checkpointing as a time/space trade
• Eager vs graph execution; tracing vs source transformation
• Operator fusion and arithmetic intensity
• Memory planning: liveness analysis and buffer reuse
• Kernel dispatch: type/device/layout dispatch, and its cost

Primary-Source Readings

Budget: 12 hours.

Karpathy's micrograd is worth reading after milestone 1, as a check on your scalar design — 150 lines, and it will make you feel good about how much of it you independently invented.

Experiments

E8 is the headline of Phase II because you can derive the answer in advance. Any gap between the predicted 2× and the measured value is information, and chasing it down is exactly the bottleneck-analysis skill the whole track trains.

E4 is the most useful in practice. Gradient checkpointing recomputes activations instead of storing them, giving \(O(\sqrt{n})\) memory for one extra forward pass. Measure the actual frontier on your framework; it is the technique that decides whether a model fits in memory.

Benchmarks and Metrics

Correctness Tests

1. Finite-difference gradient check for every op, every shape, including broadcast pairs. Use central differences with \(h \approx 10^{-4}\) in float64.
2. PyTorch agreement on P01's Transformer block, 1e-5 relative.
3. Gradient accumulation: a tensor used \(k\) times receives the sum of \(k\) contributions. Test with \(k = 3\) — this bug makes gradients too small by an exact integer factor, which is the tell.
4. Broadcast backward shapes match the forward input shapes exactly, for every pair.
5. View semantics: mutation through a view is visible in the base; gradients flow correctly through views.
6. Optimizer trajectory matches PyTorch's for 100 steps to 1e-6.
7. Serialization round-trips bit-exactly.
8. Eager and graph modes agree bit-for-bit.
9. Fusion preserves semantics exactly — not approximately. Assert bitwise equality where the op order is unchanged, and document any reassociation that changes results.
10. Numerical stability: softmax with logits of ±10⁴ produces no NaN.

Failure Tests

The in-place test is the most valuable. It is the bug that produces silently wrong gradients, which is the worst failure mode in an ML framework because the model still trains, just worse.

Expected Difficulties

1. Broadcasting backward is the single biggest source of bugs. The rule: sum the gradient over every axis that was broadcast, keeping dimensions where the input had size 1. Write it once, test it exhaustively against every shape pair, and never hand-write it again.
2. Gradient accumulation is easy to miss and produces gradients too small by an integer factor — which looks like a learning-rate problem and gets "fixed" by raising the learning rate. Test 3 catches it.
3. Matching PyTorch to 1e-5 is genuinely hard. Different summation orders give different float32 results. If you cannot reach 1e-5, try float64 — if it passes there, the discrepancy is accumulation order, not a bug, and you should say so.
4. Recursion limits on deep graphs. Iterative topological sort from milestone 1.
5. Phase II may show fusion is not worth it at your sizes. That is a result.
6. The temptation to build a GPU backend is strong and it is an extension. Phase II is four weeks.

Scope Boundaries

In scope: dense CPU tensors, reverse-mode AD, ~25 ops, optimizers, a small model zoo, graph capture, fusion, memory planning, CPU threading, compiled kernels.

Out of scope: GPU (extension); distributed training; sparse tensors; complex numbers; higher-order derivatives; a full compiler with autotuning; ONNX or interoperability; quantization (P14); dynamic shapes with recompilation.

Permitted-library line: numpy for raw storage and BLAS matmul is allowed — matmul is not the mechanism under study, autodiff is. But you must implement its backward. torch is allowed only as the correctness oracle in tests, never imported by library code.

Deliverables

1. microgradpp/ (or your name) — the framework, with a model zoo
2. REPORT.md centred on E8 (fusion, derived vs measured) and E6 (the dispatch crossover)
3. The gradient-check harness — reusable, and the most valuable standalone piece
4. Notebook entries for E4, E6, E8
5. A documented op table: op, forward, backward rule, derivation reference

Exit Criteria

Phase I:

• Gradient checks pass for every op and every broadcast shape pair
• P01's Transformer block gradients match PyTorch to 1e-5, all parameters
• Optimizer trajectories match PyTorch to 1e-6 over 100 steps
• In-place-corruption detection implemented and tested
• E2 complete: forward vs reverse cost crossover measured
• E4 complete: gradient checkpointing frontier measured
• Phase I report written

Phase II:

• Eager and graph modes agree bit-for-bit
• E6 complete: dispatch/arithmetic crossover size identified
• E8 complete: fusion measured against the derived 2×, gap explained
• E9 complete: peak memory reduced by buffer reuse, quantified
• E10 complete: CPU scaling curve with the plateau explained
• Achieved GFLOP/s compared honestly against the machine's ceiling
• Phase II report written with a falsified prediction

Extension Ideas

• GPU backend (Metal on your machine, or CUDA if available). The natural bridge to P14 — and P14's roofline analysis applies directly.
• Forward-mode AD alongside reverse, with the crossover measured rather than derived. Also enables cheap Jacobian-vector products for second-order methods.
• A real fusion compiler: pattern matching over the IR, a cost model, autotuned tile sizes.
• Gradient checkpointing with automatic policy selection — choose recompute points from a memory budget. Genuinely useful and lightly explored.

Connections

Backward: P01 supplies the model, the gradients, and the correctness target.

Forward:

• → P14: the blocked-matmul kernel and its measured GFLOP/s become the accelerator simulator's CPU baseline. Fusion's arithmetic-intensity argument is P14's central theme
• → P11-II: eager vs graph is AST-walk vs bytecode, one level up. Compare your two measurements explicitly in the report — the parallel is exact
• → P15: dynamic embedding generation, and the "custom tensor operations" component

References

• Baydin, A. G., Pearlmutter, B. A., Radul, A. A., Siskind, J. M. Automatic Differentiation in Machine Learning: a Survey. JMLR 18(153), 2018.
• Griewank, A., Walther, A. Evaluating Derivatives: Principles and Techniques of Algorithmic Differentiation, 2nd ed. SIAM, 2008.
• Paszke, A. et al. PyTorch: An Imperative Style, High-Performance Deep Learning Library. NeurIPS 2019.
• Abadi, M. et al. TensorFlow: A System for Large-Scale Machine Learning. OSDI 2016.
• Chen, T., Xu, B., Zhang, C., Guestrin, C. Training Deep Nets with Sublinear Memory Cost. arXiv:1604.06174, 2016.
• Chen, T. et al. TVM: An Automated End-to-End Optimizing Compiler for Deep Learning. OSDI 2018.
• Bradbury, J. et al. JAX: composable transformations of Python+NumPy programs. 2018. For the functional-transformation alternative to tape-based AD.
• Wengert, R. E. A simple automatic derivative evaluation program. CACM 7(8), 1964. Reverse-mode AD, sixty years ago, in two pages.
• Karpathy, A. micrograd. github.com/karpathy/micrograd. Read after milestone 1.