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

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"
export NETWORK="default"

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

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

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

• 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

4. accel_fsdp_gemma4_config.yaml ファイルを作成します。この構成ファイルは、チューニング ジョブを複数の 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_backward_prefetch: BACKWARD_PRE
     fsdp_cpu_ram_efficient_loading: true
     fsdp_forward_prefetch: true
     fsdp_offload_params: false
     fsdp_reshard_after_forward: true
     fsdp_sharding_strategy: FULL_SHARD
     fsdp_state_dict_type: FULL_STATE_DICT
     fsdp_sync_module_states: true
     fsdp_transformer_layer_cls_to_wrap: Gemma4TextDecoderLayer
     fsdp_version: 1
   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"
           - "100"
           resources:
             limits:
               nvidia.com/gpu: "8"
               memory: "1000Gi"
               ephemeral-storage: "350Gi"
             requests:
               nvidia.com/gpu: "8"
               memory: "1000Gi"
               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, prepare_model_for_kbit_training, get_peft_model
   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=bool, default="False",
           help="Trust remote code when loading tokenizer")
       parser.add_argument(
           "--use_fast", type=bool, 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=bool, 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,
       )
   
       model = prepare_model_for_kbit_training(model)
       peft_config = LoraConfig(
           lora_alpha=args.lora_alpha,
           lora_dropout=args.lora_dropout,
           r=args.lora_r,
           bias="none",
           target_modules="all-linear",
           task_type="CAUSAL_LM",
       )
       print("Applying PEFT configuration...")
       model = get_peft_model(model, peft_config)
       model.print_trainable_parameters()
       # --- 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,
           report_to="tensorboard",
       )
       # --- 7. Create Trainer and Start Training ---
       trainer = SFTTrainer(
           model=model,
           args=training_args,
           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}")
       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 分かかります。