google/paxml

これは何？

主な機能

• 01データクラスとFiddleを基盤とした、柔軟に設定できる実験とハイパーパラメータ。
• 02pjit/SPMD、pmap、TPUのマルチスライス構成に対応した並列化。
• 0310億、160億パラメータのモデルやGPT-3 XL相当のモデル向けC4学習設定例。
• 04SeqIO、Lingvoアダプター、独自の入力処理など、複数のデータ入力方法。
• 05TPU v4 PodにおけるModel FLOPs Utilizationのベンチマーク。

ユースケース

• Google Cloud TPU v4 Podでの大規模なdecoder-only言語モデルの学習。
• モデルや並列化の設定を変更しながら、再現可能なJAXベースの実験を実行。
• モデルやチップの規模を変えて学習効率（MFU）をベンチマーク。

クイックスタート

1. 1

   Cloud TPU VM を作成する（先に環境変数を設定）

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

   TPU VM に SSH 接続する

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

   TPU 向けの JAX とともに Pax を PyPI からインストールする

   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 VM を前提としています。
• !セットアップには、GCP プロジェクトやストレージバケットなどの Google Cloud リソースが必要です。

トピック

よくある質問

paxml とは？

Paxは、主に言語モデルを対象に、TPU上で大規模な機械学習実験を設定・実行するためのJAXベースのフレームワークです。 Pax（Paxml）は、JAX上で機械学習実験を設定・実行するためのフレームワークです。実験や並列化を柔軟に設定できる設計で、decoder-only Transformer言語モデルなどの大規模モデルの学習を想定しています。モデル、タスク、学習器は、データクラスベースのハイパーパラメータで定義します。これらは関連ライブラリの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 VM を作成する（先に環境変数を設定）: gcloud compute tpus tpu-vm create $TPU_NAME --zone=$ZONE --version=$VERSION --project=$PROJECT --accelerator-type=$ACCELERATOR TPU VM に SSH 接続する: gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE TPU 向けの JAX とともに Pax を PyPI からインストールする: 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