Pax: a Jax framework for configuring and running large model experiments

What Pax is, and what it sits on top of

Pax is a framework to configure and run machine learning experiments on top of Jax, and the repository description adds that it allows for advanced and fully configurable experimentation and parallelization, and that it has demonstrated industry leading model flop utilization rates. That last claim is the one to treat with care, because the repository does not include the measurements behind it, so what you can verify from the documentation is the shape rather than the number. The useful way to think about Pax is as a configuration layer over Jax rather than a new autodiff engine. The training logic is Jax, and what Pax adds is a declarative way to describe a model, how it is partitioned across accelerators, and how it is trained, so that two experiments differ by a config object rather than by a code change.

The quickstart starts with a TPU VM, not with pip

The quickstart is written for Cloud TPU, and the first step is creating a TPU VM with 8 cores from a corp machine, which the documentation says is sufficient for the more exhaustive material.

export ZONE=us-central2-b
export VERSION=tpu-vm-v4-base
export PROJECT=<your-project>
export ACCELERATOR=v4-8
export TPU_NAME=paxml

#create a TPU VM
gcloud compute tpus tpu-vm create $TPU_NAME \
--zone=$ZONE --version=$VERSION \
--project=$PROJECT \
--accelerator-type=$ACCELERATOR

The variables are exported first so the create command reads them. If you are using TPU Pod slices there is a separate guide, and the commands are run from a local machine with the worker=all option, which is what reaches every slice at once.

gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE \
--worker=all --command="<commmands>"

The quickstart sections then assume a single-host TPU, so you ssh in and run the rest there.

Installing Pax, and the requirements file that saves you

After ssh-ing into the VM you install the paxml stable release from PyPI, or the development version from GitHub. The PyPI route upgrades pip and installs paxml together with the TPU extra of jax, pointing at the libtpu release index.

python3 -m pip install -U pip
python3 -m pip install paxml jax[tpu] \
-f https://storage.googleapis.com/jax-releases/libtpu_releases.html

There is a specific escape hatch for dependency trouble, and the documentation is unusually helpful here. If you hit issues with transitive dependencies in the native Cloud TPU VM environment, navigate to the release branch named after the version you are installing and download the requirements.txt from the paxml/pip_package directory of that branch, which contains the exact versions of all transitive dependencies used in the native Cloud TPU VM environment where the release is built and tested. Install it with no deps, so the pinned list is what you get.

git clone -b rX.Y.Z https://github.com/google/paxml
pip install --no-deps -r paxml/paxml/pip_package/requirements.txt

For the development version, the instructions are to install the dev version of praxis first and install both repositories in editable mode, so that editing the code takes effect without reinstalling.

Running an example model, and the two parallelism paths

The examples are launched through a main.py with an experiment name and a job log directory, and the two forms in the documentation differ in which parallelism mechanism they use. One uses pjit, which is the SPMD path.

# example model using pjit (SPMD)
python3 .local/lib/python3.8/site-packages/paxml/main.py \
--exp=tasks.lm.params.lm_cloud.LmCloudSpmd2BLimitSteps \
--job_log_dir=gs://<your-bucket>

The other uses pmap, and carries an extra flag for using tensorstore.

# example model using pmap
python3 .local/lib/python3.8/site-packages/paxml/main.py \
--exp=tasks.lm.params.lm_cloud.LmCloudTransformerAdamLimitSteps \
--job_log_dir=gs://<your-bucket> \
--pmap_use_tensorstore=True

The experiment names point at classes in the tasks.lm.params packages, which is where the config objects live. The job log directory being a gs:// bucket is not decoration either, since the logs, checkpoints and metrics of a pod-scale run are the output you actually want to keep.

The reference convergence runs, and what they tell you

The documentation includes sample convergence runs on the c4 dataset, and these are the most useful thing in the repository for anyone deciding whether to use it, because they are the configs you would otherwise have to invent. A one billion parameter model on c4 on a TPU v4-8 uses the config C4Spmd1BAdam4Replicas from the c4.py file of the language model task parameters.

python3 .local/lib/python3.8/site-packages/paxml/main.py \
--exp=tasks.lm.params.c4.C4Spmd1BAdam4Replicas \
--job_log_dir=gs://<your-bucket>

A sixteen billion parameter model on c4 on a TPU v4-64 uses C4Spmd16BAdam32Replicas, and the naming convention itself is informative, since the number after Adam is the replica count, which is the axis you scale along.

python3 .local/lib/python3.8/site-packages/paxml/main.py \
--exp=tasks.lm.params.c4.C4Spmd16BAdam32Replicas \
--job_log_dir=gs://<your-bucket>

The GPT3-XL model on c4 on a TPU v4-128 uses the pipeline parallel config C4SpmdPipelineGpt3SmallAdam64Replicas.

python3 .local/lib/python3.8/site-packages/paxml/main.py \
--exp=tasks.lm.params.c4.C4SpmdPipelineGpt3SmallAdam64Replicas \
--job_log_dir=gs://<your-bucket>

Each run comes with a loss curve and a log perplexity graph in the documentation, so you can see what a healthy run looks like before you spend TPU hours finding out.

Running the notebooks on a TPU VM

The tutorials in the docs folder are Jupyter notebooks, and there is a step by step procedure for running them on a TPU VM, which starts with port forwarding when you ssh in.

`gcloud compute tpus tpu-vm ssh $TPU_NAME --project=$PROJECT_NAME --zone=$ZONE --ssh-flag="-4 -L 8080:localhost:8080"`

Then you install Jupyter Notebook on the VM and pin markupsafe to a specific version.

Then you install Jupyter Notebook on the VM with `pip install notebook` and pin markupsafe to a specific version with `pip install markupsafe==2.0.1`.

After that you export the Jupyter path, copy the example notebooks to the VM with the scp variant of the gcloud tpu-vm command, start the notebook server on port 8080 without a browser, and open the local port in your own browser and paste the token the server prints. One practical note is at the end of the list, and it is the kind that costs an afternoon if you do not know it. If you need to start a second notebook while the first is still holding the TPUs, you can kill the Python processes to free the accelerators.

GPU support, and where Pax actually stands now

The GPU story is a fork in the road and the project says so directly. NVIDIA has released an updated version of Pax with H100 FP8 support and broad GPU performance improvements, and the documentation points at the NVIDIA Rosetta repository in the JAX-Toolbox project for details and usage instructions. In other words, on GPU you are pointed at a different distribution of the same framework rather than at the one in this repository. The version history in this repository is the other thing to weigh. The releases are paxml 1.4.0 from April 2024, 1.3.1 from February 2024 and 1.3.0 from February 2024, while the last push to the repository was on 2026-10-15, so the tree is being worked on well after the last tagged release. The FAQ gives the same advice in a different form, pointing you at the requirements.txt in the branch matching the stable release you are installing, and using the stable 0.4.0 release and its r0.4.0 branch as the worked example.

Editorial conclusion

Pax fits a team training a large language model on Cloud TPU and wants the parallelism strategy, the data pipeline and the optimiser to be described in configuration rather than in a training loop, and it fits anyone who wants to start from a known-good reference run rather than writing one. It does not fit a GPU-first team without reading the extra work, because the examples and the quickstart are all TPU oriented, and NVIDIA maintains a separate Pax distribution with H100 FP8 support and performance work in the JAX-Toolbox repository. It also does not fit a project that wants a current dependency set, because the released versions here are 1.4.0 from April 2024 and 1.3.1 and 1.3.0 from February 2024, while the last push to the repository was on 2026-10-15, so the tree and the releases are far apart. Before you build on it, read RELEASE.md, and pin the release branch whose requirements file you intend to use.

Frequently asked questions

Official sources

1. google/paxml on GitHub
2. Issues
3. License: Apache-2.0
4. README
5. Releases