google/paxml
Pax is a Jax-based machine learning framework for training large scale models. Pax allows for advanced and fully configurable experimentation and parallelization, and has demonstrated industry leading model flop utilization rates.
Built on Praxis (a configuration-based layer abstraction), Pax uses JAX's `pjit` (SPMD) and `pmap` primitives to express complex distributed training strategies across TPU Pods and GPU clusters without explicit communication code. The framework emphasizes declarative experiment configuration, enabling researchers to specify model architecture, data pipeline, and parallelization strategy through composable Python configs that can be version-controlled and reproduced.
550 stars. Actively maintained with 5 commits in the last 30 days. Available on PyPI.
Stars
550
Forks
70
Language
Python
License
Apache-2.0
Category
Last pushed
Mar 12, 2026
Commits (30d)
5
Dependencies
20
Get this data via API
curl "https://pt-edge.onrender.com/api/v1/quality/llm-tools/google/paxml"
Open to everyone — 100 requests/day, no key needed. Get a free key for 1,000/day.
Related tools
axolotl-ai-cloud/axolotl
Go ahead and axolotl questions
JosefAlbers/PVM
Phi-3.5 for Mac: Locally-run Vision and Language Models for Apple Silicon
metriccoders/one-line-llm-tuner
This repository is the source code for fine tuning any LLM in just one line 🔥
Nano-Collective/nanotune
A simple, interactive CLI for fine-tuning small language models on Apple Silicon. No YAML...
MoHussein197/dgx-spark-finetune-llm
🔧 Fine-tune large language models efficiently on NVIDIA DGX Spark with LoRA adapters and...