Google は AI 技術を使用して、コンテンツをご希望の言語に翻訳しています。AI 翻訳には誤りが含まれる場合があります。

Ray と TorchTPU を使用してマルチホスト TPU でモデルをファインチューニングする

Autopilot Standard

このチュートリアルでは、Ray Core、Ray Train、PyTorch Fully Sharded Data Parallel（FSDP2）、TorchTPU を使用して、Google Kubernetes Engine（GKE）のマルチホスト Cloud TPU スライスで分散 PyTorch ワークロードを実行し、Gemma 4 モデルをファインチューニングする方法について説明します。

このチュートリアルは、Kubernetes のマルチホスト TPU スライスで分散 PyTorch トレーニングワークロードをオーケストレートする機械学習（ML）エンジニア、プラットフォーム管理者、オペレーター、データおよび AI スペシャリストを対象としています。 Google Cloud のコンテンツで使用されている一般的なロールとタスクの例の詳細については、一般的な GKE ユーザーロールとタスクをご覧ください。

このチュートリアルを開始する前に、次の内容をよく理解しておいてください。

背景

このセクションでは、このチュートリアルで使用されている重要なテクノロジーについて説明します。

TPU

Tensor Processing Unit（TPU）は、機械学習ワークロード用に Google がカスタム設計したアクセラレータです。このチュートリアルは、TPU Trillium（v6e）と Ironwood（TPU7x）の両方をサポートしています。マルチホスト TPU スライスでは、複数の TPU ホストが高速チップ間相互接続（ICI）ネットワークを介して直接通信し、単一のチップまたは TPU ホストのメモリ容量を超えるモデルのトレーニングに必要な高帯域幅メモリ（HBM）と低レイテンシの集団通信を提供します。

TorchTPU

TorchTPU は、Google と Meta の共同開発による Cloud TPU 向けのオープンソース PyTorch バックエンドです。TorchTPU を使用すると、最小限のコード変更で既存の PyTorch モデルを TPU ハードウェアで実行し、Accelerated Linear Algebra（XLA）コンパイラを介してトレーニングと推論のワークロードを単一チップからマルチスライス TPU Pod にスケールできます。TorchTPU は、TPU ICI ネットワーク上でデバイスごとに 1 つのプロセスを実行するモデルを使用し、eager execution、torch.compile、torch.distributed、DeviceMesh、DTensor、Fully Sharded Data Parallel（FSDP2）などの標準の PyTorch API をサポートしています。

TPU での KubeRay と Ray

KubeRay は、Ray アプリケーションをデプロイして管理するための Kubernetes オペレーターです。Ray on GKE アドオンを有効にすると、マネージド KubeRay オペレーターと KubeRay TPU Webhook がデプロイされます。KubeRay が google.com/tpu リソースをリクエストするワーカー Pod を作成すると、Webhook は、マルチホストの調整に必要な Pod ホスト名、ヘッドレス Service、TPU 環境変数を自動的に構成します。

Ray は、起動時に各ノードの TPU トポロジを検出し、TPU スライス間で分散ワークロードをスケジュールします。分散 TorchTPU ワークロードは、次の 2 つのレベルでオーケストレートできます。

Ray Core は、テンソル並列計算や強化学習（RL）パイプラインなどのカスタム分散実行ループに、トポロジを考慮したプレースメントグループとアクタースケジューリングを提供します。
Ray Train は、PyTorch Distributed Data Parallel（DDP）、PyTorch FSDP2、JAX 単一プログラム、複数データ（SPMD）など、マネージド分散データ並列トレーニングとモデル並列トレーニングを提供します。これには、組み込みのチェックポイント機能と自動障害復旧機能が含まれます。

目標

このチュートリアルでは、次の方法を説明します。

TPU Trillium（v6e）または Ironwood（TPU7x）で Autopilot または Standard の GKE クラスタを作成します。
TorchTPU ベースイメージの上に Ray と Hugging Face のトレーニングライブラリをレイヤ化したコンテナイメージをビルドします。
RayCluster カスタムリソースをデプロイして、ヘッドノードと TPU ワーカーノードをプロビジョニングします。
Ray Core を使用して、TPU スライス全体で PyTorch DeviceMesh と DTensor を使用して分散テンソル並列計算を実行します。
Ray Train を使用して、マルチホスト TPU スライスで PyTorch FSDP2 を使用して Gemma 4 モデルをファインチューニングする。

費用

のドキュメントでは、課金対象である次の Google Cloudコンポーネントを使用します。

料金計算ツールを使うと、予想使用量に基づいて費用の見積もりを算出できます。

このチュートリアルに記載されているタスクの完了後、作成したリソースを削除すると、それ以上の請求は発生しません。詳細については、クリーンアップをご覧ください。

始める前に

作業を始める前に、次のタスクが完了していることを確認してください。

Artifact Registry API、Cloud Storage API、Google Kubernetes Engine API を有効にします。

API を有効にする

このタスクに Google Cloud CLI を使用する場合は、gcloud CLI をインストールして初期化します。gcloud CLI をインストール済みの場合は、gcloud components update コマンドを実行して最新のバージョンを取得します。以前のバージョンの gcloud CLI では、このドキュメントのコマンドを実行できない場合があります。
注: gcloud CLI がすでにインストールされている場合には、必ず compute/region プロパティを設定してください。主にゾーンクラスタを使用する場合は、代わりに compute/zone を設定します。デフォルトのロケーションを設定することで、gcloud CLI のエラー（One of [--zone, --region] must be supplied: Please specify location など）を防止できます。クラスタのロケーションが設定したデフォルトと異なる場合は、特定のコマンドでロケーションの指定が必要になることがあります。

選択したゾーンで、プロジェクトに TPU Trillium（v6e）または Ironwood（TPU7x）の容量に十分な割り当てがあることを確認します。詳細については、Cloud TPU の割り当てをご覧ください。
アカウントに次の IAM ロールがあることを確認します。
- Kubernetes Engine 管理者（roles/container.admin）
- Artifact Registry 管理者（roles/artifactregistry.admin）
- Storage 管理者（roles/storage.admin）

環境を準備する

このチュートリアルでは、Cloud Shell を使用して Google Cloudでホストされているリソースを管理します。Cloud Shell には、このチュートリアルに必要なソフトウェア（kubectl、docker、python3、gcloud CLI など）がプリインストールされています。

Cloud Shell を使用して環境を設定するには、次の操作を行います。

Google Cloud コンソールで（Cloud Shell をアクティブにする）をクリックして、Google Cloud コンソールで Cloud Shell セッションを起動します。このアクションにより、 Google Cloud コンソールの下部ペインでセッションが起動します。
kubernetes-engine-samples リポジトリのクローンを作成し、このチュートリアルのサンプルディレクトリに移動します。
```
git clone https://github.com/GoogleCloudPlatform/kubernetes-engine-samples.git
cd kubernetes-engine-samples/ai-ml/gke-ray/raytrain/torchtpu
```
Ray CLI がインストールされたローカル Python 仮想環境を作成して有効にします。
```
python3 -m venv ~/ray-env
source ~/ray-env/bin/activate
pip install "ray[default]==2.58.0"
```
Gemma 4 モデルの重みをダウンロードするには、Hugging Face で Google のライセンス契約に同意します。
1. Hugging Face の Gemma 4 モデルページに移動します。
2. ログインし、[同意してリポジトリにアクセス] をクリックしてライセンス条項に同意します。
3. Hugging Face アカウントの設定で、Read 権限を持つユーザーアクセストークンを生成します。
Google Cloud プロジェクト、クラスタ、Cloud Storage バケット、Hugging Face トークンの環境変数を設定します。
```
export PROJECT_ID=PROJECT_ID
export CLUSTER_NAME=CLUSTER_NAME
export REGION=REGION
export ZONE=ZONE
export ARTIFACT_REGISTRY=ARTIFACT_REGISTRY
export BUCKET_NAME=BUCKET_NAME
export HF_TOKEN=HF_TOKEN
export IMAGE_URI=${REGION}-docker.pkg.dev/${PROJECT_ID}/${ARTIFACT_REGISTRY}/ray-torchtpu:latest
gcloud config set project ${PROJECT_ID}
```
次のように置き換えます。
- PROJECT_ID: 実際の Google Cloud プロジェクト ID。
- CLUSTER_NAME: GKE クラスタの名前（ray-torchtpu-cluster など）。
- REGION: Artifact Registry リポジトリ、Cloud Storage バケット、TPU 容量が配置されているリージョン。
- ZONE: TPU の容量が使用可能なゾーン。TPU バージョンごとに使用可能なリージョンとゾーンについては、TPU のリージョンとゾーンをご覧ください。
- ARTIFACT_REGISTRY: Artifact Registry Docker リポジトリの名前（例: ray-torchtpu）。
- BUCKET_NAME: Ray Train チェックポイントを保存する Cloud Storage バケットの名前（PROJECT_ID-ray-torchtpu など）。
- HF_TOKEN: Hugging Face アクセストークン。

GKE クラスタを作成する

マネージド Ray Operator アドオンを有効にして、GKE Autopilot クラスタまたは Standard クラスタで分散 TorchTPU トレーニングワークロードを実行できます。

GKE クラスタを作成し、kubectl コマンドラインツールの認証情報を構成します。
Autopilot
1. Ray Operator アドオンを有効にして、リージョン GKE Autopilot クラスタを作成します。
  
  gcloud container clusters create-auto ${CLUSTER_NAME} \ --project=${PROJECT_ID} \ --region=${REGION} \ --cluster-version=1.36 \ --enable-ray-operator
  
  後のステップで RayCluster カスタムリソースをデプロイすると、GKE Autopilot は、Pod リソースリクエストとノードセレクタに一致する CPU ヘッドノードとマルチホスト TPU スライスを自動的にプロビジョニングします。
2. クラスタと通信するように kubectl コマンドラインツールを構成します。
  
  gcloud container clusters get-credentials ${CLUSTER_NAME} \ --project=${PROJECT_ID} \ --region=${REGION}
標準
1. Ray ヘッド Pod、イメージストリーミング、GKE Dataplane V2、Workload Identity Federation for GKE を有効にして、CPU ノードプールを含むゾーン GKE Standard クラスタを作成します。
  
  gcloud container clusters create ${CLUSTER_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --cluster-version=1.36 \ --machine-type=n2-standard-8 \ --num-nodes=1 \ --addons=RayOperator \ --enable-image-streaming \ --enable-dataplane-v2 \ --workload-pool=${PROJECT_ID}.svc.id.goog
  注: GKE RayOperator アドオンを使用せずにオープンソースの KubeRay をセルフマネージドする場合は、TorchTPU 環境変数の挿入用に KubeRay と KubeRay TPU Webhook バージョン 1.4.0 以降をインストールします。
2. クラスタと通信するように kubectl コマンドラインツールを構成します。
  
  gcloud container clusters get-credentials ${CLUSTER_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE}
3. TPU Trillium（v6e）の 4 つのホストにまたがる 4x4 トポロジ、または Ironwood（TPU7x）の 2 つのホストにまたがる 2x2x2 トポロジを使用して、16 デバイスのマルチホスト TPU スライスノードプールを作成します。
  
  Trillium（v6e）
  gcloud container node-pools create tpu-multi-host-pool \ --project=${PROJECT_ID} \ --cluster=${CLUSTER_NAME} \ --zone=${ZONE} \ --machine-type=ct6e-standard-4t \ --tpu-topology=4x4 \ --num-nodes=4
  
  Ironwood（TPU7x）
  gcloud container node-pools create tpu-multi-host-pool \ --project=${PROJECT_ID} \ --cluster=${CLUSTER_NAME} \ --zone=${ZONE} \ --machine-type=tpu7x-standard-4t \ --tpu-topology=2x2x2 \ --num-nodes=2
  
  ヒント: 特定の TPU 予約を使用するには、--reservation-affinity=specific フラグと --reservation=RESERVATION_NAME フラグを node-pools create コマンドに追加します。ここで、RESERVATION_NAME は TPU 予約の名前です。Spot VM を使用するには、--spot フラグを追加します。

マルチホスト TPU ノード間で共有される Ray Train チェックポイントを保存する Cloud Storage バケットを作成し、default Kubernetes ServiceAccount リソースにバケットへのアクセス権を付与します。

gcloud storage buckets create gs://${BUCKET_NAME} --location=${REGION}

PROJECT_NUMBER=$(gcloud projects describe ${PROJECT_ID} --format="value(projectNumber)")
gcloud storage buckets add-iam-policy-binding gs://${BUCKET_NAME} \
    --member="principal://iam.googleapis.com/projects/${PROJECT_NUMBER}/locations/global/workloadIdentityPools/${PROJECT_ID}.svc.id.goog/subject/ns/default/sa/default" \
    --role="roles/storage.objectAdmin"

Ray Pod がゲート付きの Gemma 4 の重みをダウンロードできるように、Hugging Face アクセストークンを使用して Kubernetes Secret を作成します。
```
kubectl create secret generic hf-secret \
    --from-literal=hf_api_token=${HF_TOKEN}
```

コンテナイメージをビルドする

TorchTPU と libtpu ライブラリを含むビルド済みの TorchTPU ベースイメージに、Ray と Hugging Face トレーニングライブラリをインストールするコンテナイメージをビルドします。

Artifact Registry Docker リポジトリを作成し、Docker 認証を構成します。

gcloud artifacts repositories create ${ARTIFACT_REGISTRY} \
    --repository-format=docker \
    --location=${REGION}

gcloud auth configure-docker ${REGION}-docker.pkg.dev,us-docker.pkg.dev

サンプルディレクトリの Dockerfile を確認します。

FROM us-docker.pkg.dev/ml-oss-artifacts-transient/torch-tpu-docker-container/torch-tpu:nightly

# Install wget for KubeRay health probes, Ray (Core, Dashboard, and Train), and Hugging Face libraries for fine-tuning
RUN apt-get update && apt-get install -y --no-install-recommends wget \
    && rm -rf /var/lib/apt/lists/* \
    && pip install --no-cache-dir \
    "ray[default,train]==2.58.0" \
    "transformers>=5.5.0" \
    datasets

Ray TorchTrainer クラスを使用したシングルスライストレーニングと、Ironwood（TPU7x）でのチップレットごとのスケジューリングには、Ray バージョン 2.58.0 以降が必要です。Gemma 4 のファインチューニングには、transformers ライブラリバージョン 5.5.0 以降が必要です。

イメージをビルドして Artifact Registry に push します。
```
docker build -t ${IMAGE_URI} .
docker push ${IMAGE_URI}
```

RayCluster カスタムリソースをデプロイする

CPU ヘッド Pod、Ray 共有メモリオブジェクトストア用のインメモリ /dev/shm ボリューム、マルチホスト TPU ワーカーグループをプロビジョニングする RayCluster カスタムリソースをデプロイします。

サンプルリポジトリで、TPU バージョンの RayCluster マニフェストを確認します。

Trillium（v6e）

apiVersion: ray.io/v1
kind: RayCluster
metadata:
  name: ray-torchtpu-cluster
spec:
  headGroupSpec:
    rayStartParams:
      num-cpus: "0"
    template:
      spec:
        containers:
        - name: ray-head
          image: ${IMAGE_URI}
          env:
          - name: HF_TOKEN
            valueFrom:
              secretKeyRef:
                name: hf-secret
                key: hf_api_token
          ports:
          - containerPort: 6379
            name: gcs
          - containerPort: 8265
            name: dashboard
          - containerPort: 10001
            name: client
          resources:
            limits:
              cpu: "4"
              memory: "16Gi"
            requests:
              cpu: "4"
              memory: "16Gi"
  workerGroupSpecs:
  - groupName: v6e-16
    replicas: 1
    numOfHosts: 4
    template:
      spec:
        containers:
        - name: ray-worker
          image: ${IMAGE_URI}
          env:
          - name: HF_TOKEN
            valueFrom:
              secretKeyRef:
                name: hf-secret
                key: hf_api_token
          resources:
            limits:
              cpu: "64"
              memory: "200Gi"
              google.com/tpu: "4"
            requests:
              cpu: "64"
              memory: "200Gi"
              google.com/tpu: "4"
          volumeMounts:
          - mountPath: /dev/shm
            name: dshm
        volumes:
        - name: dshm
          emptyDir:
            medium: Memory
            sizeLimit: 100Gi
        nodeSelector:
          cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
          cloud.google.com/gke-tpu-topology: 4x4

Ironwood（TPU7x）

apiVersion: ray.io/v1
kind: RayCluster
metadata:
  name: ray-torchtpu-cluster
spec:
  headGroupSpec:
    rayStartParams:
      num-cpus: "0"
    template:
      spec:
        containers:
        - name: ray-head
          image: ${IMAGE_URI}
          env:
          - name: HF_TOKEN
            valueFrom:
              secretKeyRef:
                name: hf-secret
                key: hf_api_token
          - name: RAY_TPU_RESOURCE_PER_CHIP
            value: "2"
          ports:
          - containerPort: 6379
            name: gcs
          - containerPort: 8265
            name: dashboard
          - containerPort: 10001
            name: client
          resources:
            limits:
              cpu: "4"
              memory: "16Gi"
            requests:
              cpu: "4"
              memory: "16Gi"
  workerGroupSpecs:
  - groupName: tpu7x-16
    replicas: 1
    numOfHosts: 2
    resources:
      # Advertises 8 logical TPU cores (2 chiplets per physical chip) per host.
      TPU: "8"
    template:
      spec:
        containers:
        - name: ray-worker
          image: ${IMAGE_URI}
          env:
          - name: HF_TOKEN
            valueFrom:
              secretKeyRef:
                name: hf-secret
                key: hf_api_token
          - name: RAY_TPU_RESOURCE_PER_CHIP
            value: "2"
          resources:
            limits:
              cpu: "64"
              memory: "250Gi"
              google.com/tpu: "4"
            requests:
              cpu: "64"
              memory: "250Gi"
              google.com/tpu: "4"
          volumeMounts:
          - mountPath: /dev/shm
            name: dshm
        volumes:
        - name: dshm
          emptyDir:
            medium: Memory
            sizeLimit: 128Gi
        nodeSelector:
          cloud.google.com/gke-tpu-accelerator: tpu7x
          cloud.google.com/gke-tpu-topology: 2x2x2

Ironwood（TPU7x）では、各物理チップに 2 つのチップレットが含まれており、4 チップホストあたり 8 つの論理 TPU コアが提供されます。このマニフェストでは、resources セクションで TPU フィールドの値を "8" に設定し、RAY_TPU_RESOURCE_PER_CHIP 環境変数を "2" に設定して、ホストごとに 8 つの論理 TPU リソースをアドバタイズしてスケジュールします。

マニフェストを適用して RayCluster カスタムリソースを作成します。

Trillium（v6e）

envsubst < ray-cluster.tpu-v6e-multihost.yaml | kubectl apply -f -

Ironwood（TPU7x）

envsubst < ray-cluster.tpu-v7x-multihost.yaml | kubectl apply -f -

RayCluster ヘッド Pod とすべての TPU ワーカー Pod が Running 状態になるまで待ちます。

kubectl get pods -l ray.io/cluster=ray-torchtpu-cluster

出力は次のようになります。

Trillium（v6e）

NAME                                         READY   STATUS    RESTARTS   AGE
ray-torchtpu-cluster-head-pdpk6              1/1     Running   0          2m30s
ray-torchtpu-cluster-v6e-16-worker-brwpc     1/1     Running   0          2m30s
ray-torchtpu-cluster-v6e-16-worker-dx4w7     1/1     Running   0          2m30s
ray-torchtpu-cluster-v6e-16-worker-h4d4w     1/1     Running   0          2m30s
ray-torchtpu-cluster-v6e-16-worker-t2g24     1/1     Running   0          2m30s

Ironwood（TPU7x）

NAME                                         READY   STATUS    RESTARTS   AGE
ray-torchtpu-cluster-head-k9x2m              1/1     Running   0          2m30s
ray-torchtpu-cluster-tpu7x-16-worker-7cwpq   1/1     Running   0          2m30s
ray-torchtpu-cluster-tpu7x-16-worker-m98r9   1/1     Running   0          2m30s

いずれかのワーカー Pod の環境変数を調べて、KubeRay TPU Webhook がマルチホスト TPU 調整変数と TorchTPU トポロジ変数を挿入したことを確認します。

kubectl get pods -l ray.io/node-type=worker \
    -o=jsonpath='{range .items[0].spec.containers[0].env[*]}{.name}={.value}{"\n"}{end}' \
    | grep -E "TPU_WORKER_ID|TPU_WORKER_HOSTNAMES|TPU_NAME|TORCH_TPU_TOPOLOGY|TORCH_TPU_SLICEBUILDER_ADDRESSES"

出力は次のようになります。

Trillium（v6e）

TPU_WORKER_HOSTNAMES=v6e-16-0-0.ray-torchtpu-cluster-headless,v6e-16-0-1.ray-torchtpu-cluster-headless,v6e-16-0-2.ray-torchtpu-cluster-headless,v6e-16-0-3.ray-torchtpu-cluster-headless
TORCH_TPU_TOPOLOGY=4,4,1
TORCH_TPU_SLICEBUILDER_ADDRESSES=v6e-16-0-0.ray-torchtpu-cluster-headless:8471,v6e-16-0-0.ray-torchtpu-cluster-headless:8472,v6e-16-0-0.ray-torchtpu-cluster-headless:8473,v6e-16-0-0.ray-torchtpu-cluster-headless:8474,v6e-16-0-1.ray-torchtpu-cluster-headless:8471,v6e-16-0-1.ray-torchtpu-cluster-headless:8472,v6e-16-0-1.ray-torchtpu-cluster-headless:8473,v6e-16-0-1.ray-torchtpu-cluster-headless:8474,v6e-16-0-2.ray-torchtpu-cluster-headless:8471,v6e-16-0-2.ray-torchtpu-cluster-headless:8472,v6e-16-0-2.ray-torchtpu-cluster-headless:8473,v6e-16-0-2.ray-torchtpu-cluster-headless:8474,v6e-16-0-3.ray-torchtpu-cluster-headless:8471,v6e-16-0-3.ray-torchtpu-cluster-headless:8472,v6e-16-0-3.ray-torchtpu-cluster-headless:8473,v6e-16-0-3.ray-torchtpu-cluster-headless:8474
TPU_WORKER_ID=3
TPU_NAME=v6e-16-0

Ironwood（TPU7x）

TPU_WORKER_HOSTNAMES=tpu7x-16-0-0.ray-torchtpu-cluster-headless,tpu7x-16-0-1.ray-torchtpu-cluster-headless
TORCH_TPU_TOPOLOGY=2,2,2,2
TORCH_TPU_SLICEBUILDER_ADDRESSES=tpu7x-16-0-0.ray-torchtpu-cluster-headless:8471,tpu7x-16-0-0.ray-torchtpu-cluster-headless:8472,tpu7x-16-0-0.ray-torchtpu-cluster-headless:8473,tpu7x-16-0-0.ray-torchtpu-cluster-headless:8474,tpu7x-16-0-0.ray-torchtpu-cluster-headless:8475,tpu7x-16-0-0.ray-torchtpu-cluster-headless:8476,tpu7x-16-0-0.ray-torchtpu-cluster-headless:8477,tpu7x-16-0-0.ray-torchtpu-cluster-headless:8478,tpu7x-16-0-1.ray-torchtpu-cluster-headless:8471,tpu7x-16-0-1.ray-torchtpu-cluster-headless:8472,tpu7x-16-0-1.ray-torchtpu-cluster-headless:8473,tpu7x-16-0-1.ray-torchtpu-cluster-headless:8474,tpu7x-16-0-1.ray-torchtpu-cluster-headless:8475,tpu7x-16-0-1.ray-torchtpu-cluster-headless:8476,tpu7x-16-0-1.ray-torchtpu-cluster-headless:8477,tpu7x-16-0-1.ray-torchtpu-cluster-headless:8478
TPU_WORKER_ID=1
TPU_NAME=tpu7x-16-0

Ray ヘッドサービスへのポート転送セッションを確立します。
```
kubectl port-forward svc/ray-torchtpu-cluster-head-svc 8265:8265 >/dev/null 2>&1 &
```
TPU ワーカーノードと自動検出された ray.io/tpu-* ラベルは、Ray ダッシュボードの [Cluster] タブで確認できます。

Ray Core で分散テンソル並列処理を実行する

Ray Core には、ray.util.tpu モジュールの slice_placement_group 関数が用意されています。この関数は、完全なマルチホスト TPU スライスを予約します。また、PlacementGroupSchedulingStrategy クラスは、スライス全体の論理 TPU デバイスごとに 1 つの Ray アクターをスケジュールします。

サンプルリポジトリの ray_core_collectives.py スクリプトを確認します。

import argparse
import os

import ray
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
from ray.util.tpu import (
    get_tpu_version_from_type,
    slice_placement_group,
)
import torch
import torch.distributed as dist
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.tensor.parallel import (
    ColwiseParallel,
    RowwiseParallel,
    parallelize_module,
)
import torch.nn as nn
import torch_tpu  # noqa: F401

RANDOM_SEED = 42


class FeedForward(nn.Module):
    """Two-layer feed-forward block for tensor-parallel sharding across a TPU slice."""

    def __init__(self, hidden_dim: int):
        super().__init__()
        self.w1 = nn.Linear(hidden_dim, 4 * hidden_dim, bias=False)
        self.act = nn.SiLU()
        self.w2 = nn.Linear(4 * hidden_dim, hidden_dim, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x + self.w2(self.act(self.w1(x)))


@ray.remote
class TorchTPUSliceWorker:
    """Ray actor bound to a single logical TPU device within a slice placement group."""

    def __init__(self, devices_per_host: int):
        self.node_ip = ray.util.get_node_ip_address()
        hosts = os.environ["TPU_WORKER_HOSTNAMES"].split(",")
        self.worker_id = int(os.environ.get("TPU_WORKER_ID", "0"))
        self.local_rank = int(
            ray.get_runtime_context().get_accelerator_ids()["TPU"][0]
        )
        self.rank = self.worker_id * devices_per_host + self.local_rank
        # Set per-process PyTorch env:// variables from the webhook's host-level
        # TPU metadata and Ray's assigned local device ID (TorchTrainer sets these
        # automatically in the Ray Train section below).
        os.environ.update(
            {
                "MASTER_ADDR": hosts[0],
                "MASTER_PORT": "29500",
                "WORLD_SIZE": str(len(hosts) * devices_per_host),
                "RANK": str(self.rank),
                "LOCAL_RANK": str(self.local_rank),
                "LOCAL_WORLD_SIZE": str(devices_per_host),
            }
        )

    def run_tensor_parallel_step(
        self,
        batch_size: int = 64,
        hidden_dim: int = 128,
    ) -> dict:
        dist.init_process_group(backend="tpu_dist")
        world_size = dist.get_world_size()
        device = torch.device("tpu")

        # 1. Compute single-device reference output on CPU with identical seed.
        torch.manual_seed(RANDOM_SEED)
        ref_model = FeedForward(hidden_dim).eval()
        inputs_cpu = torch.randn(batch_size, hidden_dim)
        with torch.no_grad():
            ref_output = ref_model(inputs_cpu)

        # 2. Shard the model across all TPU devices in the slice with DTensor TP.
        torch.manual_seed(RANDOM_SEED)
        tp_model = FeedForward(hidden_dim).to(device).eval()
        mesh = init_device_mesh("tpu", (world_size,))
        tp_model = parallelize_module(
            tp_model,
            mesh,
            {"w1": ColwiseParallel(), "w2": RowwiseParallel()},
        )

        with torch.no_grad():
            tp_output = tp_model(inputs_cpu.to(device)).cpu()

        torch.testing.assert_close(tp_output, ref_output, rtol=1e-3, atol=1e-2)
        max_abs_diff = float((tp_output - ref_output).abs().max().item())

        # 3. Verify direct ICI all_reduce across all ranks in the slice.
        rank_tensor = torch.full(
            (4,), float(self.rank + 1), dtype=torch.float32, device=device
        )
        dist.all_reduce(rank_tensor, op=dist.ReduceOp.SUM)
        all_reduce_sum = rank_tensor.cpu().tolist()

        dist.destroy_process_group()
        return {
            "rank": self.rank,
            "local_rank": self.local_rank,
            "worker_id": self.worker_id,
            "node_ip": self.node_ip,
            "all_reduce_sum": all_reduce_sum,
            "max_abs_diff": max_abs_diff,
        }


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Run distributed TorchTPU tensor parallelism across a Ray TPU slice."
    )
    parser.add_argument(
        "--topology",
        type=str,
        default="4x4",
        help="TPU physical topology.",
    )
    parser.add_argument(
        "--accelerator-type",
        type=str,
        default="TPU-V6E",
        help="Ray TPU accelerator type (TPU-V6E, TPU-V7X, etc.).",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    ray.init()

    slice_handle = slice_placement_group(
        topology=args.topology,
        accelerator_version=get_tpu_version_from_type(args.accelerator_type),
        num_slices=1,
        resources_per_bundle={"TPU": 1},
    )
    ray.get(slice_handle.placement_group.ready())

    devices_per_host = slice_handle.num_bundles // slice_handle.num_hosts
    workers = [
        TorchTPUSliceWorker.options(
            resources={"TPU": 1},
            scheduling_strategy=PlacementGroupSchedulingStrategy(
                placement_group=slice_handle.placement_group,
                placement_group_bundle_index=bundle_idx,
            ),
        ).remote(devices_per_host)
        for bundle_idx in range(slice_handle.num_bundles)
    ]
    results = sorted(
        ray.get([w.run_tensor_parallel_step.remote() for w in workers]),
        key=lambda r: r["rank"],
    )

    expected_sum = sum(range(1, slice_handle.num_bundles + 1))
    for r in results:
        print(
            f"Rank {r['rank']:2d} (worker_id={r['worker_id']}, "
            f"local_rank={r['local_rank']}, ip={r['node_ip']}): "
            f"all_reduce_sum={r['all_reduce_sum']}"
        )
        assert r["all_reduce_sum"] == [float(expected_sum)] * 4

    print(
        f"All {len(results)} TorchTPU ranks verified DTensor TP parity "
        f"(max_abs_diff={results[0]['max_abs_diff']:.6f}) and "
        f"all_reduce sum = {expected_sum:.1f}"
    )
    for w in workers:
        ray.kill(w)
    slice_handle.shutdown()


if __name__ == "__main__":
    main()

このスクリプトは、TPU スライス全体で分散テンソル並列処理を 3 つのステップで実行します。

slice_placement_group 関数は、16 デバイスの TPU スライスを予約し、各 TPU デバイスに 1 つの TorchTPUSliceWorker アクターをスケジュールします。
各アクターは、ホストインデックスとローカルデバイス ID からグローバルランクを計算し、tpu_dist プロセスグループを初期化します。
init_device_mesh 関数と parallelize_module 関数は、16 台の TPU デバイスすべてにフィードフォワードレイヤをシャーディングし、単一デバイスの CPU リファレンスに対して出力を検証し、ICI で all_reduce コレクティブを実行します。

TPU バージョンの Ray Core ジョブを送信します。

Trillium（v6e）

ray job submit \
    --address http://localhost:8265 \
    --working-dir . \
    --runtime-env-json '{"excludes": ["ray-env", ".git", "*.whl"]}' \
    -- python3 ray_core_collectives.py --topology=4x4 --accelerator-type=TPU-V6E

Ironwood（TPU7x）

ray job submit \
    --address http://localhost:8265 \
    --working-dir . \
    --runtime-env-json '{"excludes": ["ray-env", ".git", "*.whl"]}' \
    -- python3 ray_core_collectives.py --topology=2x2x2 --accelerator-type=TPU-V7X

出力は次のようになります。

Trillium（v6e）

Rank  0 (worker_id=0, local_rank=0, ip=10.72.3.20): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  1 (worker_id=0, local_rank=1, ip=10.72.3.20): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  2 (worker_id=0, local_rank=2, ip=10.72.3.20): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  3 (worker_id=0, local_rank=3, ip=10.72.3.20): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  4 (worker_id=1, local_rank=0, ip=10.72.14.20): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  5 (worker_id=1, local_rank=1, ip=10.72.14.20): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  6 (worker_id=1, local_rank=2, ip=10.72.14.20): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  7 (worker_id=1, local_rank=3, ip=10.72.14.20): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  8 (worker_id=2, local_rank=0, ip=10.72.7.20): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  9 (worker_id=2, local_rank=1, ip=10.72.7.20): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank 10 (worker_id=2, local_rank=2, ip=10.72.7.20): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank 11 (worker_id=2, local_rank=3, ip=10.72.7.20): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank 12 (worker_id=3, local_rank=0, ip=10.72.6.20): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank 13 (worker_id=3, local_rank=1, ip=10.72.6.20): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank 14 (worker_id=3, local_rank=2, ip=10.72.6.20): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank 15 (worker_id=3, local_rank=3, ip=10.72.6.20): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
All 16 TorchTPU ranks verified DTensor TP parity (max_abs_diff=0.002470) and all_reduce sum = 136.0
------------------------------------------
Job 'raysubmit_HFTkKrjSnfza3ZFD' succeeded
------------------------------------------

Ironwood（TPU7x）

Rank  0 (worker_id=0, local_rank=0, ip=172.16.14.6): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  1 (worker_id=0, local_rank=1, ip=172.16.14.6): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  2 (worker_id=0, local_rank=2, ip=172.16.14.6): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  3 (worker_id=0, local_rank=3, ip=172.16.14.6): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  4 (worker_id=0, local_rank=4, ip=172.16.14.6): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  5 (worker_id=0, local_rank=5, ip=172.16.14.6): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  6 (worker_id=0, local_rank=6, ip=172.16.14.6): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  7 (worker_id=0, local_rank=7, ip=172.16.14.6): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  8 (worker_id=1, local_rank=0, ip=172.16.15.7): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  9 (worker_id=1, local_rank=1, ip=172.16.15.7): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank 10 (worker_id=1, local_rank=2, ip=172.16.15.7): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank 11 (worker_id=1, local_rank=3, ip=172.16.15.7): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank 12 (worker_id=1, local_rank=4, ip=172.16.15.7): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank 13 (worker_id=1, local_rank=5, ip=172.16.15.7): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank 14 (worker_id=1, local_rank=6, ip=172.16.15.7): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank 15 (worker_id=1, local_rank=7, ip=172.16.15.7): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
All 16 TorchTPU ranks verified DTensor TP parity (max_abs_diff=0.002470) and all_reduce sum = 136.0
------------------------------------------
Job 'raysubmit_MrJbAe3hvARBcBN2' succeeded
------------------------------------------

Ray Train と PyTorch FSDP2 を使用して Gemma 4 モデルをファインチューニングする

Ray Train の TorchTrainer クラスは、マルチホスト TPU スライス全体で TPU スライスの予約、ワーカーの配置、torch.distributed の初期化を自動化します。

サンプルリポジトリの ray_train_llm_finetune.py スクリプトを確認します。このスクリプトは、databricks/databricks-dolly-15k 指示データセットで google/gemma-4-E2B-it モデルをファインチューニングします。

import argparse
import os
import tempfile
import time

from datasets import load_dataset
import ray
import ray.train
from ray.train import (
    Checkpoint,
    CheckpointConfig,
    RunConfig,
    ScalingConfig,
)
from ray.train.torch import TorchConfig, TorchTrainer
import torch
from torch.distributed import ReduceOp, all_reduce
from torch.distributed.checkpoint.state_dict import (
    StateDictOptions,
    get_model_state_dict,
)
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.fsdp import MixedPrecisionPolicy, fully_shard
from torch.utils.data import DataLoader
import torch_tpu  # noqa: F401
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    DataCollatorForLanguageModeling,
)


def create_dataloader(
    tokenizer: AutoTokenizer,
    batch_size: int,
    dataset_split: str = "train[:1024]",
    max_seq_len: int = 512,
) -> DataLoader:
    """Loads and tokenizes the Dolly 15k dataset for causal LM fine-tuning."""
    dataset = load_dataset(
        "databricks/databricks-dolly-15k", split=dataset_split
    )
    tokenized = dataset.map(
        lambda x: tokenizer(
            f"Instruction: {x['instruction']}\nResponse: {x['response']}",
            truncation=True,
            max_length=max_seq_len,
            padding="max_length",
        ),
        remove_columns=dataset.column_names,
    )
    dataloader = DataLoader(
        tokenized,
        batch_size=batch_size,
        shuffle=True,
        drop_last=True,
        collate_fn=DataCollatorForLanguageModeling(
            tokenizer=tokenizer, mlm=False
        ),
    )
    return ray.train.torch.prepare_data_loader(dataloader)


def train_func(config: dict) -> None:
    """Per-worker distributed FSDP2 training loop executed on each logical TPU core."""
    ctx = ray.train.get_context()
    world_rank = ctx.get_world_rank()
    world_size = ctx.get_world_size()
    device = ray.train.torch.get_device()

    model_id = config["model_id"]
    dataset_split = config["dataset_split"]
    batch_size = config["batch_size_per_device"]
    num_epochs = config["num_epochs"]
    max_seq_len = 512

    tokenizer = AutoTokenizer.from_pretrained(model_id)
    dataloader = create_dataloader(
        tokenizer=tokenizer,
        batch_size=batch_size,
        dataset_split=dataset_split,
        max_seq_len=max_seq_len,
    )

    torch.manual_seed(42)
    model = AutoModelForCausalLM.from_pretrained(
        model_id,
        dtype=torch.bfloat16,
        attn_implementation="eager",
        vision_config=None,
        audio_config=None,
    ).to(device)
    model.config.get_text_config().use_cache = False

    mesh = init_device_mesh("tpu", (world_size,))
    mp_policy = MixedPrecisionPolicy(
        param_dtype=torch.bfloat16,
        reduce_dtype=torch.float32,
    )
    text_model = model.model.language_model
    text_model.embed_tokens_per_layer.requires_grad_(False)
    for layer in text_model.layers:
        fully_shard(layer.mlp, mesh=mesh, mp_policy=mp_policy)
    fully_shard(
        model,
        mesh=mesh,
        mp_policy=mp_policy,
        ignored_params=set(text_model.embed_tokens_per_layer.parameters()),
    )

    optimizer = torch.optim.AdamW(
        model.parameters(), lr=2e-5, weight_decay=0.01
    )
    global_batch_size = batch_size * world_size

    for epoch in range(num_epochs):
        dataloader.sampler.set_epoch(epoch)
        model.train()
        epoch_start = time.perf_counter()
        running_loss = torch.zeros(1, device=device, dtype=torch.float32)

        for batch in dataloader:
            optimizer.zero_grad()
            outputs = model(
                input_ids=batch["input_ids"],
                labels=batch["labels"],
            )
            outputs.loss.backward()
            optimizer.step()
            running_loss += outputs.loss.detach().float()

        all_reduce(running_loss, op=ReduceOp.AVG)
        avg_loss = running_loss.item() / len(dataloader)
        epoch_time = time.perf_counter() - epoch_start
        tokens_per_sec = (
            len(dataloader) * global_batch_size * max_seq_len
        ) / epoch_time

        metrics = {
            "epoch": epoch + 1,
            "loss": round(avg_loss, 4),
            "epoch_time_s": round(epoch_time, 2),
            "tokens_per_s": round(tokens_per_sec, 1),
        }
        if world_rank == 0:
            print(
                f"Epoch {epoch + 1}/{num_epochs} | loss={metrics['loss']:.4f} | "
                f"epoch_time={metrics['epoch_time_s']:.2f}s | "
                f"tokens/s={metrics['tokens_per_s']:.1f}"
            )

        full_state_dict = get_model_state_dict(
            model,
            options=StateDictOptions(
                full_state_dict=True,
                cpu_offload=True,
                ignore_frozen_params=True,
            ),
        )
        if world_rank == 0:
            with tempfile.TemporaryDirectory() as tmpdir:
                torch.save(full_state_dict, os.path.join(tmpdir, "model.pt"))
                ray.train.report(
                    metrics, checkpoint=Checkpoint.from_directory(tmpdir)
                )
        else:
            ray.train.report(metrics)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Fine-tune a Hugging Face CausalLM on TPUs with Ray Train and TorchTPU FSDP2."
    )
    parser.add_argument(
        "--model-id",
        type=str,
        default="google/gemma-4-E2B-it",
        help="Hugging Face model ID to fine-tune.",
    )
    parser.add_argument(
        "--topology",
        type=str,
        default="4x4",
        help="TPU slice topology.",
    )
    parser.add_argument(
        "--accelerator-type",
        type=str,
        default="TPU-V6E",
        help="Ray TPU accelerator type (TPU-V6E, TPU-V7X, etc.).",
    )
    parser.add_argument(
        "--num-workers",
        type=int,
        default=16,
        help="Total TPU worker processes across the slice.",
    )
    parser.add_argument(
        "--dataset-split",
        type=str,
        default="train[:1024]",
        help="Hugging Face dataset split to load.",
    )
    parser.add_argument(
        "--num-epochs",
        type=int,
        default=3,
        help="Number of fine-tuning epochs.",
    )
    parser.add_argument(
        "--batch-size-per-device",
        type=int,
        default=4,
        help="Per-device training batch size.",
    )
    parser.add_argument(
        "--storage-path",
        type=str,
        default=os.environ.get("STORAGE_PATH"),
        required="STORAGE_PATH" not in os.environ,
        help="Cloud Storage or shared file system path for checkpoints.",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    ray.init()
    run_name = f"torchtpu_gemma4_fsdp2_sft_{int(time.time())}"

    print(
        f"Starting Ray Train FSDP2 SFT job: model={args.model_id}, "
        f"topology={args.topology}, accelerator_type={args.accelerator_type}, "
        f"num_workers={args.num_workers}"
    )

    trainer = TorchTrainer(
        train_loop_per_worker=train_func,
        train_loop_config={
            "model_id": args.model_id,
            "dataset_split": args.dataset_split,
            "num_epochs": args.num_epochs,
            "batch_size_per_device": args.batch_size_per_device,
        },
        torch_config=TorchConfig(backend="tpu_dist"),
        scaling_config=ScalingConfig(
            use_tpu=True,
            num_workers=args.num_workers,
            topology=args.topology,
            accelerator_type=args.accelerator_type,
            resources_per_worker={"TPU": 1},
        ),
        run_config=RunConfig(
            name=run_name,
            storage_path=args.storage_path,
            checkpoint_config=CheckpointConfig(
                num_to_keep=1,
                checkpoint_score_attribute="loss",
                checkpoint_score_order="min",
            ),
        ),
    )

    result = trainer.fit()
    print(f"Training completed! Final metrics: {result.metrics}")
    print(f"Saved checkpoint: {result.checkpoint}")


if __name__ == "__main__":
    main()

このスクリプトは、次のコンポーネントを使用して TorchTrainer クラスを構成します。

ScalingConfig クラスは、トレーニングジョブの分散ワーカートポロジを構成します。use_tpu=True、num_workers=16、topology、accelerator_type の各パラメータを設定すると、16 デバイスの TPU スライスが予約され、TPU コアごとに 1 つのワーカーがスケジュールされます。
TorchConfig クラスは、ワーカー全体で torch.distributed プロセスグループを構成します。backend="tpu_dist" パラメータを設定すると、train_func 関数が実行される前に、16 個のワーカーすべてが TPU ICI ネットワーク経由で接続されます。
create_dataloader 関数は、指示データセットをダウンロードしてトークン化し、PyTorch DataLoader オブジェクトを使用してトークンシーケンスをバッチ処理します。また、ray.train.torch.prepare_data_loader 関数を呼び出して、各ワーカーが割り当てられた TPU デバイスで各グローバルバッチの互いに素なシャードを受け取るようにします。
train_func 関数は、bfloat16 精度でモデルを読み込み、各フィードフォワードブロックとルートモデルで PyTorch FSDP2 fully_shard 関数を呼び出して、16 個の TPU コアすべてにパラメータ、勾配、オプティマイザーの状態をシャーディングします。
Checkpoint クラスと ray.train.report 関数は、各エポックの終了時にランク 0 で完全なモデル状態辞書を収集し、チェックポイントを Cloud Storage バケットにアップロードします。

TPU バージョンの TorchTrainer FSDP2 ファインチューニングジョブを送信します。

Trillium（v6e）

ray job submit \
    --address http://localhost:8265 \
    --working-dir . \
    --runtime-env-json '{"excludes": ["ray-env", ".git", "*.whl"]}' \
    -- python3 ray_train_llm_finetune.py \
        --topology=4x4 \
        --accelerator-type=TPU-V6E \
        --storage-path=gs://${BUCKET_NAME}/ray_train_torchtpu

Ironwood（TPU7x）

ray job submit \
    --address http://localhost:8265 \
    --working-dir . \
    --runtime-env-json '{"excludes": ["ray-env", ".git", "*.whl"]}' \
    -- python3 ray_train_llm_finetune.py \
        --topology=2x2x2 \
        --accelerator-type=TPU-V7X \
        --storage-path=gs://${BUCKET_NAME}/ray_train_torchtpu

出力は次のようになります。

Trillium（v6e）

Starting Ray Train FSDP2 SFT job: model=google/gemma-4-E2B-it, topology=4x4, accelerator_type=TPU-V6E, num_workers=16
(TrainController pid=5407) Attempting to start training worker group of size 16 with the following resources: [{'TPU': 1, 'accelerator_type:TPU-V6E': 0.001}] * 16
(RayTrainWorker pid=19733, ip=10.72.3.20) Setting up process group for: env:// [rank=0, world_size=16]
(RayTrainWorker pid=19733, ip=10.72.3.20) Epoch 1/3 | loss=3.1402 | epoch_time=28.17s | tokens/s=18613.5
(RayTrainWorker pid=19733, ip=10.72.3.20) Checkpoint successfully created at: Checkpoint(filesystem=gcs, path=my-project-ray-torchtpu/ray_train_torchtpu/torchtpu_gemma4_fsdp2_sft_1791864427/checkpoint_2026-10-13_04-08-10.762175)
(RayTrainWorker pid=19733, ip=10.72.3.20) Epoch 2/3 | loss=2.3318 | epoch_time=26.07s | tokens/s=20112.8
(RayTrainWorker pid=19733, ip=10.72.3.20) Checkpoint successfully created at: Checkpoint(filesystem=gcs, path=my-project-ray-torchtpu/ray_train_torchtpu/torchtpu_gemma4_fsdp2_sft_1791864427/checkpoint_2026-10-13_04-10-02.548088)
(RayTrainWorker pid=19733, ip=10.72.3.20) Epoch 3/3 | loss=2.2050 | epoch_time=25.52s | tokens/s=20545.2
(RayTrainWorker pid=19733, ip=10.72.3.20) Checkpoint successfully created at: Checkpoint(filesystem=gcs, path=my-project-ray-torchtpu/ray_train_torchtpu/torchtpu_gemma4_fsdp2_sft_1791864427/checkpoint_2026-10-13_04-11-48.701853)
Training completed! Final metrics: {'epoch': 3, 'loss': 2.205, 'epoch_time_s': 25.52, 'tokens_per_s': 20545.2}
Saved checkpoint: Checkpoint(filesystem=gcs, path=my-project-ray-torchtpu/ray_train_torchtpu/torchtpu_gemma4_fsdp2_sft_1791864427/checkpoint_2026-10-13_04-11-48.701853)
------------------------------------------
Job 'raysubmit_zbm4pnFZitA1YECA' succeeded
------------------------------------------

Ironwood（TPU7x）

Starting Ray Train FSDP2 SFT job: model=google/gemma-4-E2B-it, topology=2x2x2, accelerator_type=TPU-V7X, num_workers=16
(TrainController pid=1162) Attempting to start training worker group of size 16 with the following resources: [{'TPU': 1, 'accelerator_type:TPU-V7X': 0.001}] * 16
(RayTrainWorker pid=257, ip=172.16.26.6) Setting up process group for: env:// [rank=0, world_size=16]
(RayTrainWorker pid=257, ip=172.16.26.6) Epoch 1/3 | loss=3.1407 | epoch_time=26.28s | tokens/s=19953.3
(RayTrainWorker pid=257, ip=172.16.26.6) Checkpoint successfully created at: Checkpoint(filesystem=gcs, path=my-project-ray-torchtpu/ray_train_torchtpu/torchtpu_gemma4_fsdp2_sft_1791867004/checkpoint_2026-10-13_04-52-00.026930)
(RayTrainWorker pid=257, ip=172.16.26.6) Epoch 2/3 | loss=2.3235 | epoch_time=16.68s | tokens/s=31429.9
(RayTrainWorker pid=257, ip=172.16.26.6) Checkpoint successfully created at: Checkpoint(filesystem=gcs, path=my-project-ray-torchtpu/ray_train_torchtpu/torchtpu_gemma4_fsdp2_sft_1791867004/checkpoint_2026-10-13_04-53-33.533341)
(RayTrainWorker pid=257, ip=172.16.26.6) Epoch 3/3 | loss=2.1964 | epoch_time=16.98s | tokens/s=30884.6
(RayTrainWorker pid=257, ip=172.16.26.6) Checkpoint successfully created at: Checkpoint(filesystem=gcs, path=my-project-ray-torchtpu/ray_train_torchtpu/torchtpu_gemma4_fsdp2_sft_1791867004/checkpoint_2026-10-13_04-55-01.508219)
Training completed! Final metrics: {'epoch': 3, 'loss': 2.1964, 'epoch_time_s': 16.98, 'tokens_per_s': 30884.6}
Saved checkpoint: Checkpoint(filesystem=gcs, path=my-project-ray-torchtpu/ray_train_torchtpu/torchtpu_gemma4_fsdp2_sft_1791867004/checkpoint_2026-10-13_04-55-01.508219)
------------------------------------------
Job 'raysubmit_zSNnCXhHkwKiSmJh' succeeded
------------------------------------------

トレーニング出力には次の情報が表示されます。

エポック 1 には、最初の XLA グラフのコンパイルが含まれます。エポック 2 と 3 では、/dev/shm/torch_tpu_cache からキャッシュに保存された実行可能ファイルを再利用し、TPU Trillium（v6e）では 1 秒あたり 20,000 トークン、Ironwood（TPU7x）では 1 秒あたり 31,000 トークンに達します。
16 個の FSDP2 ワーカーランクすべてで同期されたグローバルクロスエントロピー損失は、databricks/databricks-dolly-15k 指示データセットの 3 つのファインチューニングエポック全体で、エポック 1 の 3.14 からエポック 3 の 2.20 に減少しています。
各エポックの終了時に、ランク 0 は完全なモデル状態辞書を収集し、チェックポイントを Cloud Storage バケットにアップロードします。CheckpointConfig(num_to_keep=1, checkpoint_score_attribute="loss", checkpoint_score_order="min") オブジェクトで構成されているように、損失が最も低いチェックポイントを保持します。

ワークロードのパフォーマンスをモニタリングする

Cloud Monitoring で TPU TensorCore 使用率、HBM 使用率、Ray クラスタ指標をモニタリングするには、Ray on GKE クラスタのログと指標を収集して表示すると Ray ダッシュボードで TPU 指標を表示するをご覧ください。

Ray ダッシュボードにアクセスする

組み込みの Ray ダッシュボードで Ray アクター、プレースメントグループ、トレーニングジョブのステータスを検査する手順は次のとおりです。

Ray ヘッドサービスへのポート転送がまだ実行されていない場合は、バックグラウンドで開始します。
```
kubectl port-forward svc/ray-torchtpu-cluster-head-svc 8265:8265 >/dev/null 2>&1 &
```
Cloud Shell で、 [ウェブでプレビュー] をクリックし、[ポート 8265 でプレビュー] を選択するか、ローカルブラウザで http://localhost:8265 を開きます。
[Jobs]、[Cluster]、[Metrics] タブをクリックして、ワーカーログ、プレースメントグループの状態、TPU ノードラベル、[TPU Utilization by Node] グラフを確認します。

クリーンアップ

このチュートリアルで使用したリソースについて Google Cloud アカウントに課金されないようにするには、RayCluster カスタムリソース、GKE クラスタ、Cloud Storage バケット、Artifact Registry リポジトリを削除します。

バックグラウンドのポート転送プロセスを停止し、RayCluster カスタムリソースと Kubernetes Secret を削除します。

pkill -f "kubectl port-forward" || true
kubectl delete raycluster ray-torchtpu-cluster
kubectl delete secret hf-secret --ignore-not-found

GKE クラスタを削除します。

Autopilot

gcloud container clusters delete ${CLUSTER_NAME} \
    --region=${REGION}

標準

gcloud container clusters delete ${CLUSTER_NAME} \
    --zone=${ZONE}

Cloud Storage バケットとそのチェックポイントの内容を削除します。
```
gcloud storage rm --recursive gs://${BUCKET_NAME}
```

Artifact Registry リポジトリを削除します。

gcloud artifacts repositories delete ${ARTIFACT_REGISTRY} \
    --location=${REGION}

次のステップ

Cloud TPU での PyTorch の詳細を確認する。
TorchTPU のドキュメントを確認する。
Ray と TorchTPU を使用してマルチホスト TPU で LLM をサービングする方法を確認する。
DWS と Kueue を使用して TPU で AI トレーニングを最適化する方法を確認する。
GKE の TPU についてを確認する。
Ray Train のドキュメントを確認する。

Ray と TorchTPU を使用してマルチホスト TPU でモデルをファインチューニングする コレクションでコンテンツを整理 必要に応じて、コンテンツの保存と分類を行います。

背景

TPU

TorchTPU

TPU での KubeRay と Ray

目標

費用

始める前に

環境を準備する

GKE クラスタを作成する

Autopilot

標準

Trillium（v6e）

Ironwood（TPU7x）

コンテナ イメージをビルドする

RayCluster カスタム リソースをデプロイする

Trillium（v6e）

Ironwood（TPU7x）

Trillium（v6e）

Ironwood（TPU7x）

Trillium（v6e）

Ironwood（TPU7x）

Trillium（v6e）

Ironwood（TPU7x）

Ray Core で分散テンソル並列処理を実行する

Trillium（v6e）

Ironwood（TPU7x）

Trillium（v6e）

Ironwood（TPU7x）

Ray Train と PyTorch FSDP2 を使用して Gemma 4 モデルをファインチューニングする

Trillium（v6e）

Ironwood（TPU7x）

Trillium（v6e）

Ironwood（TPU7x）

ワークロードのパフォーマンスをモニタリングする

Ray ダッシュボードにアクセスする

クリーンアップ

Autopilot

標準

次のステップ

Ray と TorchTPU を使用してマルチホスト TPU でモデルをファインチューニングする

コンテナイメージをビルドする

RayCluster カスタムリソースをデプロイする