Repolode
Submit
google

google/paxml

Pax is a Jax-based framework for configuring and running large-scale machine learning experiments, mainly language models, on TPUs.

View on GitHub
Stars
560
Forks
73
Open issues
30
Last commit
11 days ago

What is it?

Pax (Paxml) is a framework for configuring and running machine learning experiments on top of Jax. It is built around fully configurable experiments and parallelization, and is aimed at training large models such as decoder-only Transformer language models. Models, tasks and learners are defined with dataclass-based hyperparameters, which are built on the companion Praxis layer library and support Fiddle configuration.

The project targets Google Cloud TPU VMs, including single-host v4 setups, TPU Pod slices, and multislice configurations. It ships example experiment configs for training on the C4 dataset at 1B, 16B and larger parameter counts, and includes a benchmark of Model FLOPs Utilization (MFU) across weak-scaling runs on TPU v4 Pods. Data inputs can come from SeqIO, Lingvo, or custom BaseInput subclasses, with guidance on multihost sharding and evaluation.

The repository notes that NVIDIA maintains a separate updated Pax version with GPU support, referenced from its own repository.

Key features

  • 01Configurable experiments and hyperparameters built on dataclasses and Fiddle.
  • 02Parallelism support for pjit/SPMD, pmap, and multislice TPU setups.
  • 03Example C4 training configs for 1B, 16B and GPT-3 XL style models.
  • 04Multiple data input options, including SeqIO, Lingvo adaptors, and custom inputs.
  • 05Benchmarks of Model FLOPs Utilization on TPU v4 Pods.

Use cases

  • Training large decoder-only language models on Google Cloud TPU v4 pods.
  • Running reproducible JAX-based experiments with configurable model and parallelism settings.
  • Benchmarking training efficiency (MFU) at different model and chip scales.

Quick start

  1. 1

    Create a Cloud TPU VM (set the environment variables first)

    gcloud compute tpus tpu-vm create $TPU_NAME --zone=$ZONE --version=$VERSION --project=$PROJECT --accelerator-type=$ACCELERATOR
  2. 2

    SSH into the TPU VM

    gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE
  3. 3

    Install Pax from PyPI together with JAX for TPU

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

    Run an example test model

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

Things to know

  • Designed for Cloud TPU hardware; the README's setup and examples assume TPU VMs.
  • Setup depends on Google Cloud resources such as a GCP project and a storage bucket.

Topics

#jax#large-language-models#tpu#distributed-training#parallelism#machine-learning#c4#gpt#llm#model-flops

FAQ

What is paxml?
Pax is a Jax-based framework for configuring and running large-scale machine learning experiments, mainly language models, on TPUs. Pax (Paxml) is a framework for configuring and running machine learning experiments on top of Jax. It is built around fully configurable experiments and parallelization, and is aimed at training large models such as decoder-only Transformer language models. Models, tasks and learners are defined with dataclass-based hyperparameters, which are built on the companion Praxis layer library and support Fiddle configuration.
Is paxml free to use? What license does it have?
paxml is open source under the Apache-2.0 license. Check the license terms before using it in commercial products.
Is paxml still maintained?
Actively maintained. The last commit to paxml was on Oct 26, 2026.
What can I use paxml for?
Training large decoder-only language models on Google Cloud TPU v4 pods. Running reproducible JAX-based experiments with configurable model and parallelism settings. Benchmarking training efficiency (MFU) at different model and chip scales.
How do I install paxml?
Create a Cloud TPU VM (set the environment variables first): gcloud compute tpus tpu-vm create $TPU_NAME --zone=$ZONE --version=$VERSION --project=$PROJECT --accelerator-type=$ACCELERATOR SSH into the TPU VM: gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE Install Pax from PyPI together with JAX for TPU: python3 -m pip install paxml jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html Run an example test model: python3 .local/lib/python3.8/site-packages/paxml/main.py --exp=tasks.lm.params.lm_cloud.LmCloudSpmd2BLimitSteps --job_log_dir=gs://<your-bucket>
Where can I get paxml?
The source code and installation instructions are on GitHub: https://github.com/google/paxml

Summary generated by AI from the project's README. GitHub stats refresh daily.

Related projects