Jcm Benchmark

Benchmarking jcm

Builds on jcm-run (config groups, Hydra traps) and the machine skill for your site — devbox-jcm-runs or derecho-jcm-runs. This skill is about getting a throughput number that is actually true, which is harder than it looks.

Run it

PY=/home/dwatsonparris/micromamba/envs/jcm/bin/python

# Short: 1 month (30 d), ~6 chunks. The default for any A/B.
$PY tools/benchmark.py --preset t63-echam-rrtmgp --months 1 --gpu 1 \
    --label baseline

# Long: 12 months. Use for provisioning and drift, not for A/B.
$PY tools/benchmark.py --preset t63-echam-jam --months 12 --gpu 3 \
    --chunk-days 30 --label jam-year

# A/B a library version without touching the shared editable install
$PY tools/benchmark.py --preset t63-echam-rrtmgp --months 1 --gpu 1 \
    --label rrtmgp-perf \
    --pythonpath /data/dwatsonparris/jax-rrtmgp-worktrees/perf

Presets are in tools/benchmark.py:PRESETS. Each carries the validated stable override set for its grid, not just physics=/grid= — see "Configuration" below. Results land in /scr/dwatsonparris/benchmarks/<label>/ as report.md, result.json, run.log and gpu.csv.

Runs take tens of minutes; launch with run_in_background and watch with a Monitor on run.log, or the tool call will time out.

"Which term is it?" — tools/profile_terms.py

benchmark.py answers how fast; it cannot say where the time went, because a step is one fused XLA module. For that, same presets, different tool:

$PY tools/profile_terms.py --preset ma-t63-l47 --gpu 3

It prints ms/step and ms/call for the dynamical core, each bridge direction and every physics term, by joining a profiler trace to the jax.named_scope labels that jcm.profiling puts in the HLO. Three rules when quoting it:

• It disables CUDA graph capture so that kernels stay individually attributable, so its total step time is high by design (264 vs 236 ms at T63L47). Never quote it as throughput — that is benchmark.py's number.
• Read the "mixed kernels" percentage first. That is device time in fusions spanning two components, charged whole to one of them; when it is large the per-term split is soft and only the coarse dynamics/physics/bridge division is safe to quote.
• Do not lengthen the window to "get a better average". The profiler's event buffer holds ~1e6 events and a T63L47 JAM step emits ~19,000 kernels, so ~20 steps is the ceiling there. A one-day window recorded 20 of its 120 steps and reported a 4× undercount. The tool now fails instead, but the instinct is the trap. It is also unnecessary: kernel shapes are static and there are no data-dependent branches, so a step costs the same regardless of state — the window only needs to span whole radiation sub-cycles, which the default does.

Methodology — the part that matters

Never quote jcm's N sim days/hr log line. It is cumulative including compile time, so it understates throughput by 2-5x early and drifts upward for the whole run. A 5.3x regression was once filed against jax-rrtmgp on the strength of chunk 1 of a run that settled 22x faster; it had to be retracted. tools/benchmark.py ignores that line and parses Wall: Xs this chunk. The reduction lives in tools/chunk_timing.py, shared with the Derecho skill's settled_rate.py so a run cannot get two different answers depending on which tool read it.

Discard chunk 1 — it contains compilation. For T63L47 ECHAM+JAM, compile alone is several minutes.

Require convergence. Chunk times settle over 3-4 chunks (XLA autotuning, cache warming, allocator). A rate is quoted only once the last two chunks agree within 3% — judged on the last two rather than the first agreeing pair, because a noisy series can match by accident early while still drifting, and what matters is that the run had settled by the time it ended. Otherwise the result is marked NOT CONVERGED; if you see that flag, run more chunks rather than reporting the number. A 30-day run at --chunk-days 5 gives 6 chunks, which is comfortably enough; 10 days gives 2 and will correctly refuse.

GPU utilisation is not a convergence signal. XLA's autotuner keeps the device at 95%+ while it is still choosing kernels. "The GPU is pegged" is not evidence a chunk time is steady-state. Only chunk-to-chunk agreement is.

Run on a genuinely free GPU. Always. No exceptions.

This is the hardest rule here. Simulations should always run on free cards; benchmarks especially, because a contended card does not fail loudly — it returns a plausible-looking number that is simply wrong, and you will not be able to tell from the report. Every other precaution in this document is wasted if the card was shared.

On a shared, unscheduled box, verify the target card is genuinely idle — both no compute apps and near-zero resident memory, since either signal alone looks idle for a parked job:

python tools/gpu_util.py          # every GPU and its tenants
python tools/gpu_util.py --free   # free indices; exit 1 if none

tools/benchmark.py calls the same check as a hard pre-flight gate and refuses to start on a busy card. On a scheduled machine (Derecho) this is the scheduler's job instead — see derecho-jcm-runs.

Never stack two runs on one card. It invalidates both timings and can OOM the other tenant.

Do not start other GPU work anywhere on the box mid-A/B. This is the non-obvious one. Even on a different, genuinely free card, a concurrent job competes for host CPU, PCIe bandwidth and the allocator. If it overlaps only the candidate arm and not the baseline arm, it contaminates exactly one side of the comparison — an asymmetry that shows up as a fake speedup or regression and is invisible in the report. Halving your wall time is not worth a result you then cannot trust. Let the whole interleaved sequence finish on one card before starting anything else.

The corollary: parallelising an A/B across cards is not a shortcut. If you must (because a card is going away), run both arms on card A and both on card B, and report the pairs separately — never baseline on A against candidate on B.

Compare like with like on chunk size. Every chunk boundary costs a host sync, a health check and a netCDF write. The short benchmark uses 5-day chunks (to get ~6 chunks and so detect convergence) while a 12-month run uses 30-day chunks, so the short number carries slightly more per-day write overhead than production. That is a couple of percent, and it cancels entirely in an A/B at the same --chunk-days — but do not quote a 5-day-chunk number as the production throughput.

The run must be a real one. This is a production configuration made shorter and more instrumented, not a reduced-physics proxy: real orography, real SSTs, the packaged CMIP6 ozone, JW init and the production sponge. Check the log line forcing.ozone_file=auto resolved to .../t63/ozone.nc. On a hybrid grid an unresolvable auto now raises rather than degrading, so this cannot pass silently; on a sigma grid it still warns about the ANALYTIC profile, and there the radiation is seeing ~7.6x the tropospheric ozone column and the benchmark is measuring the wrong workload.

A NaN'd run is not a benchmark. The tool reports nan_any and exits non-zero. Timing from a run that blew up mid-flight is meaningless — fix the configuration and re-run rather than quoting the chunks before the blow-up.

Configuration

The presets are not physics=X grid=Y — they carry the whole known-stable override set, because a T63L47 run from an isothermal cold start with no sponge goes NaN within days, which silently destroys the benchmark. Each T63 preset pins init=jw init.rh=0.0, real terrain and forcing from file, and run=longrun — which carries ECHAM's upper sponge. Ozone comes from ozone_file: auto, the shipped default, which resolves the packaged CMIP6 climatology; the preset deliberately does not override it.

Adding a preset: put the complete validated override set in PRESETS, and verify it completes the short benchmark NaN-free before using it for comparisons.

--save-interval is clamped to --chunk-days by the tool: a chunk with zero output times dies in to_xarray() with an opaque IndexError: index 0 is out of bounds for axis 0 with size 0.