JAX with Paxml

Run JAX and Paxml on AWS EC2 GPU-accelerated instances for distributed language model training

View on GitHub

Overview

This use case provides a sample Dockerfile (jax_paxml.Dockerfile) to run JAX and Paxml on AWS EC2 GPU-accelerated instances. It enables distributed training of large language models using JAX’s SPMD (Single Program, Multiple Data) programming model.

Container Description

The reference Dockerfile performs the following steps:

  • Provides JAX built for NVIDIA CUDA devices, using a recent NVIDIA CUDA image as the parent image.
  • Removes unnecessary networking packages that might conflict with AWS technologies.
  • Installs EFA user-space libraries. The kernel drivers are intentionally not built during docker build, and the self-tests are skipped, as both steps are expected to fail when run during container build.
  • Installs the NCCL recommended version.
  • Installs aws-ofi-nccl so that NCCL can utilize EFA.
  • Installs JAX.
  • Installs Paxml.
  • Installs Praxis.

The Dockerfile pins its component versions (EFA installer, NCCL, aws-ofi-nccl, JAX, Praxis, and Paxml) through build arguments. Rather than reproducing volatile version strings here, treat the jax_paxml.Dockerfile as the source of truth for the exact pinned values. The image tag used in the commands below is paxml:jax-0.4.18-1.2.0, matching the tag used by the upstream README and expected by the Slurm launch script.

Build the Container

Build the JAX container as follows:

# Build a container image
DOCKER_BUILDKIT=1 docker build --progress=plain -f jax_paxml.Dockerfile -t paxml:jax-0.4.18-1.2.0 .

# Verify the image has been built
docker images

Convert the Container to Enroot Format

Convert the container image to the enroot squash format used by the Slurm launch script. The output filename /fsx/paxml_jax-0.4.18-1.2.0.sqsh matches the --container-image path that jax.sbatch imports at run time, so keep it consistent:

# Convert to enroot format. Remove any existing .sqsh first (idempotent), otherwise enroot
# refuses to run when the output .sqsh file already exists.
rm -f /fsx/paxml_jax-0.4.18-1.2.0.sqsh
enroot import -o /fsx/paxml_jax-0.4.18-1.2.0.sqsh dockerd://paxml:jax-0.4.18-1.2.0

Run on Slurm

Once the container is converted to the enroot format, you can run the lm_clouds example of Pax on a Slurm cluster. The following command submits a job to the Slurm cluster to train a 2B-parameter transformer-based SPMD language model on synthetic data (the jax.sbatch script sets the training config to the LmCloudSpmd2B family):

sbatch jax.sbatch