Model performance (probjax.nn)¶

End-to-end forward and train-step cost of probjax.nn models and of the scan_layers / remat_policy options, measured on the RTX 3090 (gpu_nn.json, kernel_nn_gpu.json) and on CPU (kernel_nn_cpu.json). Kernel-level numbers are on the attention and kernel pages, setup on the overview. A train step is value_and_grad of a scalar loss plus an optimizer update. Times are medians in milliseconds. OOM: out of device memory. fail†: XLA found no kernel configuration that fits during autotuning (NOT_FOUND), so the case did not run. unsup.: the implementation rejects the shape or backend. error: another failure (the JSON row has the message). —: not measured.

Transformer and DiffusionTransformer¶

The standard profile (model width 256, 8 heads of 32, 4 layers), batch 8, with each attention_fn. The default dot_product_attention dispatches to flex_attention on NVIDIA GPUs (rules).

Batch 8, width 256, 4 layers, RTX 3090. cuDNN has no fp32 path.

At 4096 tokens in bf16 the default train step takes 50.1 ms against 87.5 ms with cuDNN attention; jax.nn XLA attention is 225 at that size. At this head size (32) the probjax kernel is faster than cuDNN in both the forward and the train step.

Scan over layers and rematerialization¶

Transformer and DiffusionTransformer take remat_policy=None | 'full' | 'nothing_saveable' | <jax.checkpoint_policies name> | callable and scan_layers=False. scan_layers=True stacks the layer parameters and runs the blocks under jax.lax.scan; it changes the parameter tree, so checkpoints of an unscanned model do not load into a scanned one. Both are opt-in.

24 layers, width 512, 4096 tokens, batch 2, bf16, RTX 3090; device peak in a fresh process. The last row uses einsum attention (implementation='xla'), which materializes the scores.

remat_policy='full' cuts the device peak of the 24-layer bf16 model from 5,212 to 2,072 MB for 260 instead of 212 ms per step (the forward runs twice), and is what lets the einsum-attention model train at all at this length. scan_layers=True compiles in 1.07 s instead of 7.32 s, but without remat it keeps every layer's activations as scan residuals and runs out of memory. For deep stacks use scan_layers=True, remat_policy='full'; dots_saveable sits between the two.

On CPU (24 layers, 512 tokens), remat_policy='full' is also faster, not only smaller: 1,505 vs 1,755 ms per step, because the smaller working set stays in cache.

SSMModel cells¶

SSMModel(16, 128, 16, 4) with each recurrent cell, batch 8, fp32. Mamba2Cell is the full Mamba-2 block on ssd (expand=2, short causal convolution, gated RMSNorm, head_dim=64, state_dim=64); SSDCell is a multi-head SSD layer with the Mamba-2 decay parameterization; MambaCell uses state_dim = model_dim by default; GatedDeltaNetCell and LRUCell use their defaults.

SSMModel(16, 128, 16, 4), batch 8, fp32, RTX 3090.

For the kernels underneath see SSD, Mamba and Gated DeltaNet.

UNet¶

UNet(3, (64, 128, 256)), batch 8, RTX 3090; device peak in a fresh process:

bf16 train step at resolution 64 with remat or donated buffers. Remat costs time without reducing temporaries at this size:

MoE layer¶

Dropless MoELayer, 8 experts, top-2, d_model 512, d_hidden 1024, one column per dropless_impl. More code paths, memory and the largest token counts are on the kernel page:

Training-loop overhead¶

Per-step time of an MLP trained with a bare jitted step and with fit() on an in-memory array and on a stream. fit() recompiles on every call: a second call with the same model and data costs nearly as much wall time as the first (last two columns).

How to reproduce¶

# Model grid (Transformer, DiffusionTransformer, SSM cells, UNet, MoE).
python -m docs.scripts.benchmark_gpu_nn --output docs/benchmarks/gpu_nn.json

# scan_layers / remat_policy, UNet variants, training loop; --peak adds device peaks.
python -m docs.scripts.benchmark_kernel_nn --device gpu --output docs/benchmarks/kernel_nn_gpu.json
JAX_PLATFORMS=cpu python -m docs.scripts.benchmark_kernel_nn --device cpu \
  --output docs/benchmarks/kernel_nn_cpu.json

# Rebuild the figures and this page from the JSON.
python -m docs.scripts.benchmark_figures