jax on the GPU runner is >2.7x slower than its CPU for a small float64 model (sargent_surico NUTS cell)

[mmcky]

Summary

A jax + numpyro cell in lectures/sargent_surico.md could not finish in 2400 seconds on our g4dn.2xlarge (T4) CI runner, and takes 157 seconds on a laptop CPU. Pinning jax to the CPU with jax.config.update('jax_platform_name', 'cpu') made the lecture build, cutting it from a hard timeout to 14m52s on the same runner.

So the GPU is not merely a poor fit for this workload — it is at least 2.7x slower than the CPU sitting next to it in the same box, and that is a floor, since we only ever observed "more than 2400s", never the actual GPU time.

@HumphreyYang — flagging you as the person most likely to know where to look. The pin is already merged into the PR branch as a pragmatic fix, so nothing is blocked; the question is what the underlying cause is and how many other lectures it touches.

The measurements

All from PR #1020, lecture sargent_surico.md, notebook execution step (Build Download Notebooks):

configuration	full step	that lecture alone	the NUTS cell
jax on GPU (default)	58m02s	~47m30s	timed out at 2400s
jax pinned to CPU	26m10s	14m52s	comfortably inside the cap
local laptop, CPU only	—	5m09s (52 cells)	157s

The timeout was confirmed from the execution-reports-html artifact of run 30674947077, which names the cell: CellTimeoutError: A cell timed out while it was being executed, after 2400 seconds with a preview beginning kernel = NUTS(ss_model, target_accept_prob=0.8, dense_mass=True,.

Runner CPU is roughly 5x slower than the laptop on this lecture's non-jax cells, so the comparison is CPU-to-CPU consistent — the GPU is the outlier, not the machine.

The workload

The cell is a numpyro NUTS run over a small linear state-space model:

• 10x10 matrices, a lax.scan Kalman filter over about 100 quarterly observations
• jax.config.update('jax_enable_x64', True) — the whole model runs in float64
• NUTS(..., dense_mass=True, max_tree_depth=8), so up to 256 leapfrog steps per iteration
• MCMC(..., num_warmup=400, num_samples=400, num_chains=4, chain_method='sequential')

Hypotheses, none confirmed

1. float64 on a T4. Inference-class GPUs have deliberately crippled FP64 — on a T4 it is a small fraction of FP32 throughput. This is my leading guess but I have not measured it.
2. Kernel launch overhead. At 10x10, each op is far too small to fill the device, so the run may be dominated by launch latency rather than arithmetic. A lax.scan of tiny ops is close to the worst case.
3. chain_method='sequential'. Four chains run one after another. 'vectorized' would batch them, and it is worth knowing whether that alone closes the gap on GPU.
4. Compilation. Possible but unlikely to account for tens of minutes on its own.

These are not mutually exclusive, and 1 and 2 together would be enough to explain the size of the gap.

What would settle it

A profiled run on an actual T4 — something we cannot do from a laptop. Useful experiments, roughly in order of cost:

• time the cell on GPU with jax_enable_x64 off, to size the FP64 contribution
• time it on GPU with chain_method='vectorized'
• capture a jax profile and look at kernel time versus gap time, which separates hypothesis 1 from 2
• time a single loglik_jax gradient evaluation on CPU and GPU, which isolates the model from the sampler

Why this matters beyond one lecture

If it is FP64 or small-op overhead, the rule generalises: any lecture pairing jax_enable_x64 with a small model may be slower on our GPU runners than on their CPUs. We would want to know which of the existing jax lectures are in that category, and whether the guidance should be to pin small-model lectures to CPU by default and reserve the GPU for lectures with genuinely large arrays.

There is also a CI-hygiene angle. This failure cost about two hours of on-demand g4dn.2xlarge time across two runs, because a failed notebook is not cached and every builder re-executes it — the same cell timed out twice in one job, once in the notebook build and again in the HTML build. Worth considering whether the per-cell timeout: 2400 in lectures/_config.yml should be lower, so a pathological cell fails fast instead of burning 40 minutes per builder.