NVIDIA/JAX-Toolbox

JAX-Toolbox

What it solves

JAX Toolbox provides a centralized resource for developers to run, optimize, and scale JAX-based machine learning workloads on NVIDIA GPUs. It addresses the complexity of configuring high-performance environments by offering pre-built container images, optimized libraries, and technical guidance for training and inference.

How it works

The project serves as a hub for NVIDIA's JAX-on-GPU ecosystem. It provides nightly and stable container images (via NGC and GitHub Container Registry) that include optimized versions of JAX, XLA, and related libraries like TransformerEngine. These containers are pre-configured with the necessary CUDA and NCCL drivers to ensure high performance on NVIDIA hardware. The repository also hosts documentation, technical blogs, and performance-tuning guides for scaling large models.

Who it’s for

This project is for machine learning engineers, researchers, and infrastructure developers who are building or training large-scale models using JAX and need to ensure peak performance on NVIDIA GPU clusters.

Highlights

  • Provides production-ready container images for JAX and MaxText, including nightly builds and stable monthly releases.
  • Offers deep-dive technical documentation on profiling, optimizing, and scaling JAX workloads.
  • Includes performance-focused guides for training large language models (LLMs) and long-context models.
  • Features resources for advanced GPU features like FP4/FP8 precision, custom kernel development with CuTe, and distributed training strategies.

Related

  • Project
    AI-Hypercomputer/maxtext

    A high-performance, scalable LLM library and reference implementation written in JAX for training and post-training popular open-source models on TPUs and GPUs.

  • Project
    NVIDIA/TransformerEngine

    NVIDIA Transformer Engine is an open‑source library that speeds up Transformer model training on NVIDIA GPUs by providing fused kernels and automatic mixed‑precision APIs (FP8, MXFP8, NVFP4, BF16, FP16) for PyTorch and JAX. It handles scaling‑factor management, quantization metadata, and distributed communication, letting users adopt low‑precision training without changing their code. Integrated with major LLM frameworks (Hugging Face, Lightning, Megatron‑LM, etc.), it targets large‑scale language, vision, and MoE models on Hopper/Blackwell/Rubin hardware.

  • Project
    google/tunix

    A JAX-based post-training library for LLMs that provides scalable support for supervised fine-tuning and reinforcement learning on TPUs.

  • Project
    openxla/xla

    XLA is an open-source machine learning compiler that optimizes models from PyTorch, TensorFlow, and JAX for high-performance execution on GPUs, CPUs, and ML accelerators.

  • Project
    NVIDIA/apex

    A collection of NVIDIA-maintained utilities for PyTorch that streamline mixed precision and distributed training through optimized CUDA extensions.