Repository files navigation

JAX Toolbox provides a public CI, Docker images for popular JAX libraries, and optimized JAX examples to simplify and enhance your JAX development experience on NVIDIA GPUs. It supports JAX libraries such asMaxText andPallas.

We support and test the following JAX frameworks and model architectures. More details about each model and available containers can be found in their respective READMEs.

Components	Container	Build	Test
	ghcr.io/nvidia/jax:base		[no tests]
	ghcr.io/nvidia/jax:jax
	ghcr.io/nvidia/jax:equinox		[tests disabled]
	ghcr.io/nvidia/jax:upstream-t5x		[tests disabled]
	ghcr.io/nvidia/jax:t5x
	ghcr.io/nvidia/jax:maxtext
	ghcr.io/nvidia/jax:axlearn

In all cases,ghcr.io/nvidia/jax:XXX points to latest nightly build of the container forXXX. For a stable reference, useghcr.io/nvidia/jax:XXX-YYYY-MM-DD.

In addition to the public CI, we also run internal CI tests on H100 SXM 80GB and A100 SXM 80GB.

TheJAX image is embedded with the following flags and environment variables for performance tuning of XLA and NCCL:

There are various other XLA flags users can set to improve performance. For a detailed explanation of these flags, please refer to theGPU performance doc. XLA flags can also be tuned per workload. For example, each script includes a directoryxla_flags.

For a list of previously used XLA flags that are no longer needed, please also refer to theGPU performance page.

Seethis page for more information about how to profile JAX programs on GPU.

docker run -it --shm-size=1g ...

slurmstepd: error: pyxis:     [INFO] Authentication succeededslurmstepd: error: pyxis:     [INFO] Fetching image manifest listslurmstepd: error: pyxis:     [INFO] Fetching image manifestslurmstepd: error: pyxis:     [ERROR] URL https://ghcr.io/v2/nvidia/jax/manifests/<TAG> returned error code: 404 Not Found

• AWS
  • Add EFA integration
  • SageMaker code sample
• GCP
  • Getting started with JAX multi-node applications with NVIDIA GPUs on Google Kubernetes Engine
• Azure
  • Accelerating AI applications using the JAX framework on Azure’s NDm A100 v4 Virtual Machines
• OCI
  • Running a deep learning workload with JAX on multinode multi-GPU clusters on OCI

• JAX | NVIDIA NGC Container
• Slurm and OpenMPI zero config integration
• Adding custom GPU ops
• Triaging regressions

• Equinox for JAX: The Foundation of an Ecosystem for Science and Machine Learning
• Scaling Grok with JAX and H100
• JAX Supercharged on GPUs: High Performance LLMs with JAX and OpenXLA
• What's New in JAX | GTC Spring 2024
• What's New in JAX | GTC Spring 2023