step_jit = nnx.jit(train_step) # the configuration 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)')Case Study: Optimizing a Transformer
Cumulative Optimization of a GPT
Apply the chapter’s profile–change–verify procedure to the GPT of ch. 11. Each bar changes one mechanism and inherits the previous settings.
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 sequence will look nothing like it. That contrast is the lesson.
Use Fixed Shapes, Synchronization, and Fresh Batches
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.dataloader reduced throughput by about one quarter → 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; use fresh data for each timing.
Whole-Step JIT Defines the Practical JAX Baseline
Un-jitted JAX is included only as a reference point; practical training uses a compiled step.
first (compiling) step: 7.3 s
R1 jit: 255278 tokens/s (19x)nnx.jit on the whole step (loss + grads + update) is one transformation and roughly 20×. Every subsequent ratio starts from this bar, not from R0.
Classify the Compiled Step
Ask the compiler what it plans to run:
XLA counts 765 GFLOP per step; at 32.1 ms per step that is 24 TFLOP/s achievedAbout a third of the tf32 roof (§13.2): not overhead-starved, not saturated — a mixed workload. Cheapest next configuration on §13.4’s sequence: bf16 — double the roof, half the bytes.
bf16 Requires Consistent Compute Dtypes
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: 233844 tokens/s (0.92x)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.
Interaction between Batch Size and Precision
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 sequence 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): 440731 tokens/s (1.88x), 45 TFLOP/s achieved
fp32 control at 512: 201228 tokens/s -> bf16 pays 2.19x here
planned temp: bf16 8.9 GiB (fp32 control: 15.1 GiB)The two changes interact. 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 configurations cumulatively. Checkpointing then reduces throughput by about one tenth while saving memory that is not limiting this run, as in the PyTorch tab.
Aggregate Results
configurations = ['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(configurations, [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 configurations')
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.73x (and R0 -> R3: 32x)Whole-step JIT is about 20× faster than the un-jitted diagnostic. Bf16 changes little at batch 64; bf16 and batch 512 together produce the next increase. Checkpointing reduces throughput while reducing memory. A 300-step run checks that loss still decreases under the final configuration.
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 201k (0.79x of one GPU)
R5 k=4: no-overlap floor 499k tokens/s, measured 298k (1.17x of one GPU)Throughput is near the floor at k=2 and below it at k=4 as effective fabric bandwidth falls. Because data parallelism requires three annotations, the fast configuration can be rerun in a loop; efficiency then improves as the grown compute term predicts.
Further Optimization Techniques
modded-nanoGPT’s speedrun = this chapter’s contents stacked to the ceiling: compiled kernels, FlashAttention, a better optimizer (Muon), fp8. Each record a new configuration.
measure → classify → fix → re-measure. Two budgets, three regimes, one loop — now shown on a real model, seven sections deep.