The Differentiable Scaling Book
August 2026
Abstract
How to Scale Your Model is a book about JAX, and the soul of JAX is the gradient. This page rebuilds the book's quantitative content, rooflines, collective costs, transformer arithmetic, training and inference regimes, as one computation graph, then differentiates it with a reverse-mode autodiff written in the page source. Click any derived value to see the gradient of that value with respect to every knob: which constraint binds, which hardware upgrade is worth money, and which is worthless, at your exact operating point. A verifier in section 7 re-derives thirteen of the book's own numbers to show the graph is honest. Everything on this page computes live.
1Hardware
Specs as printed in the book's TPU and GPU chapters. Every field is an input node of the graph; edit any of them. W_ici is the bidirectional collective bandwidth per mesh axis.
2Model
Gated-MLP transformer with grouped-query attention, counted as in the transformers chapter: P_layer = 3DF + 2DH(N_q+N_kv), N = L·P_layer + 2VD. Layernorms are dropped (they are ~10⁻⁵ of N).
3Training
Mesh of X data/FSDP shards × Y tensor shards. Per-layer costs from the training chapter; collectives cost V/(W_ici·axes) per the sharding chapter, comms overlap with math, so a layer takes max(T_math, T_comms). With Y>1 tensor parallelism takes one ICI axis and FSDP the rest.
The training chapter's thresholds, drawn exhaustively: chip count held fixed, per-chip batch on the horizontal, tensor parallelism on the vertical, brightness = MFU. The ring is your current configuration; click anywhere to move it.
4Inference
Greedy generation on Y chips of the same type, tensor-parallel, per the inference chapter: each step streams the weights and the KV cache, so T_step = max(T_flops, T_weights + T_KV) + T_comms. Generation is memory-bound below the chip's critical batch; prefill is compute-bound.
Batch against chip count, brightness = fraction of the step spent on useful FLOPs. The cheap-tokens region is bright; the low-latency region is bottom-left. Click to move the ring.
5The gradient
Reverse-mode automatic differentiation through everything above. Bars show ∂log(output)/∂log(input): the % change in the focused output per % change in each input. max() routes gradient only through its binding branch, so slack constraints report exactly zero, the regime structure of the book, recovered as sparsity in the gradient.
6Goal-seek
The inverse problem, solved with the same gradients: pick a target, free some knobs, and the page runs gradient descent in log-space, u ← u − η·(log y − log y*)·e, where e is the elasticity.
7Verifier
Thirteen statements from the book (and two vendors' ground truth) recomputed by this page's graph, live. If you edit the hardware table above, these do not change; they always use the book's printed specs.
| statement | recomputed here | as stated | via | source |
|---|
8What this model ignores
- All times are roofline lower bounds: max(compute, comms) with perfect overlap and bandwidth-optimal collectives. The only overhead kept is a 1 µs latency floor per inference collective.
- The matmul-efficiency knob scales peak FLOPs; bandwidths are taken at 100% of spec.
- Attention dot-products are counted causally (half of 12·B·S·N_q·H per layer); layernorms, rotary embeddings and biases are dropped.
- Pipeline parallelism and cross-pod DCN are not modeled. The mesh is a single ICI or NVLink domain; choose the H100 SuperPod entry to price cross-node InfiniBand instead.
- Training memory is 10 bytes per parameter under Adam (bf16 weights + fp32 moments) plus the book's checkpointed-activation term 2·L·B·(D+2F)/chips.
- GPUs are modeled as one collective "axis" at the egress bound V/W; the (N−1)/N factor and NVLink SHARP in-network reductions are dropped.
- X, Y, L are differentiated as continuous; read their gradients as marginal rates for doubling arguments, then round to powers of two.
- KV cache is dense GQA: no paged sharing, no sliding window, no prefix cache.