Case Study: Making a Transformer Fast

The Method, on a Real Model

Nothing new here — that is the point. Six sections built a method and a toolbox; now run the whole loop on ch. 11’s GPT and take it down a waterfall, one rung per technique, each measured and attributed.

Subject: d2l.GPT, width 512 (~19M params), char Time Machine, context 128, end-to-end tokens/s. Same model and metric as the PyTorch tab — the ladder will look nothing like it. That contrast is the lesson.

Measure Honestly First

Traps that corrupted early drafts — JAX edition:

  • a changed shape is not a retrace, it is a recompile → whole batches only, one shape
  • the stock tf.data loader taxed the fast rungs by ~¼ → stage the 5 MB corpus on device, shuffle there
  • no block_until_ready, no measurement — you timed the enqueue (§13.1)
  • never re-time one cached batch: the GPU serves it from cache. Honest data is fresh data.

Rungs 0 → 1: The Only Real Baseline

Un-jitted JAX is a strawman — measure it once, briefly, to price the mistake; nobody trains there.

step_jit = nnx.jit(train_step)         # the rung is one transformation

t0 = time.perf_counter()
jax.block_until_ready(step_jit(model, optimizer, *next(stream)))
print(f'first (compiling) step: {time.perf_counter() - t0:.1f} s')
fwd = nnx.jit(lambda m, X: m(X))
X, Y = next(stream)
assert jnp.allclose(fwd(model, X), model(X), atol=1e-2, rtol=1e-3)
tput1 = throughput(lambda X, Y: step_jit(model, optimizer, X, Y), stream)
print(f'R1 jit: {tput1:.0f} tokens/s ({tput1 / tput0:.0f}x)')
first (compiling) step: 7.0 s
R1 jit: 253478 tokens/s (19x)

nnx.jit on the whole step (loss + grads + update) is one transformation and roughly 20×. Every honest ratio below starts from this bar, not from R0.

Classify the Compiled Step

Ask the compiler what it plans to run:

compiled = step_jit.lower(model, optimizer, X, Y).compile()
flops = compiled.cost_analysis()['flops']
t_step = 64 * 128 / tput1
print(f'XLA counts {flops / 1e9:.0f} GFLOP per step; at {1e3 * t_step:.1f} '
      f'ms per step that is {flops / t_step / 1e12:.0f} TFLOP/s achieved')
XLA counts 765 GFLOP per step; at 32.3 ms per step that is 24 TFLOP/s achieved

About a third of the tf32 roof (§13.2): not overhead-starved, not saturated — a mixed workload. Cheapest next rung on §13.4’s ladder: bf16 — double the roof, half the bytes.

Rung 2: bf16 Is a Discipline, Not a Cast

Cast every array to bf16 → measures exactly like fp32. Silent. The receipt: the embedding still emits fp32 — flax modules remember a compute dtype; rope’s fp32 trig promotes q, k right back.

def to_bf16(model):
    """A bf16 compute copy: cast every floating array AND the compute
    dtype each flax module remembers from construction."""
    copy = bf16_arrays(model)
    for _, m in nnx.iter_graph(copy):
        if isinstance(m, nnx.Module) and hasattr(m, 'dtype'):
            m.dtype = jnp.bfloat16
    return copy

class Bf16GPT(d2l.GPT):
    """ch. 11's GPT with one dtype leak patched: rope's fp32 trigonometry
    promotes q and k back to fp32 (and the attention kernel rightly
    refuses mixed inputs), so cast them back after rope."""
    class CausalAttention(d2l.GPT.CausalAttention):
        def __call__(self, X, *_):
            B, T, D = X.shape
            q, k, v = jnp.split(self.W_qkv(X), 3, axis=-1)
            q, k, v = (u.reshape(B, T, self.num_heads, -1)
                       for u in (q, k, v))
            if self.rope:
                q, k = self._rope(q), self._rope(k)
            q, k = q.astype(v.dtype), k.astype(v.dtype)
            Y = jax.nn.dot_product_attention(q, k, v, is_causal=True)
            return self.W_o(Y.reshape(B, T, D)), None

def make_bf16_gpt():
    return Bf16GPT(vocab_size, num_hiddens=512, num_blks=6,
                   rngs=nnx.Rngs(0))

model16 = make_bf16_gpt()
opt16 = nnx.Optimizer(model16, optax.adamw(1e-3), wrt=nnx.Param)

@nnx.jit
def step_bf16(model, optimizer, X, Y):
    def loss_fn(model):
        logits = to_bf16(model)(X)         # bf16 compute copy of the step;
        return optax.softmax_cross_entropy_with_integer_labels(
            logits.reshape(-1, vocab_size).astype(jnp.float32),  # fp32 loss
            Y.reshape(-1)).mean()
    loss, grads = nnx.value_and_grad(loss_fn)(model)   # grads arrive fp32
    optimizer.update(model, grads)
    return loss

print('the receipt, again:', to_bf16(model16).token_emb(X).dtype)
plan16 = step_bf16.lower(model16, opt16, X, Y).compile().memory_analysis()
print(f'planned temp: bf16 {plan16.temp_size_in_bytes / 2**20:.0f} MiB '
      f'(fp32: {compiled.memory_analysis().temp_size_in_bytes / 2**20:.0f} '
      f'MiB)')
tput2 = throughput(lambda X, Y: step_bf16(model16, opt16, X, Y), stream)
assert all(p.dtype == jnp.float32                      # fp32 masters intact
           for p in jax.tree.leaves(nnx.state(model16, nnx.Param)))
print(f'R2 +bf16, threaded: {tput2:.0f} tokens/s ({tput2 / tput1:.2f}x)')
the receipt, again: bfloat16
planned temp: bf16 1149 MiB (fp32: 1989 MiB)
R2 +bf16, threaded: 242897 tokens/s (0.96x)

Threaded for real (receipt: bfloat16; the planned temporaries drop ~40%) — and still flat: at batch 64 the wall is per-step dispatch, which precision cannot touch. Masters and grads stay fp32. No GradScaler.

Rungs 3–4: Rungs Interact

big_stream = batches(data, batch_size=512)   # same corpus, bigger gathers
Xb, Yb = next(big_stream)

# The bf16-512 step (~9 GiB) fits alongside the pool the ladder has already
# grown; the fp32-512 control (~15 GiB) did not, which is why it was clocked
# up front (`ctl512`, `plan32_gib`). Here we run bf16-512 and read its plan.
tput3 = throughput(lambda X, Y: step_bf16(model16, opt16, X, Y),
                   big_stream, warmup=30, timed=30)
c16 = step_bf16.lower(model16, opt16, Xb, Yb).compile()
plan16_gib = c16.memory_analysis().temp_size_in_bytes / 2**30
tf512 = c16.cost_analysis()['flops'] / (512 * 128 / tput3) / 1e12
print(f'R3 +batch-up (512): {tput3:.0f} tokens/s ({tput3 / tput2:.2f}x), '
      f'{tf512:.0f} TFLOP/s achieved')
print(f'fp32 control at 512: {ctl512:.0f} tokens/s '
      f'-> bf16 pays {tput3 / ctl512:.2f}x here')
print(f'planned temp: bf16 {plan16_gib:.1f} GiB '
      f'(fp32 control: {plan32_gib:.1f} GiB)')
R3 +batch-up (512): 434987 tokens/s (1.79x), 44 TFLOP/s achieved
fp32 control at 512: 216489 tokens/s -> bf16 pays 2.01x here
planned temp: bf16 8.9 GiB (fp32 control: 15.1 GiB)

Neither rung pays alone. fp32 at 512 is slower than at 64 — the ~15 GiB plan is traffic, a memory wall. bf16 at 64 was flat — a dispatch wall. Together: most of 2×, ~43 TFLOP/s. Measure ladders cumulatively. Checkpointing then costs ~a tenth of the speed for memory the plan says we did not need — the same negative rung as the PyTorch tab.

The Waterfall

rungs = ['R0\nun-jitted', 'R1\njit', 'R2\n+bf16', 'R3\n+batch', 'R4\n+ckpt']
tputs = [tput0, tput1, tput2, tput3, tput4]
print(f'cumulative, R1 -> R3: {tputs[3] / tputs[1]:.2f}x '
      f'(and R0 -> R3: {tputs[3] / tputs[0]:.0f}x)')
d2l.plt.figure(figsize=(6, 3.5))
bars = d2l.plt.bar(rungs, [t / 1e3 for t in tputs],
                   color=['#7f7f7f', '#1f77b4', '#1f77b4', '#2ca02c', '#d62728'])
d2l.plt.ylabel('throughput (k tokens/s)')
d2l.plt.title('Making a Transformer fast: cumulative rungs')
for b, t in zip(bars, tputs):
    d2l.plt.text(b.get_x() + b.get_width() / 2, t / 1e3,
                 f'{t / tputs[1]:.2f}x', ha='center', va='bottom')
d2l.plt.show()

cumulative, R1 -> R3: 1.72x (and R0 -> R3: 33x)

jit is the cliff (~20×). bf16 stalls at batch 64 — the flat bar that taught the most. Batch-up cashes both rungs at once; checkpointing is red, as in the other tab. Same model, same card — a different bottleneck sequence, found by the same loop. A 300-step run confirms the fast configuration still learns.

Predict, Then Measure: Data Parallel

~76 MB of gradients vs §13.5’s measured psum β ⇒ the two-term cost model prices k=2 at no gain at all. Then measure: mesh + shard the batch + the unchanged jitted step — no launcher, no sidecar.

n_bytes = 4 * n_params
beta = 4.5e9   # psum, effective bytes/device/s on this box (13.5)
t_comm, t_cmp = 2 * n_bytes / beta, 64 * 128 / tput1
print(f'per step: t_comm ~ {1e3 * t_comm:.0f} ms vs '
      f't_compute ~ {1e3 * t_cmp:.0f} ms')
for k in (k for k in (2, 4) if k <= jax.local_device_count()):
    floor = k * 64 * 128 / (t_cmp + t_comm)
    meas = dp_tokens(k, step_jit, make_gpt, stream)
    print(f'R5 k={k}: no-overlap floor {floor / 1e3:.0f}k tokens/s, '
          f'measured {meas / 1e3:.0f}k ({meas / tput1:.2f}x of one GPU)')
per step: t_comm ~ 34 ms vs t_compute ~ 32 ms
R5 k=2: no-overlap floor 249k tokens/s, measured 199k (0.79x of one GPU)
R5 k=4: no-overlap floor 497k tokens/s, measured 297k (1.17x of one GPU)

On the floor at k=2, below it at k=4 (the staged fabric sags). And because DP is three annotations, stacking it on the fast config is a loop — efficiency recovers exactly as the grown compute term predicts.

The Lore, and the Ladder Beyond

modded-nanoGPT’s speedrun = this chapter’s contents stacked to the ceiling: compiled kernels, FlashAttention, a better optimizer (Muon), fp8. Each record a new rung.

measure → classify → fix → re-measure. Two budgets, three regimes, one loop — now shown on a real model, seven sections deep.