Tools & calculators

Takeaway. Backprop never writes down a global formula for dL/dw. It seeds dL/dL = 1 at the loss and then, at every edge, multiplies the gradient arriving from downstream by that edge's local derivative — the only thing each node has to know. Because every leaf here is reached by exactly one path, the products collapse to dL/dw = 2e a′(s) x, dL/dx = 2e a′(s) w, dL/db = 2e a′(s), dL/dy = −2e; when a node feeds several consumers the arriving gradients are summed instead, which is exactly why PyTorch accumulates into .grad and why you must call zero_grad(). Choose ReLU and push s below zero to watch the gate close: a′(s) = 0 annihilates every gradient upstream of it, while dL/dy — which never crosses that edge — survives. The finite-difference column is the same check torch.autograd.gradcheck runs.

Each point is a 2-D stand-in for a row of the embedding matrix W_E: a learned direction and magnitude, not a lookup key. Cosine similarity is cos θ = (A·B) / (|A|·|B|) – the dot product with both lengths divided out, so it depends only on the angle. Drag the Scale A by k slider: A slides along its own ray, so A·B and |A−B| both change while cos θ does not move at all. That invariance is exactly why embeddings are compared by cosine – a token whose vector happens to have a large norm (frequent tokens usually do) should not automatically look "more similar" to everything. The construction overlay shows the two unit vectors on the circle |v| = 1 (cos θ is literally their dot product) and the projection of B onto A, whose signed length is |B| cos θ, so that A·B = |A| × (|B| cos θ). The analogy button draws king − man + woman as vector arithmetic: the offset from man to king transplanted onto woman, with the nearest stored embedding to the landing point highlighted – and, because the two metrics rank neighbours differently, cosine and Euclidean answers are reported separately.

The takeaway: FlashAttention never forms the S×S score matrix. It walks tiles of it, and for each tile updates only three per-row statistics that live in registers/SRAM: running max m_i, running denominator ℓ_i, and unnormalized accumulator O_i. Before adding a tile's contribution it rescales the old state by exp(m_old − m_new) — that one multiply re-bases every earlier term onto the new max, so the recurrence is exact, not an approximation: watch the max|diff| column fall to ~1e-16 the moment a row's last KV block lands. The single divide by ℓ_i happens once, at the end. Raise B_r/B_c and the SRAM tile grows as B_rB_c; lower B_c and K/V get re-read more times (more HBM traffic) — that tension is the entire kernel-design tradeoff, and the naive path stays stuck at O(S²) either way.

GQA interpolates between MHA and MQA. Group size s = 1 gives every query its own KV head (MHA: best quality, biggest cache); s = h collapses all queries onto one shared KV head (MQA: smallest cache among the head-sharing schemes, largest quality risk); any divisor of h in between trades quality for cache linearly in the KV-head count g = h / s, since the cache is 2 · g · d_h elements per token per layer (the 2 is K and V). MLA takes a different axis: instead of sharing whole KV heads it caches one low-rank latent vector per token (d_c numbers, no factor of 2) and reconstructs full per-head keys/values from it on the fly – so it can beat GQA's cache while keeping every head distinct. Real MLA also caches a single shared decoupled-RoPE key of d_r = 64 (toggle it above); the text's simplified forward pass is the content-only version. Note the “total” column: the same number is both the memory one sequence occupies and the bytes that must be streamed from HBM on every decode step – which is why shrinking the cache speeds up decoding, not just fits more of it.

The surface is the exact convex bowl L(θ) = ½(λ₁θ₁² + λ₂θ₂²) with λ₁ = 1/√κ and λ₂ = √κ, so the Hessian is exactly diag(λ₁, λ₂), the condition number is κ = λ₂/λ₁ by construction, and the eigenvalues keep geometric mean 1 for every κ. Each update is the real rule applied to the analytic gradient ∇L = (λ₁θ₁, λ₂θ₂) plus optional noise σ·N(0,1) per coordinate — SGD: θ ← θ − ηĝ; momentum: v ← βv − ηĝ, θ ← θ + v; Adam: bias-corrected first/second moments with β₂ = 0.999, ε = 1e−8. Three things worth doing. (1) With σ = 0, GD on a quadratic is a linear recursion, so its per-step contraction is exactly ρ = max_i|1 − ηλ_i| (heavy ball: the larger root modulus of z² − (1+β−ηλ)z + β, which is √β whenever the roots are complex). The dashed line on the loss chart is that prediction, L₀ρ^2t, and the measured curve settles onto its slope (with momentum it oscillates around it, because the roots are complex). Divergence is not at η = 1/λ_max — that is the guaranteed-descent bound; the real ceiling is η = 2/λ_max (2(1+β)/λ_max with momentum), and just below it you get the classic zig-zag rather than an explosion. (2) The gradient ratio |g₂/g₁| = κ·|θ₂/θ₁| is enormous, and SGD's step ratio equals it exactly (the step is proportional to the gradient) — that is the zig-zag. Switch to Adam and the step ratio sits at ≈1 no matter how large the gradient ratio: that is what "Adam preconditions the ill-conditioned directions" means in numbers. (3) Raise σ: SGD stops converging and hovers at the stationary noise floor E[L_∞] = ½∑_i λ_iη²σ²/(1 − (1−ηλ_i)²) ≈ ησ²d/4 for small η — a floor proportional to η, which is exactly why the learning rate has to decay. The noise draws come from a seeded PRNG, so Reset replays a bit-identical run.

KV bytes = 2 × layers × KV-heads × head-dim × tokens × bytes (the 2 is K and V). With grouped-query attention only the KV heads count — that's the whole memory point of GQA/MLA. Total tokens = context × concurrent requests.

KV bytes = 2 × L × H_kv × d_h × seq_len × bytes/elem — the leading 2 stores both K and V, and batch size is 1 here (one sequence). Only the number of KV heads enters the formula; query heads do not. That is exactly why GQA shrinks the cache: several query heads read the same stored KV head, so fewer bytes per token, at every layer, for the whole sequence. Prefill writes the entire prompt's KV in one parallel pass (the solid block appears at once); each decode step appends exactly one token's worth of KV, so the cache grows linearly forever — which is why a long, many-turn conversation can end up costing more memory than the model's own weights.

Compare cosine against WSD at the same T: the mean LR stat shows WSD spends far more of the budget at full speed (cosine's average is only about ½(η_max+η_min), because the cosine is decaying from step T_w onward), yet Hägele et al. (2024) find they reach the same final loss at equal compute. That is the whole argument for WSD: the long stable phase commits to no total step count, and the short decay leg (shaded; typically ~10% of T) is where the characteristic extra loss drop lives. Set the min-LR ratio to 0 to reproduce MiniCPM's practice of decaying all the way down. Because the lowest-LR tokens are also the most influential, the decay window is the natural seam for a data-mixture switch — this is exactly why mid-training/annealing is scheduled inside the WSD decay phase in the capstone build (see Optimizer & Schedule: Muon + MuonClip and WSD and Mid-Training).

Activation memory for one pipeline stage can roughly double the model-state figure at steady state (chapter Worked Example: 11 GB state + 3.2 GB activations + 2 GB buffers). Treat green (<50%) as the safe target; amber means budget activations carefully. This models full ZeRO-3 / FSDP sharding across DP × TP × PP; ZeRO-1 (chapter Worked Example) shards only across TP × PP and gives a larger per-GPU figure.

Approximate: weight matrices only (ignores biases, LayerNorm, tied vs. untied LM head). 1 MAC = 2 FLOPs. FFN assumed d_ff = 4d (vanilla); gated SwiGLU adds ~50% to MLP. Matches the chapter's 24d² + 4sd per-token-per-layer accounting.

Each vector is one RoPE dimension pair rotating by angle m·θ_k. Drag m: high-frequency pairs (left) spin many times; low-frequency pairs (right) barely move – the same geometric frequency ladder as sinusoidal encodings. Hover / focus a cell for its exact numbers.

Compute-optimal means splitting a fixed budget C = 6·N·D between parameters N and tokens D. Chinchilla found the split scales as N,D ∝ √C at a ratio of ~20 tokens per parameter (adjust the ratio for over-/under-training, e.g. inference-aware training uses more). GPU-hours convert via C = GPU-hrs × 3600 × effective FLOP/s.

The heatmap is the transformer: each row is one query's probability distribution over all keys (random here, learned in a trained model), and the context vector below is literally that distribution used as mixing weights over the value vectors — so the output can never leave the convex hull of V. Watch the row sum stay pinned at 1.000: softmax normalizes over keys, never over queries. Turn the scale off and the score spread jumps by exactly √8 ≈ 2.83×, the rows sharpen toward one-hot, and the effective-key count collapses — that is softmax saturating into its low-gradient regime, which is the whole reason for the 1/√d_k factor. This toy has no positional encoding, so two copies of the same word get identical embeddings; only the causal mask makes their rows differ.

Verification is exact, not approximate: a drafted token x ~ q is accepted with probability min(1, p(x)/q(x)); on rejection the replacement is resampled from the residual (p−q)⁺. Summed over both paths a token is emitted with probability min(p,q) + (p−q)⁺ = p(x) — the target's own distribution — no matter how bad the drafter q is. α changes how many tokens come out per parallel pass, never which distribution they come from. With i.i.d. acceptance, E[tokens per target call] = (1−α^γ+1)/(1−α) and cost = 1+γc target-equivalent passes, so speedup = E/(1+γc). The γ-th extra draft slot buys only α^γ expected tokens but always costs c, so once α^γ+1/c drops below the current speedup, growing γ hurts — that is the γ* marked on the chart. Caveats this idealization hides: acceptance is not really i.i.d. (it decays along the chain), single-chain drafting is weaker than tree drafting, and at large batch sizes decode is compute-bound so the “free” verification FLOPs stop being free.

The takeaway: byte-level BPE guarantees every possible input is encodable — it bottoms out at raw bytes, so there is no <unk> token — while learned merges compress frequent substrings into single tokens. Push the merge slider and watch bytes / token climb: that ratio is your context-window and inference-cost multiplier. Merges are chosen purely by adjacent-pair frequency in the training corpus, so a tokenizer inherits its corpus's biases (English prose compresses well; unseen scripts fall back toward one token per byte). Decoding is just concatenating each token's stored bytes, which is why encode → decode reproduces the input's bytes exactly, even though a single merged token's bytes need not be valid UTF-8 on their own.

Training FLOPs use the standard 6ND rule (forward+backward, dense). Memory uses the Adam mixed-precision 16 bytes/parameter optimizer-state budget (fp16 param+grad + fp32 master+m+v), sharded evenly across GPUs (ZeRO-3/FSDP) — activation memory is extra and not shown.