マルチホスト TPU v6e で Gemma 4 31B の SFT を実行する

このチュートリアルでは、MaxText と Cluster Toolkit を使用して、マルチホスト Tensor Processing Unit（TPU）v6e クラスタで教師ありファインチューニング（SFT）を実行する方法について説明します。Cluster Toolkit を使用してマルチホスト トレーニング ワークロードを実行し、結果を Hugging Face 形式にエクスポートしてサービングします。

ワークロードの概要

Cloud TPU と GKE でこの SFT ワークロードを実行するには、次のコンポーネントが必要です。

  • モデル: Gemma 4 31B（google/gemma-4-31b-it）。高度なオープンウェイト モデルです。このチュートリアルでは、モデルの重みを Hugging Face からダウンロードし、MaxText 互換形式に変換して、Cloud TPU v6e（v6e-32）インフラストラクチャを備えた GKE クラスタにデプロイします。
  • トレーニング方法: MaxText を使用した SFT。このワークロードは、JAX を使用して Cloud TPU v6e でスケーラブルなファインチューニング パイプラインを実行します。MaxText の組み込み SFT トレーナー（maxtext.trainers.post_train.sft.train_sft）と指示チューニング データセットを使用して、1,000 ステップにわたってモデル パラメータを最適化します。
  • データセットとテンプレート: UltraChat 200k（HuggingFaceH4/ultrachat_200k）。大規模で高品質なマルチターン ダイアログ データセット。ワークロードは、このデータセットを使用してモデルをファインチューニングし、会話機能と指示に従う機能を改善します。
  • インフラストラクチャとフレームワーク: Cloud TPU v6e（v6e-32、4x8 トポロジ）で、Cluster Toolkit を使用して管理される GKE。フレームワークは、JAX バックエンドを備えた MaxText（tpu_post_training:0.2.4）です。

目標

  • Cluster Toolkit とその依存関係をインストールします。
  • Cluster Toolkit クラスタをデプロイする。
  • Hugging Face モデルを MaxText 形式に変換する。
  • TPU で SFT トレーニング ワークロードを実行します。
  • ファインチューニングされたモデルを Hugging Face 形式に変換してサービングする。

費用

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

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

新規の Google Cloud ユーザーは無料トライアルをご利用いただける場合があります。

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

始める前に

  • このチュートリアルを使用するには、Hugging Face アクセス トークンが必要です。Hugging Face で無料アカウントに登録できます。アカウントを作成したら、アクセス トークンを生成します。

    1. [Welcome to Hugging Face] ページで、アカウントのアバターをクリックし、[アクセス トークン] を選択します。
    2. [アクセス トークン] ページで、[新しいトークンを作成] をクリックします。
    3. [Read] トークン タイプを選択し、トークンの名前を入力します。
    4. アクセス トークンが表示されます。トークンを安全な場所に保存します。

このチュートリアルを完了するために必要な権限を取得するには、プロジェクトに対する次の IAM ロールを付与するよう管理者に依頼してください。

ロールの付与については、プロジェクト、フォルダ、組織へのアクセス権の管理をご覧ください。

必要な権限は、カスタムロールや他の事前定義ロールから取得することもできます。

環境変数を設定する

環境変数を設定します。

export PROJECT="YOUR_PROJECT_ID"
export REGION="YOUR_REGION"
export ZONE="YOUR_ZONE"
export CLUSTER_NAME="YOUR_CLUSTER_NAME"
export GCS_BUCKET="YOUR_BUCKET_NAME"
export CLOUD_IMAGE_NAME="us-docker.pkg.dev/cloud-tpu-images/maxtext-images/tpu_post_training:0.2.4"
export TPU_TYPE="v6e-32"
export RESERVATION="YOUR_RESERVATION_NAME"
export HF_TOKEN="YOUR_HF_TOKEN"
export MODEL_NAME="gemma4-31b"
export CLUSTER_TOOLKIT_VERSION="v1.103.0"

次のように置き換えます。

  • YOUR_PROJECT_ID: 実際の Google Cloudプロジェクトの ID。
  • YOUR_REGION: TPU 予約が配置されているリージョン。クラスタをデプロイするリージョン。
  • YOUR_ZONE: TPU 予約が配置され、クラスタをデプロイするゾーン。
  • YOUR_CLUSTER_NAME: GKE クラスタの名前。名前は、先頭を英小文字にして、英小文字、数字、ハイフンのみで構成される必要があります。長さは 20 文字以下です。これにより、この名前に基づくサービス アカウント ID が 30 文字の制限内に収まります。
  • YOUR_BUCKET_NAME: Cloud Storage バケットのグローバルに一意の名前。
  • YOUR_RESERVATION_NAME: TPU v6e 予約の名前。
  • YOUR_HF_TOKEN: Hugging Face アクセス トークン。

デフォルトでは、スクリプトは MODEL_NAME を gemma4-31b（google/gemma-4-31b-it の MaxText モデル名）に設定します。

Cluster Toolkit の依存関係をインストールする

Linux または macOS のクライアントまたはワークステーションからこのチュートリアルを完了するには、Cluster Toolkit のドキュメントの依存関係をインストールするの手順に沿って操作します。

Cloud Shell を使用している場合は、このセクションをスキップできます。

Cluster Toolkit をインストールする

CLUSTER_TOOLKIT_VERSION（v1.103.0）で定義されている特定のバージョンの Cluster Toolkit のビルド済みバンドルをダウンロードして、現在の作業ディレクトリに展開します。

case "$(uname -s)-$(uname -m)" in
  Linux-x86_64)              BUNDLE="gcluster_bundle_linux_amd64.tgz" ;;
  Linux-aarch64|Linux-arm64) BUNDLE="gcluster_bundle_linux_arm64.tgz" ;;
  Darwin-x86_64)             BUNDLE="gcluster_bundle_mac_amd64.tgz" ;;
  Darwin-arm64)              BUNDLE="gcluster_bundle_mac_arm64.tgz" ;;
  *) BUNDLE=""; echo "Unsupported platform: $(uname -s)-$(uname -m)" >&2 ;;
esac
mkdir -p cluster-toolkit
curl -fsSL "https://github.com/GoogleCloudPlatform/cluster-toolkit/releases/download/${CLUSTER_TOOLKIT_VERSION:?}/${BUNDLE:?}" \
  | tar -xz -C cluster-toolkit
cd cluster-toolkit

環境を認証する

Cluster Toolkit がリソースをプロビジョニングし、イメージを Artifact Registry に push できるようにするには、Google Cloud CLI を認証してアプリケーションのデフォルト認証情報を設定します。

gcloud auth login
gcloud auth application-default login --no-launch-browser
gcloud config set project "${PROJECT:?}"
gcloud auth application-default set-quota-project "${PROJECT:?}"

gcloud auth configure-docker gcr.io --quiet
gcloud auth configure-docker "${REGION:?}-docker.pkg.dev" --quiet

Cluster Toolkit クラスタを作成する

32 個の v6e TPU チップを含む Cluster Toolkit クラスタを作成してデプロイするには、次の手順を完了します。

  1. Cloud Storage バケットを作成します。

    gcloud storage buckets create "gs://${GCS_BUCKET:?}" --project="${PROJECT:?}" --location="${REGION:?}" || true

  2. ブループリント構成をステージングし、サンプルのデプロイ ブループリントを分離された作業ディレクトリにコピーします。

    mkdir -p tmp
cp examples/gke-tpu-v6e/gke-tpu-v6e-advanced.yaml tmp/

  3. tmp/gke-tpu-v6e-advanced.yaml を更新して、クラスタ ノードプール サービス アカウント（gke-np-sa）に、Cloud Storage バケットにチェックポイントを書き込むためのストレージ管理者ロール（roles/storage.admin）を付与します。システム ノードプールのマシンタイプを e2-standard-8 に切り替え、システム ノードプールをターゲット ゾーンに固定して、リージョンのマシンタイプの在庫切れを回避します。

    # Grant the GKE Node Pool Service Account storage.admin access for checkpoint writes
sed -i "s/- storage.objectViewer/- storage.admin/" tmp/gke-tpu-v6e-advanced.yaml
# Change n2-standard-8 to e2-standard-8 to avoid GCE_STOCKOUT in some regions
sed -i "s/n2-standard-8/e2-standard-8/" tmp/gke-tpu-v6e-advanced.yaml
# Pin system node pool to a single zone to avoid cross-zone stockouts
sed -i '/system_node_pool_machine_type/a \      system_node_pool_zones: [$(vars.zone)]' tmp/gke-tpu-v6e-advanced.yaml

  4. gcluster deploy コマンドを使用して、ブループリント tmp/gke-tpu-v6e-advanced.yaml を使用して Cluster Toolkit クラスタをデプロイし、--vars フラグを使用して必要な変数を渡します。

    ./gcluster deploy tmp/gke-tpu-v6e-advanced.yaml \
    --vars "project_id=${PROJECT:?},deployment_name=${CLUSTER_NAME:?},region=${REGION:?},zone=${ZONE:?},num_slices=1,tpu_topology=4x8,authorized_cidr=0.0.0.0/0,reservation=${RESERVATION:-}" \
    --download-dependencies \
    -l IGNORE \
    --auto-approve -w

モデルを MaxText 形式に変換する

MaxText 形式でモデルをトレーニングするには、Hugging Face 形式から MaxText 形式に変換する必要があります。

  1. 以降のコマンドを簡素化するには、gcluster job config コマンドを使用して、デフォルトのプロジェクト、クラスタ、ロケーションを構成します。

    # Configure gcluster Defaults
./gcluster job config set project "${PROJECT:?}"
./gcluster job config set cluster "${CLUSTER_NAME:?}"
./gcluster job config set location "${REGION:?}"

  2. gcluster job submit コマンドを使用して、モデルを Hugging Face 形式から MaxText 形式に変換し、Cloud Storage バケットに保存します。

    ./gcluster job submit --name hf-to-mt \
    --compute-type "${TPU_TYPE:?}" \
    --num-slices 1 \
    --image "${CLOUD_IMAGE_NAME:?}" \
    --await-job-completion \
    --command "[ \"\$JOB_COMPLETION_INDEX\" != \"0\" ] || \
      python3 -m maxtext.checkpoint_conversion.to_maxtext \
        model_name=${MODEL_NAME:?} \
        hf_access_token=${HF_TOKEN:?} \
        base_output_directory=gs://${GCS_BUCKET:?}/${MODEL_NAME}/max-text-format/ \
        scan_layers=True \
        use_multimodal=False \
        skip_jax_distributed_system=true \
        checkpoint_storage_use_zarr3=0 \
        checkpoint_storage_use_ocdbt=0 \
        hardware=cpu \
        --lazy_load_tensors=True"

    このコマンドには --await-job-completion フラグが含まれています。これは、変換が完了するまで現在のターミナルで待機することを意味します。

  3. 省略可: ジョブの実行中に変換ログをストリーミングするには、gcluster を抽出した作業ディレクトリで別のターミナルを開き、gcluster job logs コマンドを実行します。

    # Use the list command to check status
./gcluster job list

# Check progress of the job (--main-only targets the coordinator pod (Job Index 0, Pod Index 0) to avoid duplicate logs from other workers)
./gcluster job logs hf-to-mt --main-only -f

  4. 元のターミナルで、変換されたモデルファイルが Cloud Storage バケットで使用可能であることを確認します。

    gcloud storage ls "gs://${GCS_BUCKET:?}/${MODEL_NAME:?}/max-text-format/"

トレーニング ワークロードを開始する

  1. 変換プロセスが完了したら、SFT ワークロードを開始します。

    ./gcluster job submit --name sft \
    --compute-type "${TPU_TYPE:?}" \
    --num-slices 1 \
    --image "${CLOUD_IMAGE_NAME:?}" \
    --await-job-completion \
    --command "JAX_PLATFORMS=tpu,cpu ENABLE_PJRT_COMPATIBILITY=true JAX_TRACEBACK_FILTERING=off LIBTPU_INIT_ARGS=' --xla_tpu_scoped_vmem_limit_kib=61440 --xla_tpu_bf16_emission_mode=NATIVE_EMISSION --xla_tpu_enable_sparse_core_collective_offload_all_reduce=true --xla_tpu_use_single_sparse_core_for_all_gather_offload=true ' \
      python3 -m maxtext.trainers.post_train.sft.train_sft \
      run_name=sft \
      base_output_directory=gs://${GCS_BUCKET:?}/${MODEL_NAME:?}/trained/ \
      model_name=${MODEL_NAME} \
      load_parameters_path=gs://${GCS_BUCKET}/${MODEL_NAME}/max-text-format/0/items/ \
      hf_access_token=${HF_TOKEN:?} \
      dataset_type=hf \
      hf_path=HuggingFaceH4/ultrachat_200k \
      per_device_batch_size=1 steps=1000 \
      profiler=xplane \
      checkpoint_storage_use_zarr3=0 \
      checkpoint_storage_use_ocdbt=0 \
      skip_jax_distributed_system=False"

    このコマンドには --await-job-completion フラグが含まれています。これは、トレーニング ジョブが完了するまで現在のターミナルで待機することを意味します。

  2. 省略可: ジョブの実行中にライブ トレーニング ログをストリーミングするには、gcluster を抽出した作業ディレクトリで別のターミナルを開き、gcluster job logs コマンドを実行します。

    # Use the list command to check status
./gcluster job list

# Check progress of the job (--main-only targets the coordinator pod (Job Index 0, Pod Index 0) to avoid duplicate logs from other workers)
./gcluster job logs sft --main-only -f

    最初の XLA コンパイルが完了すると、コーディネーターは 1, 000 ステップの実行全体で各ステップの出力の進行状況を completed step: <step>, seconds: ..., TFLOP/s/device: ..., loss: ... の形式でログに記録します。トレーニング損失は、ステップ 5 の約 1.57 からステップ 1,000 の約 0.82 に減少します。次のセクションに進む前に、ステップ 1000 が完了し、./gcluster job list コマンドの出力に sft のステータスが Succeeded と表示されるまで待ちます。

  3. 元のターミナルで、SFT トレーニングのチェックポイントが Cloud Storage バケットに生成されていることを確認します。

    gcloud storage ls "gs://${GCS_BUCKET:?}/${MODEL_NAME:?}/trained/sft/checkpoints/1000/model_params/"

トレーニング済みモデルを Hugging Face 形式に変換する

  1. トレーニング ワークロードが完了したら、モデルを Hugging Face 形式に戻します。

    ./gcluster job submit --name mt-to-hf \
    --compute-type "${TPU_TYPE:?}" \
    --num-slices 1 \
    --image "${CLOUD_IMAGE_NAME:?}" \
    --await-job-completion \
    --command "[ \"\$JOB_COMPLETION_INDEX\" != \"0\" ] || \
      python3 -m maxtext.checkpoint_conversion.to_huggingface \
        model_name=${MODEL_NAME:?} \
        hf_access_token=${HF_TOKEN:?} \
        load_parameters_path=gs://${GCS_BUCKET:?}/${MODEL_NAME}/trained/sft/checkpoints/1000/model_params/ \
        base_output_directory=gs://${GCS_BUCKET}/${MODEL_NAME}/hf-trained/ \
        skip_jax_distributed_system=true \
        hardware=cpu \
        scan_layers=True \
        use_multimodal=False \
        weight_dtype=bfloat16 \
        --override_model_architecture"

    このコマンドには --await-job-completion フラグが含まれています。これは、変換が完了するまで現在のターミナルで待機することを意味します。

  2. 省略可: ジョブの実行中にライブ変換ログをストリーミングするには、gcluster を抽出した作業ディレクトリで別のターミナルを開き、gcluster job logs コマンドを実行します。

    # Use the list command to check status
./gcluster job list

# Check progress of the job (--main-only targets the coordinator pod (Job Index 0, Pod Index 0) to avoid duplicate logs from other workers)
./gcluster job logs mt-to-hf --main-only -f

# The trained model is now available in gs://${GCS_BUCKET}/${MODEL_NAME}/hf-trained/

  3. 元のターミナルで、最終的な Hugging Face Safetensors の重みが Cloud Storage バケットに生成されていることを確認します。

    gcloud storage ls -l --readable-sizes "gs://${GCS_BUCKET:?}/${MODEL_NAME:?}/hf-trained/"

クリーンアップ

追加料金が発生しないようにするには、gcluster destroy コマンドを使用して、このチュートリアルで作成したリソースを削除します。

./gcluster destroy "${CLUSTER_NAME:?}" --robust --auto-approve
gcloud storage rm -r "gs://${GCS_BUCKET:?}"

# To delete the local deployment folder and staged blueprint
rm -rf .ghpc "${CLUSTER_NAME:?}" tmp

次のステップ

  • Cloud TPU の詳細については、Cloud TPU の概要をご覧ください。
  • v6e-32 TPU のアーキテクチャと構成の詳細については、TPU v6e をご覧ください。