TPU7x（Ironwood）を使用してモデルをトレーニングする

このドキュメントでは、TPU7x リソースをプロビジョニングする方法について説明します。また、MaxText と XPK を使用してトレーニングワークロードをデプロイする例を示します。

TPU7x は、Ironwood ファミリーの最初のリリースであり、 Google Cloudの第 7 世代 TPU です。Ironwood 世代は、大規模な AI トレーニングと推論向けに設計されています。詳細については、TPU7x をご覧ください。

TPU7x 向けに最適化されたその他の例については、GitHub の Ironwood TPU のトレーニングレシピをご覧ください。

TPU をプロビジョニングする

TPU7x は、次の方法でプロビジョニングして管理できます。

GKE: GKE を使用して、TPU をコンテナ化された ML ワークロード用のアクセラレータのプールとしてプロビジョニングして管理できます。Google Cloud CLI を使用して、GKE クラスタインスタンスを手動で作成し、既存の GKE 本番環境を正確にカスタマイズまたは拡張します。詳細については、GKE での TPU についてをご覧ください。
GKE と XPK: XPK は、GKE でのクラスタの作成とワークロードの実行を簡素化するコマンドラインツールです。これは、ML の実務担当者が Kubernetes の深い専門知識を必要とせずに TPU をプロビジョニングしてトレーニングジョブを実行できるように設計されています。XPK を使用すると、GKE クラスタをすばやく作成し、概念実証とテスト用のワークロードを実行できます。詳細については、XPK GitHub リポジトリをご覧ください。
GKE と TPU Cluster Director: TPU Cluster Director は、予約容量すべてにアクセスできる（ホールドバックのない）All Capacity モードの予約で使用できます。これにより、TPU ハードウェアのトポロジ、使用率、健全性ステータスを完全に把握できます。詳細については、All Capacity モードの概要をご覧ください。

MaxText と XPK を使用してトレーニングワークロードをデプロイする

Accelerated Processing Kit（XPK）を使用して、コンセプト実証とテスト用の GKE クラスタを作成します。XPK は、ML ワークロードのプロビジョニング、管理、実行を簡素化するように設計されたコマンドラインツールです。

以降のセクションでは、MaxText と XPK を使用してトレーニングワークロードをデプロイする方法について説明します。

始める前に

始める前に、次の手順を完了します。

課金が有効になっている Google Cloud プロジェクトがあることを確認します。
TPU7x へのアクセス権を取得します。詳細については、アカウントチームにお問い合わせください。
XPK で使用するアカウントに、XPK GitHub リポジトリに記載されているロールがあることを確認します。

XPK と依存関係をインストールする

XPK をインストールします。XPK GitHub リポジトリの手順に沿って操作します。
管理者から提供された手順で Docker をインストールするか、公式のインストール手順でインストールします。インストールが完了したら、次のコマンドを実行して Docker を構成し、インストールをテストします。
```
gcloud auth configure-docker
sudo usermod -aG docker $USER # relaunch the terminal and activate venv after running this command
docker run hello-world # Test Docker
```
次の環境変数を設定します。
```
export PROJECT_ID=YOUR_PROJECT_ID
export ZONE=YOUR_ZONE
export CLUSTER_NAME=YOUR_CLUSTER_NAME
export ACCELERATOR_TYPE=YOUR_ACCELERATOR_TYPE
export RESERVATION_NAME=YOUR_RESERVATION_NAME
export BASE_OUTPUT_DIR="gs://YOUR_BUCKET_NAME"
```
次のように置き換えます。
- YOUR_PROJECT_ID: 実際の Google Cloud プロジェクト ID。
- YOUR_ZONE: クラスタを作成するゾーン。プレビュー版では、us-central1-c のみがサポートされています。
- YOUR_CLUSTER_NAME: 新しいクラスタの名前。
- YOUR_ACCELERATOR_TYPE: TPU のバージョンとトポロジ。例: tpu7x-4x4x8サポートされているトポロジの一覧については、サポートされている構成をご覧ください。
- YOUR_RESERVATION_NAME: 予約の名前。共有予約の場合は、projects/YOUR_PROJECT_NUMBER/reservations/YOUR_RESERVATION_NAME を使用します。
- YOUR_BUCKET_NAME: Cloud Storage バケットの名前。これはモデルトレーニングの出力ディレクトリになります。

既存の Cloud Storage バケットがない場合は、次のコマンドを使用して作成します。

gcloud storage buckets create ${BASE_OUTPUT_DIR} \
    --project=${PROJECT_ID} \
    --location=US \
    --default-storage-class=STANDARD \
    --uniform-bucket-level-access

単一 NIC、単一スライスクラスタを作成する

MTU を構成するの手順に沿って、ネットワーク構成を最適化します。
${CLUSTER_ARGUMENTS} 変数に値を入力します。この変数は xpk cluster create コマンドで使用します。
```
export CLUSTER_ARGUMENTS="--network=${NETWORK_NAME} --subnetwork=${SUBNET_NAME}"
```
xpk cluster create コマンドを使用して、TPU7x ノードプールを含む GKE クラスタを作成します。
```
xpk cluster create \
    --project=${PROJECT_ID} \
    --zone=${ZONE} \
    --cluster ${CLUSTER_NAME} \
    --cluster-cpu-machine-type=n1-standard-8 \
    --tpu-type=${ACCELERATOR_TYPE} \
    --reservation=${RESERVATION_NAME} \
    --custom-cluster-arguments="${CLUSTER_ARGUMENTS}"
```
--cluster-cpu-machine-type フラグを n1-standard-8（またはそれ以上）に設定すると、デフォルトのノードプールにシステム Pod（JobSet Webhook など）に十分な CPU が確保され、エラーを防ぐことができます。デフォルトでは、XPK は e2-standard-16 を使用します。一部のゾーンは特定の CPU タイプのみをサポートしているため、n1、n2、e2 のタイプを切り替える必要がある場合があります。そうしないと、割り当てエラーが発生する可能性があります。
クラスタのアップグレードを防ぐために、メンテナンスの除外を追加します。
```
gcloud container clusters update ${CLUSTER_NAME} \
    --zone=${ZONE} \
    --add-maintenance-exclusion-name="no-upgrade-next-month" \
    --add-maintenance-exclusion-start="EXCLUSION_START_TIME" \
    --add-maintenance-exclusion-end="EXCLUSION_END_TIME" \
    --add-maintenance-exclusion-scope="no_upgrades"
```
次のように置き換えます。
- EXCLUSION_START_TIME: メンテナンスの除外に選択した開始時間（YYYY-MM-DDTHH:MM:SSZ 形式）。
- EXCLUSION_END_TIME: メンテナンスの除外に選択した終了時間（YYYY-MM-DDTHH:MM:SSZ 形式）。

MaxText Docker イメージをビルドまたはアップロードする

MaxText が提供するスクリプトを使用して Docker イメージをローカルでビルドするか、ビルド済みのイメージを使用できます。

ローカルにビルドする

次のコマンドは、ローカルディレクトリをコンテナにコピーします。

# Make sure you're running on a virtual environment with python3.12. If nothing is printed, you have the correct version.
[[ "$(python3 -c 'import sys; print(f"{sys.version_info.major}.{sys.version_info.minor}")' 2>/dev/null)" == "3.12" ]] || { >&2 echo "Error: Python version must be 3.12."; false; }

# Clone MaxText
git clone https://github.com/AI-Hypercomputer/maxtext.git
cd maxtext
git checkout maxtext-tutorial-v1.0.0

# Custom Jax and LibTPU wheels
pip download libtpu==0.0.28.dev20251104+nightly -f "https://storage.googleapis.com/jax-releases/libtpu_releases.html"
pip download --pre jax==0.8.1.dev20251104 jaxlib==0.8.1.dev20251104 --index https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/

# Build the Docker image
bash docker_build_dependency_image.sh MODE=custom_wheels

コマンドが正常に実行されると、maxtext_base_image という名前のイメージがローカルに作成されます。ローカルイメージは、xpk ワークロードコマンドで直接使用できます。

画像をアップロードする（省略可）

前のセクションの手順に沿って Docker イメージをローカルでビルドしたら、次のコマンドを使用して MaxText Docker イメージをレジストリにアップロードできます。

export CLOUD_IMAGE_NAME="${USER}-maxtext-runner"
bash docker_upload_runner.sh CLOUD_IMAGE_NAME=${CLOUD_IMAGE_NAME}

このコマンドが正常に実行されると、gcr.io に gcr.io/PROJECT_ID/CLOUD_IMAGE_NAME という名前の MaxText イメージが表示されます。

MaxText トレーニングコマンドを定義する

Docker コンテナ内でトレーニングスクリプトを実行するコマンドを準備します。

MaxText 1B モデルは、約 10 億個のパラメータで言語モデルをトレーニングするために設計された MaxText フレームワーク内の構成です。このモデルは、小規模なチップスケールのテストに使用します。パフォーマンスが最適化されていません。

export MAXTEXT_COMMAND="JAX_PLATFORMS=tpu,cpu \
    ENABLE_PJRT_COMPATIBILITY=true \
    python3 src/MaxText/train.py src/MaxText/configs/base.yml \
        base_output_directory=${BASE_OUTPUT_DIR} \
        dataset_type=synthetic \
        per_device_batch_size=2 \
        enable_checkpointing=false \
        gcs_metrics=true \
        run_name=maxtext_xpk \
        steps=30"

トレーニングワークロードをデプロイする

xpk workload create コマンドを実行して、トレーニングジョブをデプロイします。MaxText ベースイメージを使用するために --base-docker-image フラグを指定するか、--docker-image フラグと使用するイメージを指定する必要があります。--enable-debug-logs フラグを含めてデバッグロギングを有効にすることもできます。

xpk workload create \
    --cluster ${CLUSTER_NAME} \
    --base-docker-image maxtext_base_image \
    --workload maxtext-1b-$(date +%H%M) \
    --tpu-type=${ACCELERATOR_TYPE} \
    --zone ${ZONE} \
    --project ${PROJECT_ID} \
    --command "${MAXTEXT_COMMAND}"
    # [--enable-debug-logs]

ワークロード名はクラスタ内で一意である必要があります。この例では、一意性を確保するために、ワークロード名に $(date +%H%M) が追加されています。