ANNOUNCEMENT: General Availability of JAX NGC containers

We are excited to announce the GA release of the latest versions of JAX, Paxml, and T5x containers on NGC. NVIDIA NGC Catalog containers include a validated set of libraries optimized for GPU performance. Before each release, the containers are scanned for security vulnerabilities and performance across a range of popular DL models and workloads. These containers may also include experimental features in order to maximize performance and compatibility with NVIDIA GPUs.

Highlights from 23.08 release:

  • Pre-train GPT models with Paxml, an LLM framework built on JAX and Praxis
  • Pre-train T5 models using T5x
  • State-of-the-art vision transformers (ViT) integrated with DALI for faster image preprocessing

Containers are available for free to pull from NGC Catalog. You can track our latest progress in JAX Toolbox where we also host nightly containers.

To learn more: