JAX Release 26.04
The NVIDIA JAX Release 26.04 is made up of two container images available on NGC: JAX and MaxText.
Contents of the JAX container
This container image contains the complete source for the following software:
- JAX:
/opt/jax - XLA:
/opt/xla - Flax:
/opt/flax - TransformerEngine:
/opt/transformer-engine
The MaxText container image is based on the JAX container. Additionally, it includes:
- MaxText:
/opt/maxtext
The JAX runtime package jaxlib is prebuilt and installed in the default Python environment (/usr/local/lib/python3.10/dist-packages/jaxlib) in the container image.
Versions of packages included in both of these containers:
- CUDA 13.2.1
- Please refer to the CUDA DL 26.04 release notes section for the list of libraries inherited from the CUDA container.
Driver Requirements
Release 26.04 is based on CUDA 13.2.1. For comprehensive and up-to-date driver compatibility information, please refer to the following documentation:
- NVIDIA CUDA Compatibility Guide - Compatibility information between CUDA versions and driver releases.
- CUDA Toolkit Release Notes - Driver version requirements and compatibility matrices.
- NVIDIA Drivers Download - Latest NVIDIA drivers.
Key Features and Enhancements
- Integrated NCCL 2.29+ one-side APIs: ncclPutSignal, ncclSignal, and ncclWaitSignal in XLA to enable custom kernels and ragged all-2-all communication to use NCCL’s symmetric memory.
- Completed Thor support in JAX/XLA with full Triton support - Enabled TCGen5 MMA in triton with CUDA 13.1.
Performance optimizations
- Enabled FP8 reduction ops in XLA-NCCL, helping improve communication throughput by ~1.6x-2x over BF16.
- Improved collective cost model of GB200/GB300 systems to improve the accuracy of latency estimators.
- In PGLE-disabled scenario, improved heuristic to expand the overlap coverage, delivering a single digit perf for dense models.
JAX Toolbox
The JAX Toolbox projects focus on achieving the best performance and convergence on NVIDIA Ampere, Hopper, and Blackwell architecture families and provide the latest deep learning models and scripts for training and fine-tuning. These examples are tested against a nightly CI as well as each NGC container release to ensure consistent accuracy and performance over time.
Nightly Containers
In addition to projects, JAX Toolbox includes nightly containers for libraries across the JAX ecosystem.
| Container | Type | Image URI |
|---|---|---|
| jax | - | ghcr.io/nvidia/jax:jax-YYYY-MM-DD |
| maxtext | LLM framework | ghcr.io/nvidia/jax:maxtext-YYYY-MM-DD |
| axlearn | LLM framework | ghcr.io/nvidia/jax:axlearn-YYYY-MM-DD |
Known Issues
- Transformer Engine version 2.14 shipped with 26.04 has an issue resulting in a nondeterministic wrong answer when using MXFP8 training and bias is present. Please upgrade to Transformer Engine version 2.14.1, where the issue was fixed.