MaxText を使用して Qwen3-14b モデルでマルチホスト教師ありファインチューニングを実行する

このチュートリアルでは、Cloud TPU で MaxText を使用して Qwen3-14b モデルで教師ありファインチューニング（SFT）を実行する手順を説明します。このチュートリアルでは、特殊なコンテナ イメージをビルドする方法、Accelerated Processing Kit（XPK）を使用して Pathways で Google Kubernetes Engine（GKE）クラスタをプロビジョニングする方法、マルチホスト トレーニング ワークロードを実行する方法について説明します。

目標

• トレーニング後の最適化されたカスタム MaxText コンテナ イメージをビルドする方法について学習します。
• Pathways を有効にして XPK を使用して GKE クラスタをプロビジョニングする。
• Qwen3 14b モデルを Hugging Face 形式から MaxText 形式に変換します。
• Cloud TPU でマルチホスト SFT トレーニング ワークロードを実行します。
• サービング用にファインチューニングされたモデルを Hugging Face 形式に戻します。

費用

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

• Google Kubernetes Engine
• Cloud TPU
• Cloud Storage
• Artifact Registry
• Compute Engine, for the temporary build virtual machine (VM)

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

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

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

始める前に

• ユーザー アカウントまたはサービス アカウントに次のロールがあることを確認します。
  • roles/compute.admin: ビルド VM を作成します。
  • roles/artifactregistry.admin: Docker リポジトリを管理する
  • roles/storage.admin: データバケットを管理する
  • roles/container.admin: Google Kubernetes Engine クラスタの作成と管理
  • roles/iam.serviceAccountAdmin: ワークロード サービス アカウントを作成する
  • roles/resourcemanager.projectIamAdmin: Identity and Access Management（IAM）ポリシーを設定する
  • roles/iam.serviceAccountUser: サービス アカウントとして機能する
• Google Cloud CLI をインストールして初期化します。
• ワークステーションに Python 3.12 以降がインストールされていることを確認します。

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

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

環境を設定する

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

export PROJECT="YOUR_PROJECT_ID"
export REGION="YOUR_REGION"
export ZONE="YOUR_ZONE"
export CLUSTER_NAME="YOUR_CLUSTER_NAME"
export GCS_BUCKET="YOUR_GCS_BUCKET"
export CLOUD_IMAGE_NAME="$REGION-docker.pkg.dev/$PROJECT/maxtext-images/maxtext_base:latest"
export TPU_TYPE="v6e-32"
export CLUSTER_NODEPOOL_COUNT=1
export PW_CPU_MACHINE_TYPE="c4d-standard-96"
export RESERVATION="YOUR_RESERVATION_NAME"
export MODEL_NAME="qwen3-14b"
export HF_TOKEN="YOUR_HF_TOKEN"

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

• YOUR_PROJECT_ID: 実際の Google Cloud プロジェクト ID
• YOUR_REGION: 使用するリージョン
• YOUR_ZONE: 使用するゾーン
• YOUR_CLUSTER_NAME: Google Kubernetes Engine クラスタの名前
• YOUR_GCS_BUCKET: Cloud Storage バケットの一意の名前
• YOUR_RESERVATION_NAME: 容量予約
• YOUR_HF_TOKEN: Hugging Face アクセス トークン

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

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

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

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

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

   gcloud artifacts repositories create maxtext-images \
       --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.1 --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}"

Google Kubernetes Engine クラスタを作成する

Qwen3 14b モデルで SFT トレーニングを実行するには、TPU チップを搭載した Google Kubernetes Engine クラスタが必要です。Accelerated Processing Kit（XPK）をインストールし、Pathways をサポートする GKE クラスタを作成します。

# Start with creating a new virtual environment to install XPK in.
VENV_DIR=venvp3
python3 -m venv $VENV_DIR
source $VENV_DIR/bin/activate
pip install xpk==1.14.0

xpk cluster create-pathways \
  --num-slices=${CLUSTER_NODEPOOL_COUNT} \
  --tpu-type=${TPU_TYPE} \
  --pathways-gce-machine-type=${PW_CPU_MACHINE_TYPE} \
  --project=${PROJECT} \
  --zone=${ZONE} \
  --cluster=${CLUSTER_NAME} \
  --custom-cluster-arguments="--enable-ip-alias" \
  --reservation=$RESERVATION \
  --default-pool-cpu-machine-type=n4-standard-16

gcloud container clusters get-credentials $CLUSTER_NAME \
  --location=$REGION \
  --project $PROJECT

トレーニング用にモデルを準備する

CPU ベースのワークロードを使用して、ベースモデルを MaxText 形式に変換します。このタスクを複数のマシンで並行して実行しないでください。次のコマンドには、変換が 1 つの TPU ノードでのみ実行されるようにするためのチェックが含まれています。

xpk workload create \
  --workload "qwen-hf-to-mt" \
  --docker-image $CLOUD_IMAGE_NAME \
  --cluster ${CLUSTER_NAME} \
  --tpu-type=${TPU_TYPE} \
  --num-slices=1 \
  --project=${PROJECT} \
  --zone=${ZONE} \
  --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}/qwen-3-14b/max-text-format/ \
  scan_layers=True \
  use_multimodal=False \
  skip_jax_distributed_system=true \
  hardware=cpu \
  --lazy_load_tensors=True"

モデル変換の進行状況を追跡する

変換の進行状況を追跡する手順は次のとおりです。

1. GKE クラスタでスケジュールされた Pod を一覧表示するには、kubectl get pod コマンドを実行します。
2. qwen-hf-to-mt-slice-job-0-0-HASH という名前の Pod を見つけます。
3. Pod の出力をリアルタイムで検査するには、kubectl logs -f POD_NAME コマンドを実行します。

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

変換プロセスが完了したら、XPK を使用して SFT ファインチューニング ワークロードを開始できます。

xpk workload create-pathways \
  --cluster=${CLUSTER_NAME} \
  --project=${PROJECT} \
  --zone=${ZONE} \
  --docker-image=$CLOUD_IMAGE_NAME \
  --workload="qwen-training" \
  --tpu-type=${TPU_TYPE} \
  --num-slices=1 \
  --command="JAX_PLATFORMS=proxy JAX_BACKEND_TARGET=grpc://127.0.0.1:29000 ENABLE_PATHWAYS_PERSISTENCE=1 \
  python3 -m maxtext.trainers.post_train.sft.train_sft \
  run_name=sft \
  base_output_directory=gs://${GCS_BUCKET}/qwen-3-14b/trained/ \
  model_name=${MODEL_NAME} \
  load_parameters_path=gs://${GCS_BUCKET}/qwen-3-14b/max-text-format/0/items/ \
  hf_access_token=${HF_TOKEN} \
  per_device_batch_size=1 \
  steps=1000 \
  profiler=xplane \
  checkpoint_storage_use_zarr3=0 \
  checkpoint_storage_use_ocdbt=0 \
  enable_single_controller=True"

トレーニング ワークロードをモニタリングする

XPK コマンドライン インターフェース（CLI）を使用して、ワークロードのステータスをモニタリングします。

xpk workload list --cluster ${CLUSTER_NAME} --project ${PROJECT} --zone ${ZONE}

ログと TPU 使用率を表示するには、Google Cloud コンソールを使用します。次のコマンドを実行してログを表示することもできます。

kubectl logs -f qwen-training-pathways-head-0-0-HASH

HASH は、Pod 名の数値ハッシュに置き換えます。このハッシュの値を確認するには、kubectl get pod コマンドを実行して、返された Pod のリストを確認します。

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

トレーニング ワークロードが完了したら、チェックポイントを Hugging Face 形式に戻します。

xpk workload create \
  --cluster=${CLUSTER_NAME} \
  --project=${PROJECT} \
  --zone=${ZONE} \
  --docker-image=$CLOUD_IMAGE_NAME \
  --workload="qwen-mt-to-hf" \
  --tpu-type=${TPU_TYPE} \
  --num-slices=1 \
  --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}/qwen-3-14b/trained/sft/checkpoints/1000/model_params/ \
  base_output_directory=gs://$GCS_BUCKET/qwen-3-14b/hf-trained/ \
  skip_jax_distributed_system=true \
  hardware=cpu \
  scan_layers=True \
  use_multimodal=False \
  weight_dtype=bfloat16"

変換の進行状況を追跡するには、kubectl logs -f qwen-mt-to-hf-slice-job-0-0-HASH コマンドを実行します。HASH は、Pod 名の数値ハッシュに置き換えます。

変換が完了すると、gs://$GCS_BUCKET/qwen-3-14b/hf-trained/ に保存されているチューニング済みモデルを使用できるようになります。

クリーンアップ

追加料金が発生しないように、このチュートリアルで作成したリソース（Google Kubernetes Engine クラスタ、Cloud Storage バケット、Artifact Registry リポジトリなど）を削除します。

注意: このチュートリアルでトレーニングしたモデルを今後も使用する場合は、そのモデルを含む Cloud Storage バケットを削除する前に、アクセス可能なバージョンを保存してください。

このチュートリアル用に作成したリソースを削除するには、次のコマンドを実行します。

xpk cluster delete --cluster $CLUSTER_NAME --project $PROJECT --zone $ZONE --force

gcloud storage rm --recursive gs://$GCS_BUCKET

gcloud artifacts repositories delete maxtext-images --location=$REGION --project=$PROJECT --quiet

次のステップ

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