PROJECT: Characterise JAX on cloud CI runners — GPU fit, concurrency, and build guidance

[mmcky]

Why now

Two independent JAX-on-cloud failures surfaced in the space of a week, in different repos, with different root causes, at different layers of the stack. Each was invisible until it burned hours of on-demand GPU runner time, and neither would have been predicted by "it works on a laptop" or "it works on Colab".

incident	repo	symptom	root cause	status
NUTS cell timeout	QuantEcon/lecture-python.myst#1022	jax + numpyro cell >2400s on the T4, 157s on a laptop CPU; GPU ≥2.7x slower than the CPU in the same box	unconfirmed — leading hypotheses are crippled FP64 on inference-class GPUs and kernel-launch overhead on 10x10 ops	worked around by pinning the lecture to CPU
Concurrent kernel death	QuantEcon/lecture-python-programming#363	jax_intro kernel aborts with no traceback on every cold-cache build	confirmed by an instrumented run — three JAX kernels share one T4, the first preallocates JAX's default 75% (11.3GiB of 15.3GiB), the loser clamps to the ~3.4GiB remainder and aborts when a cell outgrows it	fixed with XLA_PYTHON_CLIENT_PREALLOCATE=false, validation run in progress

The two incidents share a pattern worth naming: our lecture code is written against a mental model of "a GPU makes JAX fast", but the runners are single inference-class GPUs (T4 on g4dn.2xlarge) shared by concurrent notebook execution, running workloads that are frequently small-matrix, float64, and latency-bound rather than throughput-bound. We do not currently have a tested model of when that combination wins, loses, or crashes — we find out one incident at a time.

Proposed project

A focused investigation to characterise JAX on our cloud CI runners and turn the findings into standing guidance and CI configuration, rather than per-incident firefighting.

Questions to answer

1. FP64 penalty. How large is the float64 penalty on the T4 for our actual workload shapes (Kalman filters, NUTS, small linear models)? T4 FP64 is nominally 1/32 of FP32 throughput — measure what that does to representative lecture cells, and whether jax_enable_x64 lectures should default to CPU.
2. Small-op / launch overhead. At what array size does the T4 break even with the runner's own CPU for our common op mix? A lax.scan of 10x10 ops may never amortise launch latency — find the crossover so lecture authors have a rule of thumb rather than a guess.
3. Concurrency policy. myst executes notebooks at parallelism = cpus-1, so multiple JAX kernels share one GPU. With preallocation off, what is the safe concurrency for JAX lectures, and should JAX-heavy lectures be serialised while pure-CPU lectures fan out?
4. Sampler configuration. Does chain_method='vectorized' close the GPU gap for numpyro workloads, or is sequential + CPU the right default for small models?
5. Standard CI environment. Which XLA/JAX env vars should be standard across all GPU lecture repos (XLA_PYTHON_CLIENT_PREALLOCATE, possibly XLA_PYTHON_CLIENT_MEM_FRACTION, allocator choice), and where do they belong so every repo inherits them rather than rediscovering them?
6. Instance fit. Is g4dn.2xlarge the right instance at all? Candidates to price against: g5/g6 (better FP64 story is not automatic — check), plain CPU instances for FP64-heavy lectures, or a split fleet where only genuinely large-array lectures get a GPU.
7. Failure hygiene. A pathological cell currently burns 40 minutes per builder because a failed notebook is never cached and every builder re-executes it. What per-cell timeout and caching policy makes a bad cell fail fast and fail once?

Method

A small benchmark harness that runs a matrix of representative kernels (Kalman filter scan, NUTS on the sargent_surico model, large dense linear algebra as a control, the jax_intro 50M-element ops) across {T4 GPU, runner CPU} x {float32, float64} x {1, 2, 3 concurrent processes}, with a jax profile captured for the pathological cases to separate kernel time from gap time. Run it on the actual CI runners via workflow_dispatch so the numbers are the numbers we live with. The two incident lectures become the first regression fixtures.

Deliverables

1. A short findings report with the measured crossover points and the answer to each question above, checked into QuantEcon/meta.
2. Guidance for lecture authors (a page in the style guide or manual): when to expect the GPU to help, when to pin a lecture to CPU, and the sampler/config idioms that make small models CI-safe.
3. Standard CI configuration rolled out across the five GPU-CI lecture repos — env vars, concurrency settings, timeouts — ideally via the shared infrastructure in quantecon/actions so it lands once (related: Migrate the five native lecture CI pipelines to quantecon/actions/build-lectures #348, GitHub Actions config hardening across the five GPU-CI lecture repos #347).
4. An audit of existing JAX lectures against the findings: which are in the losing region (small model + x64 on GPU) and should be pinned or restructured.

Out of scope

Rewriting lecture content for performance, GPU vendor comparisons beyond what instance selection requires, and anything about the mystmd/jupyter-book build engines (tracked separately on QuantEcon/lecture-python-programming#363).

Evidence trail

• jax on the GPU runner is >2.7x slower than its CPU for a small float64 model (sargent_surico NUTS cell) lecture-python.myst#1022 — the NUTS/T4 measurements, hypotheses and proposed experiments.
• [Experimental] Build the lectures with QuantEcon/mystmd instead of jupyter-book lecture-python-programming#363 — the preallocation-race diagnosis: instrumented run 30806359596 in that repo shows the GPU timeline (0 → 11,305MiB → 14,727 → 3,427 → 0) and a clean dmesg ruling out host OOM; commits 8af1015 (instrumentation, engine pin) and 94a7e52 (fix).

[mmcky]

Adding some analysis ahead of building the harness: priors for the seven investigation questions (so benchmark runs confirm/refute expectations rather than explore blindly), a few proposed additions to the methodology, and two fixes worth making now. Based on a review of the two incidents plus the CI configs across the lecture repos.

Reading of the two incidents

They look like one underlying mismatch in two guises:

• #1022 (NUTS >2400 s on T4, 157 s on laptop CPU). Over-determined. T4 FP64 peak is ~0.25 TFLOPS (1/32 of FP32) — the same order as the g4dn.2xlarge's own 4-core Cascade Lake host under MKL, so for float64 matmul the T4 may never beat its own host CPU, let alone a modern laptop. On top of that, NUTS on a small model is sequential leapfrog: thousands of ~5–20 µs kernel launches that a GPU can't amortise. Neither factor is fixable with a bigger g-family GPU (see Q1 below).
• #363 (three JAX kernels, silent death). 0.75 × 15.36 GiB = 11.5 GiB — matches the reported 11.3 GiB, i.e. JAX's default 75% per-process preallocation. The second/third kernel then OOMs during CUDA context init, before Python can raise — hence no traceback. Worth confirming whether the three kernels were designedconcurrency or zombie kernels left alive after an earlier cell timeout; the mitigation differs.

Priors for the investigation questions

Q	Expected answer
FP64 penalty	~2–4× for elementwise (bandwidth-bound: 2× bytes over 320 GB/s), up to ~30× for matmul (FLOP-bound). Note g5/g6 don't fix FP64: A10G/L4 are 1/64-ratio ≈ 0.47–0.49 TFLOPS, barely 2× the T4. Fix for small FP64 models is CPU pinning, not hardware.
Crossover	FP32 matmul: GPU wins from roughly 512×512. FP64 matmul: possibly never on T4. Elementwise: ~1e6+ elements to amortise launch overhead.
Concurrency	Prefer a designed budget — XLA_PYTHON_CLIENT_MEM_FRACTION=.30 for ≤3 kernels — over PREALLOCATE=false, which trades OOM-at-init for fragmentation and per-alloc overhead. Plus guaranteed kernel teardown on timeout.
Sampler defaults	chain_method="vectorized" should help GPU on small models (one launch across chains) but likely still loses to CPU for tiny models; guidance default: numpyro.set_platform("cpu") / JAX_PLATFORMS=cpu for small NUTS.
Standard env	Baseline candidates: XLA_PYTHON_CLIENT_MEM_FRACTION (or PREALLOCATE=false), JAX_COMPILATION_CACHE_DIR (see below), JAX_TRACEBACK_FILTERING=off in CI so failures produce real tracebacks.
Instance selection	Prior: g6.2xlarge (L4, SM 8.9, 24 GB) for build jobs — T4 sits at the exact floor of the jax cuda13 support matrix (SM ≥ 7.5, driver ≥ 580) and is first to fall off when it next moves. Keep g4dn for the Colab-emulation jobs (Colab free tier = same T4). Repos with token JAX usage may belong on CPU instances entirely.
Failure behaviour	execute.timeout in _config.yml is already per-cell (nbclient semantics). lecture-python.myst's 2400 s is the outlier — programming/intro use 600 s. Lower the default and raise per-lecture in frontmatter. The 40-min-per-builder repeat is capped by the ci.yml gating fix (first jb build gates); confirm all GPU repos have it.

Proposed additions to the methodology

• [ ] Add a compile-vs-runtime split to the benchmark matrix. JIT compile is CPU-bound and dominates many short lectures. Follow-through: persist JAX_COMPILATION_CACHE_DIR alongside the jupyter cache in CI — speeds every build on every instance family, likely the single biggest cheap win here.
• [ ] Make instance family a harness axis: a workflow_dispatch input interpolated into the RunsOn label (family=g4dn.2xlarge|g5.2xlarge|g6.2xlarge|c7i.2xlarge) runs the identical harness per family for ~$1 each, and answers Q6 with data.
• Report medians over ≥5 reps; separate cold (first-call, includes compile) from warm.
• [ ] Tag "GPU-narrative" lectures in the audit — lectures whose published output depends on a GPU being present (e.g. numpy_vs_numba_vs_jax, jax_intro device printouts). These constrain retiering decisions independently of performance.
• Frame the audit as Colab-parity too: Colab free tier is the same T4, so lectures in the losing region are slow for readers, not just CI.
• Deliver the standard env + runner label via a composite action in quantecon/actions, so instance/env policy is a one-line change across all GPU repos.

Two fixes worth making now (independent of the project)

• [ ] lecture-python-programming publish.yml omits spot= — RunsOn defaults to spot=price-capacity-optimized, so the publish job currently runs on interruptible spot capacity while throwaway previews run on-demand. Add spot=false(or deliberately flip previews to spot instead).
• [ ] lecture-python-programming is a CPU-tier candidate: only ~3 of 27 lectures execute JAX. A c7i.2xlarge (~$0.36/hr) is half the g4dn's cost and likely builds faster — subject to the GPU-narrative tagging above.

Adding some analysis ahead of building the harness: priors for the seven investigation questions (so benchmark runs confirm/refute expectations rather than explore blindly), a few proposed additions to the methodology, and two fixes worth making now. Based on a review of the two incidents plus the CI configs across the lecture repos.

Reading of the two incidents

They look like one underlying mismatch in two guises:

• #1022 (NUTS >2400 s on T4, 157 s on laptop CPU). Over-determined. T4 FP64 peak is ~0.25 TFLOPS (1/32 of FP32) — the same order as the g4dn.2xlarge's own 4-core Cascade Lake host under MKL, so for float64 matmul the T4 may never beat its own host CPU, let alone a modern laptop. On top of that, NUTS on a small model is sequential leapfrog: thousands of ~5–20 µs kernel launches that a GPU can't amortise. Neither factor is fixable with a bigger g-family GPU (see Q1 below).
• #363 (three JAX kernels, silent death). 0.75 × 15.36 GiB = 11.5 GiB — matches the reported 11.3 GiB, i.e. JAX's default 75% per-process preallocation. The second/third kernel then OOMs during CUDA context init, before Python can raise — hence no traceback. Worth confirming whether the three kernels were designed concurrency or zombie kernels left alive after an earlier cell timeout; the mitigation differs.

Priors for the investigation questions

Q	Expected answer
1. FP64 penalty	~2–4× for elementwise (bandwidth-bound: 2× bytes over 320 GB/s), up to ~30× for matmul (FLOP-bound). Note g5/g6 don't fix FP64: A10G/L4 are 1/64-ratio ≈ 0.47–0.49 TFLOPS, barely 2× the T4. Fix for small FP64 models is CPU pinning, not hardware.
2. Crossover	FP32 matmul: GPU wins from roughly 512×512. FP64 matmul: possibly never on T4. Elementwise: ~1e6+ elements to amortise launch overhead.
3. Concurrency	Prefer a designed budget — XLA_PYTHON_CLIENT_MEM_FRACTION=.30 for ≤3 kernels — over PREALLOCATE=false, which trades OOM-at-init for fragmentation and per-alloc overhead. Plus guaranteed kernel teardown on timeout.
4. Sampler defaults	chain_method="vectorized" should help GPU on small models (one launch across chains) but likely still loses to CPU for tiny models; guidance default: numpyro.set_platform("cpu") / JAX_PLATFORMS=cpu for small NUTS.
5. Standard env	Baseline candidates: XLA_PYTHON_CLIENT_MEM_FRACTION (or PREALLOCATE=false), JAX_COMPILATION_CACHE_DIR (see below), JAX_TRACEBACK_FILTERING=off in CI so failures produce real tracebacks.
6. Instance selection	Prior: g6.2xlarge (L4, SM 8.9, 24 GB) for build jobs — T4 sits at the exact floor of the [jax cuda13 support matrix](https://github.com/jax-ml/jax/blob/main/docs/installation.md) (SM ≥ 7.5, driver ≥ 580) and is first to fall off when it next moves. Keep g4dn for the Colab-emulation jobs (Colab free tier = same T4). Repos with token JAX usage may belong on CPU instances entirely.
7. Failure behaviour	execute.timeout in _config.yml is already per-cell (nbclient semantics). lecture-python.myst's 2400 s is the outlier — programming/intro use 600 s. Lower the default and raise per-lecture in frontmatter. The 40-min-per-builder repeat is capped by the ci.yml gating fix (first jb build gates); confirm all GPU repos have it.

Proposed additions to the methodology

• Add a compile-vs-runtime split to the benchmark matrix. JIT compile is CPU-bound and dominates many short lectures. Follow-through: persist JAX_COMPILATION_CACHE_DIR alongside the jupyter cache in CI — speeds every build on every instance family, likely the single biggest cheap win here.
• Make instance family a harness axis: a workflow_dispatch input interpolated into the RunsOn label (family=g4dn.2xlarge|g5.2xlarge|g6.2xlarge|c7i.2xlarge) runs the identical harness per family for ~$1 each, and answers Q6 with data.
• Report medians over ≥5 reps; separate cold (first-call, includes compile) from warm.
• Tag "GPU-narrative" lectures in the audit — lectures whose published output depends on a GPU being present (e.g. numpy_vs_numba_vs_jax, jax_intro device printouts). These constrain retiering decisions independently of performance.
• Frame the audit as Colab-parity too: Colab free tier is the same T4, so lectures in the losing region are slow for readers, not just CI.
• Deliver the standard env + runner label via a composite action in quantecon/actions, so instance/env policy is a one-line change across all GPU repos.

Two fixes worth making now (independent of the project)

• lecture-python-programming publish.yml omits spot= — RunsOn defaults to spot=price-capacity-optimized, so the publish job currently runs on interruptible spot capacity while throwaway previews run on-demand. Add spot=false (or deliberately flip previews to spot instead).
• lecture-python-programming is a CPU-tier candidate: only 3 of 27 lectures execute JAX. A c7i.2xlarge ($0.36/hr) is half the g4dn's cost and likely builds faster — subject to the GPU-narrative tagging above.

[mmcky]

Two things from project-compute — one offer, one dependency.

The offer. QuantEcon owns quantecon-gpu2: two TITAN RTX cards, Turing sm_75 — the same architecture generation as the g4dn's T4. It is powered down and shipping, but once recommissioned it could reproduce the incidents here on hardware we control, without consuming CI minutes or waiting on workflow_dispatch. A second machine, quantecon-gpu1, is slated for a used 24 GB Ampere card, which would let you test directly whether the FP64 penalty in question 1 is generation-specific or persists past Turing. That bears on question 6: "a newer generation fixes it" and "GPUs are the wrong shape for these workloads" imply very different instance answers.

Neither is available today, and the method here is deliberately to run on the actual CI runners so the numbers are the numbers we live with — so this is offered as a supplement, not a substitute.

The dependency, stated so it is visible rather than discovered later. project-compute is planning a self-hosted GPU runner pilot, and its evidential base is 538 GPU-hours a year measured across the five GPU-CI repos. If this programme pins lectures to CPU, that figure shrinks — and the Turing FP64 penalty applies identically to gpu2's Titans, so a finding here changes what insourcing is worth. We are watching rather than blocking, but it is the largest external dependency that plan has, and it is now recorded as such on our side.