A4 GKE クラスタで Gemma 4 をファインチューニングする

クラスタ化された GPU

このチュートリアルでは、 Google Cloud上のシングルホスト、マルチ GPU Google Kubernetes Engine（GKE）クラスタで Gemma 4 大規模言語モデル（LLM）をファインチューニングする方法について説明します。このクラスタは、8 個の NVIDIA B200 GPU が割り当てられた単一の A4 仮想マシン（VM）インスタンスを使用します。

このチュートリアルで説明する主なプロセスは次の 2 つです。

1. GKE Autopilot を使用して、高パフォーマンスの GKE クラスタをデプロイします。このデプロイの一環として、必要なソフトウェアがプリインストールされたカスタム VM イメージを作成します。
2. クラスタがデプロイされたら、このチュートリアルに付属するスクリプトのセットを使用して、分散型ファインチューニング ジョブを実行します。このジョブは、Hugging Face Accelerate ライブラリを活用します。

重要: このチュートリアルを完了するには、A4 VM を作成する容量を予約しておく必要があります。将来の日時に AI Hypercomputer で容量を予約するオプションの詳細については、使用オプションを選択するをご覧ください。

このチュートリアルは、LLM のトレーニングに Google Cloud で GKE クラスタをデプロイすることに関心のある ML エンジニア、研究者、プラットフォーム管理者、オペレーター、データおよび AI スペシャリストを対象としています。

目標

• Hugging Face を使用して Gemma 4 モデルにアクセスします。

• 環境を準備します。

• A4 GKE クラスタを作成してデプロイする。

• 完全にシャーディングされたデータ並列処理（FSDP）で Hugging Face Accelerate ライブラリを使用して、Gemma 4 モデルをファインチューニングします。

• ジョブをモニタリングします。

• クリーンアップする。

費用

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

• Compute Engine
• Google Kubernetes Engine (GKE)
• Filestore
• Cloud Storage
• Cloud Logging

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

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

始める前に

1. 必要な API が有効になっていない場合は、有効にします。

   API を有効にするために必要なロール

   API を有効にするには、serviceusage.services.enable 権限が必要です。プロジェクトを作成した場合は、オーナーロール（roles/owner）を通じてこの権限がすでに付与されている可能性があります。それ以外の場合は、Service Usage 管理者ロール（roles/serviceusage.serviceUsageAdmin）を通じてこの権限を取得できます。ロールを付与する方法を確認する。

   gcloud services enable compute.googleapis.com container.googleapis.com file.googleapis.com logging.googleapis.com cloudresourcemanager.googleapis.com servicenetworking.googleapis.com

2. Google Cloud プロジェクトのデフォルトのサービス アカウントを有効にします。

   export PROJECT_NUMBER="$(gcloud projects describe "YOUR_PROJECT_ID" --format "value(project_number)")"
   gcloud iam service-accounts enable "${PROJECT_NUMBER}[e-mail]" \
       --project=YOUR_PROJECT_ID

3. デフォルトのサービス アカウントがファインチューニング ワークロードを実行するために必要なロールを付与します。

   ROLES=("roles/aiplatform.user"
          "roles/container.developer"
          "roles/storage.objectUser")
   for role in "${ROLES[@]}"; do
     gcloud projects add-iam-policy-binding YOUR_PROJECT_ID \
         --member="serviceAccount:${PROJECT_NUMBER}[e-mail]" \
         --role="${role}"
   done
   unset ROLES

   注: この手順では、このチュートリアルで必要な事前定義ロールのみが付与されます。編集者（roles/editor）などの基本ロールを付与すると、サービス アカウントに必要以上の権限が付与され、組織のポリシーによって付与がブロックされる可能性があります。詳細については、サービス アカウントの使用に関するベスト プラクティスをご覧ください。

4. ユーザー アカウントのローカル認証情報を作成します。

   gcloud auth application-default login

   注: ローカルシェルと外部 ID プロバイダ（IdP）を使用しているときに、上記のコマンドを実行した後に認証エラーが発生した場合は、フェデレーション ID を使用して gcloud CLI にログインします。

5. プロジェクトで OS Login を有効にします。

   gcloud compute project-info add-metadata \
       --metadata=enable-oslogin=TRUE \
       --project=YOUR_PROJECT_ID

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

• Kubernetes Engine 管理者（roles/container.admin）
• Compute 管理者（roles/compute.admin）
• ストレージ管理者 （roles/storage.admin）
• Artifact Registry 管理者（roles/artifactregistry.admin）
• Cloud Build 編集者 （roles/cloudbuild.builds.editor）
• サービス アカウント ユーザー（roles/iam.serviceAccountUser）
• サービス アカウント管理者 （roles/iam.serviceAccountAdmin）
• プロジェクト IAM 管理者 （roles/resourcemanager.projectIamAdmin）
• Service Usage 管理者（roles/serviceusage.serviceUsageAdmin）

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

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

Hugging Face を使用して Gemma 4 にアクセスする

Hugging Face を使用して Gemma 4 にアクセスする手順は次のとおりです。

1. Hugging Face にログインする
2. Hugging Face write アクセス トークンを作成します。
   [Your Profile] > [Settings] > [Access tokens] > [+Create new token] の順にクリックします。
3. write access トークンの値をコピーして保存します。これは、このチュートリアルの後半で使用します。

環境を準備する

環境を準備するには、次の設定を行います。

export PROJECT_ID="YOUR_PROJECT_ID"
export CLUSTER_NAME="YOUR_CLUSTER_NAME"
export CLUSTER_REGION="YOUR_REGION"
export RESERVATION="YOUR_RESERVATION_NAME"
export HF_TOKEN="YOUR_HF_TOKEN"
export ARTIFACT_REPO_LOCATION="YOUR_ARTIFACT_REGISTRY_LOCATION"

gcloud config set project "${PROJECT_ID}"
gcloud config set billing/quota_project "${PROJECT_ID}"

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

• YOUR_PROJECT_ID: GKE クラスタを作成する Google Cloud プロジェクトの ID。

• YOUR_CLUSTER_NAME: 作成する GKE クラスタの名前。

• YOUR_REGION: GKE クラスタを作成するリージョン。クラスタは、予約が存在するリージョンでのみ作成できます。

• YOUR_RESERVATION_NAME: 予約済み容量の識別子。

• YOUR_HF_TOKEN: 前のセクションで作成した Hugging Face アクセス トークン。

• YOUR_ARTIFACT_REGISTRY_LOCATION: Artifact Registry リポジトリを作成する Google Cloud リージョン。イメージの pull レイテンシを最小限に抑えるには、YOUR_REGION で使用したリージョンと同じリージョンを使用することをおすすめします。

Autopilot モードの GKE クラスタを作成する

Autopilot モードで GKE クラスタを作成するには、次のコマンドを実行します。

gcloud container clusters create-auto "${CLUSTER_NAME}" \
    --project="${PROJECT_ID}" \
    --location="${CLUSTER_REGION}"

GKE クラスタの作成には時間がかかることがあります。 Google Cloud がクラスタの作成を完了したことを確認するには、 Google Cloud コンソールの [Kubernetes クラスタ] に移動します。

Hugging Face の認証情報用の Kubernetes Secret を作成する

Hugging Face の認証情報用の Kubernetes Secret を作成する手順は次のとおりです。

1. GKE クラスタと通信するように kubectl を構成します。

   gcloud container clusters get-credentials "${CLUSTER_NAME}" \
       --location="${CLUSTER_REGION}"

2. Hugging Face トークンを保存する Kubernetes Secret を作成します。

   kubectl create secret generic hf-secret \
       --from-literal=hf_api_token="${HF_TOKEN}" \
       --dry-run=client -o yaml | kubectl apply -f -

ワークロードを準備する

ワークロードを準備する手順は次のとおりです。

1. ワークロード スクリプトを作成する。

2. Docker と Cloud Build を使用してファインチューニング コンテナを作成します。

ワークロード スクリプトを作成する

ファインチューニング ワークロードで使用するスクリプトを作成するには、次の操作を行います。

1. ワークロード スクリプト用のディレクトリを作成します。このディレクトリを作業ディレクトリとして使用します。

   mkdir llm-finetuning-gemma
   cd llm-finetuning-gemma

2. Google Cloud Build を使用する cloudbuild.yaml ファイルを作成します。このファイルは、ワークロード コンテナを作成して Artifact Registry に保存します。

   steps:
   - name: 'gcr.io/cloud-builders/docker'
     args:
     - 'build'
     - '-t'
     - '$_ARTIFACT_REPO_LOCATION-docker.pkg.dev/$PROJECT_ID/gemma/finetune-gemma-gpu:2.0.0'
     - '.'
   images:
   - '$_ARTIFACT_REPO_LOCATION-docker.pkg.dev/$PROJECT_ID/gemma/finetune-gemma-gpu:2.0.0'
   options:
     logging: CLOUD_LOGGING_ONLY

3. ファインチューニング ジョブを実行する Dockerfile ファイルを作成します。

   FROM nvidia/cuda:12.8.1-cudnn-devel-ubuntu24.04
   RUN apt-get update && \
       apt-get -y install python3 python3-dev gcc python3-pip \
           python3-venv git curl vim && \
       rm -rf /var/lib/apt/lists/*
   RUN python3 -m venv /opt/venv
   ENV PATH="/opt/venv/bin:/usr/local/nvidia/bin:$PATH"
   ENV LD_LIBRARY_PATH="/usr/local/nvidia/lib64:$LD_LIBRARY_PATH"
   RUN pip3 install setuptools wheel packaging ninja
   RUN pip3 install torch torchvision torchaudio \
       --index-url https://download.pytorch.org/whl/cu128
   RUN pip3 install \
       "transformers>=4.53.3" \
       "datasets>=4.0.0" \
       "accelerate>=1.9.0" \
       evaluate==0.4.5 \
       bitsandbytes==0.46.1 \
       "trl>=0.19.1" \
       "peft>=0.16.0" \
       tensorboard==2.20.0 \
       protobuf==6.31.1 \
       sentencepiece==0.2.0
   COPY finetune.py /finetune.py
   COPY accel_fsdp_gemma4_config.yaml /accel_fsdp_gemma4_config.yaml
   CMD accelerate launch --config_file accel_fsdp_gemma4_config.yaml finetune.py

   注: このチュートリアルでは、NVIDIA と Pytorch の依存関係の最新バージョンを使用していません。新しい依存関係が必要な場合は、NVIDIA のドキュメントと Pytorch のドキュメントをご覧ください。

4. accel_fsdp_gemma4_config.yaml ファイルを作成します。この構成ファイルは、FSDP を使用して、単一ホスト上の 8 個のローカル GPU にチューニング ジョブを分割するように Hugging Face Accelerate に指示します。

   compute_environment: LOCAL_MACHINE
   debug: false
   distributed_type: FSDP
   downcast_bf16: 'no'
   enable_cpu_affinity: false
   fsdp_config:
     fsdp_activation_checkpointing: false
     fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
     fsdp_cpu_ram_efficient_loading: false
     fsdp_offload_params: false
     fsdp_reshard_after_forward: true
     fsdp_state_dict_type: FULL_STATE_DICT
     fsdp_transformer_layer_cls_to_wrap: Gemma4TextDecoderLayer
     fsdp_version: 2
   machine_rank: 0
   main_training_function: main
   mixed_precision: bf16
   num_machines: 1
   num_processes: 8
   rdzv_backend: static
   same_network: true
   tpu_env: []
   tpu_use_cluster: false
   tpu_use_sudo: false
   use_cpu: false

5. finetune.yaml ファイルを作成します。

   apiVersion: batch/v1
   kind: Job
   metadata:
     name: finetune-job
     namespace: default
   spec:
     backoffLimit: 2
     template:
       metadata:
         annotations:
           kubectl.kubernetes.io/default-container: finetuner
       spec:
         terminationGracePeriodSeconds: 600
         containers:
         - name: finetuner
           image: $IMAGE_URL
           command: ["accelerate","launch"]
           args:
           - "--config_file"
           - "accel_fsdp_gemma4_config.yaml"
           - "finetune.py"
           - "--model_id"
           - "google/gemma-4-31B-it"
           - "--output_dir"
           - "gemma-31b-text-to-sql"
           - "--per_device_train_batch_size"
           - "2"
           - "--gradient_accumulation_steps"
           - "32"
           - "--num_train_epochs"
           - "3"
           - "--learning_rate"
           - "1e-5"
           - "--save_strategy"
           - "steps"
           - "--save_steps"
           - "20"
           - "--push_to_hub"
           resources:
             limits:
               nvidia.com/gpu: "8"
               memory: "1200Gi"
               ephemeral-storage: "350Gi"
             requests:
               nvidia.com/gpu: "8"
               memory: "1200Gi"
               ephemeral-storage: "350Gi"
           env:
           - name: HF_TOKEN
             valueFrom:
               secretKeyRef:
                 name: hf-secret
                 key: hf_api_token
           volumeMounts:
           - mountPath: /dev/shm
             name: dshm
         volumes:
         - name: dshm
           emptyDir:
             medium: Memory
             sizeLimit: 64Gi
         nodeSelector:
           cloud.google.com/gke-accelerator: nvidia-b200
           cloud.google.com/reservation-name: $RESERVATION
           cloud.google.com/reservation-affinity: "specific"
           cloud.google.com/gke-gpu-driver-version: latest
         restartPolicy: OnFailure

6. finetune.py ファイルを作成します。

   import torch
   import argparse
   import subprocess
   from datasets import load_dataset
   from transformers import (
       AutoTokenizer,
       AutoModelForCausalLM,
       BitsAndBytesConfig,
       AutoConfig,
   )
   from peft import LoraConfig
   from trl import SFTTrainer, SFTConfig
   from huggingface_hub import login
   
   def get_args():
       parser = argparse.ArgumentParser()
       parser.add_argument(
           "--model_id", type=str, default="google/gemma-4-31b-it",
           help="Hugging Face model ID")
       parser.add_argument(
           "--hf_token", type=str, default=None,
           help="Hugging Face token for private models")
       parser.add_argument(
           "--trust_remote", type=lambda x: str(x).lower() == "true",
           default=False,
           help="Trust remote code when loading tokenizer")
       parser.add_argument(
           "--use_fast", type=lambda x: str(x).lower() == "true",
           default=True,
           help="Determines if a fast Rust-based tokenizer should be used")
       parser.add_argument(
           "--dataset_name", type=str,
           default="philschmid/gretel-synthetic-text-to-sql",
           help="Hugging Face dataset name")
       parser.add_argument(
           "--output_dir", type=str, default="gemma-31b-text-to-sql",
           help="Directory to save model checkpoints")
   
       # LoRA arguments
       parser.add_argument(
           "--lora_r", type=int, default=16,
           help="LoRA attention dimension")
       parser.add_argument(
           "--lora_alpha", type=int, default=32,
           help="LoRA alpha scaling factor")
       parser.add_argument(
           "--lora_dropout", type=float, default=0.05,
           help="LoRA dropout probability")
       # SFTConfig arguments
       parser.add_argument(
           "--max_length", type=int, default=1024,
           help="Maximum sequence length")
       parser.add_argument(
           "--num_train_epochs", type=int, default=3,
           help="Number of training epochs")
       parser.add_argument(
           "--per_device_train_batch_size", type=int, default=2,
           help="Batch size per device during training")
       parser.add_argument(
           "--gradient_accumulation_steps", type=int, default=32,
           help="Gradient accumulation steps")
       parser.add_argument(
           "--learning_rate", type=float, default=1e-5,
           help="Learning rate")
       parser.add_argument(
           "--logging_steps", type=int, default=10,
           help="Log every X steps")
       parser.add_argument(
           "--save_strategy", type=str, default="steps",
           help="Checkpoint save strategy")
       parser.add_argument(
           "--save_steps", type=int, default=100,
           help="Save checkpoint every X steps")
       parser.add_argument(
           "--push_to_hub", action='store_true',
           help="Push model back up to HF")
       parser.add_argument(
           "--hub_private_repo", type=lambda x: str(x).lower() == "true",
           default=True,
           help="Push to a private repo")
       return parser.parse_args()
   
   def main():
       args = get_args()
       # --- 1. Setup and Login ---
       if args.hf_token:
           login(args.hf_token)
       # --- 2. Create and prepare the fine-tuning dataset ---
       dataset = load_dataset(args.dataset_name, split="train")
       dataset = dataset.shuffle().select(range(12500))
       dataset = dataset.train_test_split(test_size=2500/12500)
       # --- 3. Configure Model and Tokenizer ---
       if torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 8:
           torch_dtype_obj = torch.bfloat16
           torch_dtype_str = "bfloat16"
       else:
           torch_dtype_obj = torch.float16
           torch_dtype_str = "float16"
       tokenizer = AutoTokenizer.from_pretrained(
           args.model_id,
           trust_remote_code=args.trust_remote,
           use_fast=args.use_fast,
       )
       if tokenizer.pad_token is None:
           tokenizer.pad_token = tokenizer.eos_token
       # --- 4. Define the Formatting Function ---
       def formatting_func(example):
           system_message = (
               "You are a text to SQL query translator. Users will ask you "
               "questions in English and you will generate a SQL query based "
               "on the provided SCHEMA."
           )
           user_prompt = (
               "Given the <USER_QUERY> and the <SCHEMA>, generate the "
               "corresponding SQL command to retrieve the desired data, "
               "considering the query's syntax, semantics, and schema "
               "constraints.\n\n<SCHEMA>\n{context}\n</SCHEMA>\n\n"
               "<USER_QUERY>\n{question}\n</USER_QUERY>\n"
           )
   
           messages = [
               {"role": "system", "content": system_message},
               {
                   "role": "user",
                   "content": user_prompt.format(
                       question=example["sql_prompt"],
                       context=example["sql_context"],
                   ),
               },
               {"role": "assistant", "content": example["sql"]}
           ]
           return tokenizer.apply_chat_template(messages, tokenize=False)
       # --- 5. Load Model and Apply PEFT ---
       config = AutoConfig.from_pretrained(args.model_id)
       config.use_cache = False
       print("Loading base model...")
       model = AutoModelForCausalLM.from_pretrained(
           args.model_id,
           config=config,
           attn_implementation="sdpa",
           torch_dtype=torch_dtype_obj,
       )
   
       peft_config = LoraConfig(
           lora_alpha=args.lora_alpha,
           lora_dropout=args.lora_dropout,
           r=args.lora_r,
           bias="none",
           # Train only the text decoder projections; skip vision/audio modules.
           target_modules=[
               "q_proj", "k_proj", "v_proj", "o_proj",
               "gate_proj", "up_proj", "down_proj",
           ],
           exclude_modules=(
               r".*(vision_tower|audio_tower|embed_vision|embed_audio).*"),
           task_type="CAUSAL_LM",
       )
       # --- 6. Configure Training Arguments ---
       training_args = SFTConfig(
           output_dir=args.output_dir,
           max_length=args.max_length,
           num_train_epochs=args.num_train_epochs,
           per_device_train_batch_size=args.per_device_train_batch_size,
           gradient_accumulation_steps=args.gradient_accumulation_steps,
           learning_rate=args.learning_rate,
           logging_steps=args.logging_steps,
           save_strategy=args.save_strategy,
           save_steps=args.save_steps,
           packing=False,
           label_names=["domain"],
           gradient_checkpointing=True,
           gradient_checkpointing_kwargs={"use_reentrant": False},
           optim="adamw_torch",
           fp16=True if torch_dtype_obj == torch.float16 else False,
           bf16=True if torch_dtype_obj == torch.bfloat16 else False,
           max_grad_norm=0.3,
           warmup_steps=0.03,
           lr_scheduler_type="constant",
           push_to_hub=args.push_to_hub,
           hub_private_repo=args.hub_private_repo,
           report_to="tensorboard",
       )
       # --- 7. Create Trainer and Start Training ---
       trainer = SFTTrainer(
           model=model,
           args=training_args,
           peft_config=peft_config,
           train_dataset=dataset["train"],
           eval_dataset=dataset["test"],
           processing_class=tokenizer,
           formatting_func=formatting_func,
       )
       print("Starting training...")
       trainer.train()
       print("Training finished.")
       # --- 8. Save the final model ---
       print(f"Saving final model to {args.output_dir}")
       if trainer.is_fsdp_enabled:
           trainer.accelerator.state.fsdp_plugin.set_state_dict_type(
               "FULL_STATE_DICT")
       trainer.save_model(args.output_dir)
       if torch.distributed.is_initialized():
           torch.distributed.destroy_process_group()
   
   if __name__ == "__main__":
       main()

Docker と Cloud Build を使用してファインチューニング コンテナを作成する

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

   gcloud artifacts repositories create gemma \
       --repository-format=docker \
       --location="${ARTIFACT_REPO_LOCATION}" \
       --description="Repository for Gemma fine tuning workload containers" || true

2. 前の手順で作成した llm-finetuning-gemma ディレクトリで、次のコマンドを実行してファインチューニング イメージを作成し、Artifact Registry に push します。

   gcloud builds submit . \
       --substitutions=_ARTIFACT_REPO_LOCATION="${ARTIFACT_REPO_LOCATION}"

3. 画像の URL をエクスポートします。この値は、このチュートリアルの後半の手順で使用します。

   IMAGE_REGISTRY="${ARTIFACT_REPO_LOCATION}-docker.pkg.dev/${PROJECT_ID}"
   export IMAGE_URL="${IMAGE_REGISTRY}/gemma/finetune-gemma-gpu:2.0.0"

ファインチューニング ワークロードを開始する

ファインチューニング ワークロードを開始するには、次の操作を行います。

1. ファインチューニング マニフェストを適用して、ファインチューニング ジョブを作成します。

   envsubst < finetune.yaml | kubectl apply -f -

   GKE Autopilot モードでクラスタを使用しているため、GPU 対応ノードの起動には数分かかることがあります。

2. 次のコマンドを実行して、ジョブをモニタリングします。

   watch kubectl get pods

3. Pod が実行されたら、ジョブのログを確認します。

   kubectl logs job.batch/finetune-job -f

   Job リソースはモデルデータをダウンロードし、8 つすべての GPU を使用してモデルをファインチューニングします。ダウンロードには 5 分ほどかかります。ダウンロードが完了すると、ファインチューニング プロセスが完了するまでに約 2 時間 30 分かかります。

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

GKE クラスタでの GPU の使用状況をモニタリングして、ファインチューニング ジョブが効率的に実行されていることを確認できます。これを行うには、ブラウザで次のリンクを開きます。

echo "https://console.cloud.google.com/kubernetes/clusters/details/${CLUSTER_REGION}/${CLUSTER_NAME}/observability?mods=monitoring_api_prod&project=${PROJECT_ID}&pageState=(\"timeRange\":(\"duration\":\"PT1H\"),\"nav\":(\"section\":\"gpu\"),\"groupBy\":(\"groupByType\":\"namespacesTop5\"))"

ワークロードをモニタリングすると、次の情報を確認できます。

• GPU の使用率: 正常なファインチューニング ジョブでは、8 個の GPU すべての使用率が上昇し、トレーニング全体を通して高いレベルで安定することが予想されます。
• ジョブの所要時間: 指定された A4 クラスタでジョブが完了するまでに約 2 時間 30 分かかります。

クリーンアップ

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

リソースを削除する

1. ファインチューニング ジョブを削除するには、次のコマンドを実行します。

   kubectl delete job finetune-job

2. GKE クラスタを削除するには、次のコマンドを実行します。

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

次のステップ

• AI 最適化 GKE クラスタを管理する

• AI / ML ワークロードにおすすめのストレージ サービスについて学習する