Hardware validation¶
The hosted CI runs on CPUs. What the accelerator stack does on a real GPU is
checked on rented machines by scripts/cloud/runpod_matrix.py, on demand,
never on a push. The questions it answers:
- whether JAX sees the card or quietly answers
CpuDevice(id=0) - whether
ppy.cudalaunches - whether
ppy.xlacompiles for the GPU - whether a sharded trainer spans two physical cards
python scripts/cloud/runpod_matrix.py --cuda # one NVIDIA GPU
python scripts/cloud/runpod_matrix.py --multigpu # two or more NVIDIA GPUs in one Pod
python scripts/cloud/runpod_matrix.py --rocm # one AMD Instinct GPU
python scripts/cloud/runpod_matrix.py --all
How a run works¶
The script needs an authenticated runpodctl (runpodctl doctor) and a
public key in ~/.ssh. It never reads or prints the API key.
For each environment it:
- asks what is in stock, and prefers a common card (an RTX 4090, an A6000, an A40, an L40S) over an expensive one
- makes a Pod named
ppy-test-<env>-<id>from a vendor image, with a four-hour termination set on it - waits for SSH
- runs
scripts/cloud/remote_test.shthere against the exact commit named (--sha, the checkout's HEAD by default). The Pod clones it, so it must be pushed.
Every step of the remote script is logged and its status recorded; nothing
is skipped silently. The results come back under
cloud-results/<stamp>/<env>/, which git ignores. The Pod is deleted in a
finally and confirmed gone. Only Pods the run created are touched.
What a run checks¶
- The environment as installed. The lock's JAX is synced, then the
accelerator plugin for that same version is installed (
jax[cuda12], or AMD's ROCm plugin).scripts/cloud/accelerator_check.pyrecords the interpreter, JAX, jaxlib, every plugin package, the driver and toolkit the vendor tools report,jax.default_backend(), and every device. It then fails, with the report written, if every device is a CPU. An import that succeeded proves nothing, and a CPU does not pass a GPU test. - A computation on the device. A jitted matrix multiplication and a
reduction, synchronized with
block_until_ready(), against a NumPy reference. With two or more devices, an array is sharded across them and reduced across the shards. - PPy's own paths.
examples/39_xlaunderppy run, with the bridge compiling for the GPU.examples/38_cudaandexamples/44_tileunderppy run, withppy explainon the kernels so a launch that fell back to the reference is seen as one.- The GPU, tile, and XLA test files.
- The lowering-limit tests, so a function that stays in Python is the one the docs say stays.
- Multi-GPU.
examples/45_multi_gpu_jaxunderppy runandpython: the batch sharded over a mesh of the cards, the gradient summed across them, the loss falling, the result held to a single-device run.scripts/cloud/multiprocess_smoke.py: one JAX process per card with a coordinator on the loopback address, assertingjax.process_count(),jax.process_index(), the device lists, and one collective across the processes.
- A few benchmarks, for a sanity check of accelerator performance: the JAX and PyTorch comparisons, and the CUDA and tile examples against CuPy and Triton. These are validation numbers from a rented machine. The tables in the READMEs come from the dedicated bench runner and are never overwritten by them.
Last run¶
The fixes-0.3 branch, 2026-09-11, all three environments at commit
e280fe5, from cloud-results/ of those runs. The numbers are those
machines' and are not the READMEs' tables.
| Environment | Hardware | Count | JAX backend | PPy backend | JAX compute | Training |
|---|---|---|---|---|---|---|
| Local | Intel Core Ultra 9 386H (CPU) | 1 | cpu | LLVM | PASS (two virtual CPU devices, tests/test_multi_device_jax.py) |
PASS on virtual devices, not a hardware pass |
| CUDA single | NVIDIA GeForce RTX 4090 (RunPod) | 1 | gpu (cuda:0) |
CUDA: kernels compiled for the device, ppy.xla on cuda:0 |
PASS | N/A |
| CUDA multi | 2 × NVIDIA GeForce RTX 4090 (RunPod) | 2 | gpu (cuda:0, cuda:1) |
CUDA | PASS: sharded across both, reduced across both | PASS under ppy run and python; one process per GPU PASS |
| ROCm single | AMD Instinct MI300X (RunPod, EU-RO-1) | 1 | gpu (rocm:0) |
ppy.xla on rocm:0; HIP is ppy emit hip only, the source compiled by the image's hipcc for gfx942, never launched by PPy |
PASS | N/A |
| ROCm multi | -- | -- | -- | -- | NOT RUN | NOT RUN |
CUDA single¶
Stack. Python 3.13.8; jax 0.11.1, jaxlib 0.11.1, jax-cuda12-plugin
0.11.1, jax-cuda12-pjrt 0.11.1, nvidia-nccl-cu12 2.31.2 (CUDA 12.9 pip
libraries); driver 580.159.04, CUDA 13.0 image, no nvcc.
jax.default_backend() = gpu, jax.devices() = [cuda:0].
JAX compute. A jitted 1024×512×256 float32 matmul at full precision against NumPy: relative error 5.9e-7, reduction exact.
PPy.
examples/38_cudaandexamples/44_tileunderppy runprintkernels compiled for a device here: True, andppy explainshows the host function boxed for itsGpuLauncheffect, as the guide says it is.examples/39_xlaprints# devices: ['cuda:0'], the bridge having compiled for the GPU.- 32 tests in
test_gpu_frontend,test_ir_gpu,test_tile,test_xlapass, as do the lowering-limit tests.
Benchmarks on that machine.
| benchmark | PPy cuda.launch |
CuPy |
|---|---|---|
| saxpy | 0.40 ms | 0.44 ms |
| block max | 1.03 | 1.02 |
| saxpy with copies | 42 | 59 |
The JAX comparison ran JAX on the GPU against PPy's fused loop on the Pod's
CPU, which measures the machines and says nothing about the compilers. The
Triton counterpart's wheel found no driver in that image
(0 active drivers), so its row is its own failure and the PPy tile row
stands alone.
CUDA multi¶
Stack. The same stack (Python 3.13.8, jax 0.11.1 with the CUDA 12
plugin, NCCL 2.31.2) on driver 580.95.05 with a CUDA 13.0 image.
jax.devices() = [cuda:0, cuda:1], jax.local_devices() the same, both
NVIDIA GeForce RTX 4090.
JAX compute. The check sharded an array over both cards and reduced across them (exact).
examples/45_multi_gpu_jax. The batch was sharded over the two cards,
and a hundred steps were dispatched without a host synchronization between
them. The train time is the steps and does not include a wait per step.
ppy run |
python |
|
|---|---|---|
train time |
573 ms | 570 ms |
| largest parameter difference | 7e-8 (the standardization native, 1.1 ms) | 5e-7 (71 ms) |
| result | PASS |
PASS |
Loss went from 1.0414 to 1.0108 on two devices, and the same on one.
scripts/cloud/multiprocess_smoke.py. Two processes, one per card,
jax.process_count() 2, each seeing one local and two global devices, the
collective sum exact in both.
PPy. 32 GPU, tile, and XLA tests pass. PPy cuda.launch saxpy 0.39 ms
against CuPy 0.44, saxpy with copies 34 against 52.
Hosts that failed¶
Two other two-GPU hosts were tried in an earlier run and are the reason the harness records what it does.
- On a two-A6000 host, NCCL's peer-to-peer transport hung (both cards at
100%, the collective never returning) and completed with
NCCL_P2P_DISABLE=1. - On a two-A40 host every check passed, but the sharded trainer's loss came out wrong and different on each run (0.52, then 0.64, against 1.04 on one device). A collective answered wrongly instead of hanging.
The same program on the two-4090 host, and on two virtual CPU devices,
agrees with the single-device run to 7e-8. The faults are in those hosts'
NCCL, and not in PPy's or JAX's arithmetic: sharding_probe.py compares
every stage with NumPy and found nothing on the host that works. A run is
therefore judged on its own transcript, never on a step's exit status alone,
and the trainer exits non-zero on FAIL.
ROCm single¶
Stack. An MI300X (gfx942) in EU-RO-1, refused earlier that day and
granted in the evening, in AMD's rocm/jax:rocm10.0-jax0.11.0-py3.12 image:
Python 3.12.3, jax 0.11.0, jaxlib 0.11.0, jax-rocm10-plugin and
jax-rocm10-pjrt 0.11.0+rocm10.0.0, ROCm 10.0.0 as the pip SDK, driver
6.10.5. jax.default_backend() = gpu, jax.devices() = [rocm:0].
JAX compute. The matmul and the reduction report a relative error of 0.0 against NumPy.
Install. PPy was installed beside that stack with the image's jax and
jaxlib as constraints (without the jax group's flax, which wants a newer
jax than the image has). The run checks that every jax* distribution is the
one that was there before the install; it is.
PPy.
examples/39_xlaprints# devices: ['rocm:0'], the bridge having compiled for the AMD card through the same PJRT path as on CUDA.examples/38_cudaandexamples/44_tilerun their CPU path there (kernels compiled for a device here: False), sinceppy.cudaandppy.tilelaunch CUDA only.- 31 GPU, tile, and XLA tests pass with one skipped, as do the lowering-limit tests.
ppy emit hiponexamples/38_cuda/saxpy.ppywrites 79 lines of HIP C++ with the__global__kernel and the<<<>>>launch. The image'shipcccompiles it for gfx942 (HIP_DEVICE_LIB_PATHpointed at the SDK's device library).
That is the extent of ROCm in PPy: ppy.hip is source only, and this run
does not make it a launch runtime.
ROCm multi¶
Not run; one MI300X was in stock.
What the runs found¶
Bugs in PPy¶
The PJRT bridge behind ppy.xla had two bugs, both fixed:
- It compiled for every device of the platform. On the two-GPU hosts its
executables expected one argument shard per device and refused the single
buffers the bridge places
(
Expected args to execute_sharded_on_local_devices to have 2 shards, got: [1, 1]). It compiles for one device now, and a test over two virtual CPU devices holds it there. - It defaulted to the
cpuplatform whatever the machine had. It follows JAX's default backend now, and refuses to call an initialization failure or an installed-but-idle accelerator plugin "cpu". That is howexamples/39_xlacame to printcuda:0androcm:0.
Changes to the harness¶
- A float32 matmul on an Ampere-class card is TF32 unless full precision is asked for, so the check asks.
- A cross-device step gets fifteen minutes and a recorded retry without peer-to-peer.
- The one-process-per-GPU step counts its transcripts, since a loop that ran nothing had once exited zero.
- A Pod that is still listed after the run fails the run.
- AMD's image runs no
sshd, so the harness starts one for the injected key (RunPod's own pattern for a plain image) and hands the container's environment to the SSH session.