TPU v6e で Gemma 4 の教師ありファインチューニングを実行する

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

目標

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

費用

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

• Google Kubernetes Engine
• Cloud TPU
• Cloud Storage
• Artifact Registry
• Compute Engine

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

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

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

重要: このチュートリアルを完了するには、v6e-32 TPU にアクセスできる必要があります。この チュートリアルでは、必要なときに容量を使用できるように、Google Cloud プロジェクトにこれらの TPU の予約がすでにあることを前提としています。将来の日時に Cloud TPU 容量を予約するオプションの詳細については、 Cloud TPU の予約についてをご覧ください。

始める前に

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

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

• Hugging Face ウェブサイトで、トレーニングするモデルのライセンス 契約に同意します。このチュートリアルでは、モデル gemma4-31b を使用します。

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

• このチュートリアルを完了するには:
  • TPU 管理者 (roles/tpu.admin)
  • サービス アカウント ユーザー （roles/iam.serviceAccountUser）
  • Compute 編集者 （roles/compute.editor）
• MaxText コンテナ イメージを準備するには:
  • Cloud Build 編集者 （roles/cloudbuild.builds.editor）
  • Artifact Registry 管理者 （roles/artifactregistry.admin）
  • ストレージ管理者 （roles/storage.admin）
  • Service Usage 管理者 （roles/serviceusage.serviceUsageAdmin）

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

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

環境変数を設定する

次のスクリプトを実行して、環境変数を設定します。

export PROJECT="YOUR_PROJECT_ID"
export REGION="YOUR_REGION"
export ZONE="YOUR_ZONE"
export CLUSTER_NAME="gke-tpu-v6e"
export REPOSITORY_NAME="YOUR_REPOSITORY_NAME"
export CLOUD_IMAGE_NAME="${REGION}-docker.pkg.dev/${PROJECT}/${REPOSITORY_NAME}/maxtext_base:latest"
export TPU_TYPE="v6e-32"
export RESERVATION="YOUR_RESERVATION_NAME"
export MODEL_NAME="gemma4-31b"
export HF_TOKEN="YOUR_HF_TOKEN"
export GCS_BUCKET="YOUR_BUCKET_NAME"

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

• YOUR_PROJECT_ID: 実際の Google Cloud プロジェクト ID。
• YOUR_REGION: クラスタをデプロイするリージョン。
• YOUR_ZONE: クラスタをデプロイするゾーン。
• YOUR_REPOSITORY_NAME: MaxText イメージの Artifact Registry リポジトリの名前。
• YOUR_RESERVATION_NAME: 予約の名前。
• YOUR_HF_TOKEN: Hugging Face アクセス トークン。
• YOUR_BUCKET_NAME: Cloud Storage バケットのグローバルに一意の名前。

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

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

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

Cluster Toolkit をインストールする

Cluster Toolkit をインストールする の手順に沿って、 Cluster Toolkit の事前ビルド済みバンドルをインストールします。

MaxText コンテナ イメージを準備する

必要な依存関係のインストールなど、MaxText コンテナ イメージを準備する手順は次のとおりです。

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

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

2. Artifact Registry リポジトリを作成します。

   gcloud artifacts repositories create ${REPOSITORY_NAME} \
       --repository-format=docker \
       --location=${REGION} \
       --project=${PROJECT} \
       --description="Docker repository for MaxText images in ${REGION}" || true

3. リポジトリのルート ディレクトリに、ファイル名 cloudbuild.yaml で次の内容のファイルを作成します。

   steps:
     - name: 'gcr.io/cloud-builders/docker'
       entrypoint: 'bash'
       args:
         - '-c'
         - |
           set -euo pipefail
   
           # 0. Install prerequisites (if needed)
           apt-get update && apt-get install -y curl || apk add curl || true
   
           # 1. Install uv
           curl -LsSf https://astral.sh/uv/install.sh | sh
           source $$HOME/.local/bin/env
   
           # 2. Setup Python environment and install MaxText runner
           uv venv --python 3.12 --seed maxtext_venv
           source maxtext_venv/bin/activate
           uv pip install maxtext[runner]==0.2.3 --resolution=lowest
   
           # 3. Build the Docker image (Cloud Build has Docker pre-configured)
           build_maxtext_docker_image WORKFLOW=post-training
   
           # 4. Tag the image properly
           docker tag maxtext_base_image ${_CLOUD_IMAGE_NAME}
   
   # Cloud Build automatically pushes images listed here
   images:
     - '${_CLOUD_IMAGE_NAME}'
   
   options:
     # We use a high-CPU machine to match the n4-standard-16 from the VM tutorial
     machineType: 'E2_HIGHCPU_32'

4. Cloud Build を使用して MaxText Docker イメージをビルドします。

   gcloud builds submit . \
       --project=${PROJECT} \
       --region=${REGION} \
       --substitutions=_CLOUD_IMAGE_NAME="${CLOUD_IMAGE_NAME}"

Cluster Toolkit クラスタを作成する

32 個の v6e TPU チップを使用して Cluster Toolkit クラスタを作成してデプロイする手順は次のとおりです。

1. gke.gcsfuse.profileUser という名前の Identity and Access Management（IAM）カスタムロールを作成します。

   # The GKE TPU v6e blueprint uses GCS Fuse CSI Storage Profiles which requires a custom IAM role.
   # If this role is not already created in your project, you must create it before deploying.
   gcloud iam roles create gke.gcsfuse.profileUser \
     --project=${PROJECT} \
     --title="GKE GCSFuse Profile User" \
     --description="Allows scanning GCS buckets for objects, retrieving bucket metadata, and creating Anywhere Caches." \
     --permissions="storage.objects.list,storage.buckets.get,storage.anywhereCaches.create,storage.anywhereCaches.get,storage.anywhereCaches.list,storage.anywhereCaches.update"
   
   

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

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

3. デフォルトでは、クラスタ ノードプールのサービス アカウントには、Cloud Storage バケットに書き込むために必要な権限がありません。ノードプール サービス アカウントが Cloud Storage バケットに書き込めるようにするには、Storage Admin ロールを付与する必要があります。このロールを付与するには、node_pool_service_account モジュールを更新して、ファイル gke-tpu-v6e-advanced.yaml を編集します。

   - id: node_pool_service_account
     source: modules/project/service-account
     settings:
       name: gke-np-sa
       project_roles:
       - logging.logWriter
       - monitoring.metricWriter
       - monitoring.viewer
       - stackdriver.resourceMetadata.writer
       - storage.admin            # Change from storage.objectViewer
       - artifactregistry.reader

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

   ./gcluster deploy examples/gke-tpu-v6e/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 \
       -w

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

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

1. Cluster Toolkit クラスタの作成が完了したら、Docker を構成します。

   # Configure docker for pulling images
   gcloud auth configure-docker gcr.io --quiet
   gcloud auth configure-docker ${REGION}-docker.pkg.dev --quiet

2. 以降のコマンドを簡略化するため、デフォルトのプロジェクト、クラスタ、ロケーションを構成します。

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

3. モデルを Hugging Face 形式から MaxText 形式に変換して Cloud Storage バケットに保存するには、次のスクリプトを実行します。

   ./gcluster job submit --name hf-to-mt \
       --cluster ${CLUSTER_NAME} \
       --project ${PROJECT} \
       --location ${REGION} \
       --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"

変換ジョブのステータスを確認するには、次のコマンドを実行します。

# Use the list command to check status
./gcluster job list \
    --cluster ${CLUSTER_NAME} \
    --project ${PROJECT} \
    --location ${REGION}

# 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 \
    --cluster ${CLUSTER_NAME} \
    --project ${PROJECT} \
    --location ${REGION}

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

変換プロセスが完了したら、次のコマンドを実行して SFT ワークロードを開始できます。

./gcluster job submit --name sft \
    --cluster ${CLUSTER_NAME} \
    --project ${PROJECT} \
    --location ${REGION} \
    --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"

トレーニング ジョブのステータスを確認するには、次のコマンドを実行します。

# Use the list command to check status
./gcluster job list \
    --cluster ${CLUSTER_NAME} \
    --project ${PROJECT} \
    --location ${REGION}

# 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 \
    --cluster ${CLUSTER_NAME} \
    --project ${PROJECT} \
    --location ${REGION}

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

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

./gcluster job submit --name mt-to-hf \
    --cluster ${CLUSTER_NAME} \
    --project ${PROJECT} \
    --location ${REGION} \
    --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"

変換ジョブのステータスを確認するには、次のコマンドを実行します。

# Use the list command to check status
./gcluster job list \
    --cluster ${CLUSTER_NAME} \
    --project ${PROJECT} \
    --location ${REGION}

# 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 \
    --cluster ${CLUSTER_NAME} \
    --project ${PROJECT} \
    --location ${REGION}

# The trained model is now available in gs://${GCS_BUCKET}/${MODEL_NAME}/hf-trained/ - though again, it's ~2x the size of the original...

クリーンアップ

追加の料金が発生しないように、このチュートリアルで作成したリソースを削除します。

gcluster destroy ${CLUSTER_NAME} --robust
gcloud storage rm -r gs://${GCS_BUCKET}
gcloud artifacts repositories delete ${REPOSITORY_NAME} --location=${REGION} --project=${PROJECT} --quiet

# To delete the local deployment folder
rm -rf .ghpc ${CLUSTER_NAME}

次のステップ

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