google/paxml

这是什么？

核心功能

• 01基于 dataclass 与 Fiddle 的可配置实验和超参数体系。
• 02支持 pjit/SPMD、pmap 以及多切片 TPU 并行方式。
• 03提供 1B、16B、GPT-3 XL 等规模的 C4 训练示例配置。
• 04支持 SeqIO、Lingvo 适配器和自定义等多种数据输入方式。
• 05提供 TPU v4 Pod 上的 Model FLOPs Utilization 基准数据。

适用场景

• 在 Google Cloud TPU v4 Pod 上训练大规模 decoder-only 语言模型。
• 用可配置的模型和并行设置开展可复现的 JAX 实验。
• 在不同模型和芯片规模下评测训练效率（MFU）。

快速上手

1. 1

   创建 Cloud TPU 虚拟机（先设置环境变量）

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

   通过 SSH 连接到 TPU 虚拟机

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

   从 PyPI 安装 Pax 及 TPU 版 JAX

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

   运行示例测试模型

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

使用前须知

• !主要面向 Cloud TPU 硬件；README 中的安装和示例假定使用 TPU 虚拟机。
• !部署依赖 Google Cloud 资源，例如 GCP 项目和存储桶。

标签

常见问题

paxml 是什么？

基于 Jax 的机器学习框架，用于配置和运行大规模模型训练实验，主要面向 TPU 上的语言模型。 Pax（Paxml）是构建在 Jax 之上的机器学习实验配置与运行框架，强调实验流程和并行策略都可以完全配置。它主要用于训练大规模模型，例如 decoder-only Transformer 语言模型。模型、任务和优化器通过基于 dataclass 的超参数定义，底层依赖配套的 Praxis 层库，并支持 Fiddle 配置。

paxml 可以免费使用吗？用的什么许可证？

paxml 以 Apache-2.0 许可证开源。用于商业产品前，请先确认许可证条款。

paxml 还在维护吗？

活跃维护中。paxml 最近一次提交是在 2026年10月26日。

paxml 适合用来做什么？

在 Google Cloud TPU v4 Pod 上训练大规模 decoder-only 语言模型。 用可配置的模型和并行设置开展可复现的 JAX 实验。 在不同模型和芯片规模下评测训练效率（MFU）。

paxml 怎么安装和运行？

创建 Cloud TPU 虚拟机（先设置环境变量）: gcloud compute tpus tpu-vm create $TPU_NAME --zone=$ZONE --version=$VERSION --project=$PROJECT --accelerator-type=$ACCELERATOR 通过 SSH 连接到 TPU 虚拟机: gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE 从 PyPI 安装 Pax 及 TPU 版 JAX: python3 -m pip install paxml jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html 运行示例测试模型: python3 .local/lib/python3.8/site-packages/paxml/main.py --exp=tasks.lm.params.lm_cloud.LmCloudSpmd2BLimitSteps --job_log_dir=gs://<your-bucket>

在哪里获取 paxml？

源代码和安装说明在 GitHub 上：https://github.com/google/paxml