在 A4 GKE 集群上微调 Gemma 3

本教程介绍如何在 Google Cloud上的多节点、多 GPU GKE 集群中微调 Gemma 3 大语言模型 (LLM)。此集群使用包含 8 个 NVIDIA B200 GPU 的 A4 虚拟机 (VM) 实例。

本教程中介绍的两个主要流程如下：

  1. 使用 GKE Autopilot 部署高性能 GKE 集群。在此部署过程中，您将创建一个预先安装了必要软件的自定义虚拟机映像。
  2. 部署集群后，您可以使用本教程随附的一组脚本来运行分布式微调作业。该作业利用了 Hugging Face Accelerate 库。

本教程适用于机器学习 (ML) 工程师、研究人员、平台管理员和运维人员，以及对在 Google Cloud 上部署 GKE 集群以训练 LLM 感兴趣的数据和 AI 专家。

目标

  • 使用 Hugging Face 访问 Gemma 3 模型。

  • 准备环境。

  • 创建并部署 A4 GKE 集群。

  • 使用 Hugging Face Accelerate 库和完全分片数据并行处理 (FSDP) 来微调 Gemma 3 模型。

  • 监控作业。

  • 清理。

费用

在本文档中，您将使用 Google Cloud的以下收费组件：

如需根据您的预计使用情况来估算费用，请使用价格计算器。

新 Google Cloud 用户可能有资格申请免费试用。

准备工作

  1. 安装 Google Cloud CLI。

  2. 配置 gcloud CLI 以使用您的联合身份。

    如需了解详情，请参阅使用联合身份登录 gcloud CLI。

  3. 如需初始化 gcloud CLI，请运行以下命令： 

    gcloud init

  4. 创建或选择 Google Cloud 项目。

    选择或创建项目所需的角色

    • 选择项目：选择项目不需要特定的 IAM 角色，您可以选择已获授角色的任何项目。
    • 创建项目：如需创建项目，您需要拥有 Project Creator 角色 (roles/resourcemanager.projectCreator)，该角色包含 resourcemanager.projects.create 权限。了解如何授予角色。

    • 创建 Google Cloud 项目：

      gcloud projects create PROJECT_ID

      将 PROJECT_ID 替换为您要创建的 Google Cloud 项目的名称。

    • 选择您创建的 Google Cloud 项目：

      gcloud config set project PROJECT_ID

      将 PROJECT_ID 替换为您的 Google Cloud 项目名称。

  5. 验证是否已为您的 Google Cloud 项目启用结算功能。

  6. 启用所需的 API：

    启用 API 所需的角色

    如需启用 API，您需要拥有 serviceusage.services.enable 权限。如果您创建了项目，则可能已经通过 Owner 角色 (roles/owner) 获得了此权限。否则，您可以通过 Service Usage Admin 角色 (roles/serviceusage.serviceUsageAdmin) 获得此权限。了解如何授予角色。 

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

  7. 向您的用户账号授予角色。对以下每个 IAM 角色运行以下命令一次： roles/compute.admin, roles/iam.serviceAccountUser, roles/cloudbuild.builds.editor, roles/artifactregistry.admin, roles/storage.admin, roles/serviceusage.serviceUsageAdmin 

    gcloud projects add-iam-policy-binding PROJECT_ID --member="user:USER_IDENTIFIER" --role=ROLE

    替换以下内容：

    • PROJECT_ID：您的项目 ID。
    • USER_IDENTIFIER：您的用户 个账号。如需查看示例，请参阅 在 IAM 政策中表示员工池用户。
    • ROLE：您向用户账号授予的 IAM 角色。
  8. 为您的 Google Cloud 项目启用默认服务账号： 
      export PROJECT_NUMBER="$(gcloud projects describe "PROJECT_ID" --format "value(project_number)")"
  gcloud iam service-accounts enable "${PROJECT_NUMBER}-compute@ \
    --project=PROJECT_ID
  9. 向默认服务账号授予 Editor 角色 (roles/editor)： 
    gcloud projects add-iam-policy-binding PROJECT_ID \
    --member="serviceAccount:PROJECT_NUMBER-compute@" \
    --role=roles/editor
  10. 为您的用户账号创建本地身份验证凭据： 
    gcloud auth application-default login
  11. 为项目启用 OS Login： 
    gcloud compute project-info add-metadata --metadata=enable-oslogin=TRUE
  12. 登录或创建 Hugging Face 账号。

使用 Hugging Face 访问 Gemma 3

如需使用 Hugging Face 访问 Gemma 3，请执行以下操作：

  1. 登录 Hugging Face
  2. 创建 Hugging Face write 访问令牌。
    依次点击您的个人资料 > 设置 > 访问令牌 > +创建新令牌
  3. 复制并保存 write access 令牌值。您将在本教程的后面部分使用该地址。

准备环境

如需准备环境，请设置以下内容：

export PROJECT_ID="YOUR_PROJECT_ID"
export CLUSTER_NAME="YOUR_CLUSTER_NAME"
export CLUSTER_REGION="YOUR_REGION"
export RESERVATION="RESERVATION_NAME"
export HF_TOKEN="YOUR_HF_TOKEN"
export ARTIFACT_REPO_LOCATION="ARTIFACT_REGISTRY_LOCATION"
export NETWORK="default"

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

替换以下内容：

  • PROJECT_ID：您要在其中创建 GKE 集群的 Google Cloud 项目 的名称。

  • RESERVATION：预留容量的标识符。

  • CLUSTER_NAME：要创建的 GKE 集群的名称。

  • CLUSTER_REGION：您要在其中创建 GKE 集群的区域。您只能在预留所在的区域中创建集群。

  • HF_TOKEN：您在上一部分中创建的 Hugging Face 访问令牌。

  • ARTIFACT_REPO_LOCATION：您要在其中创建 Artifact Registry 代码库的位置。

在 Autopilot 模式下创建 GKE 集群

如需在 Autopilot 模式下创建 GKE 集群，请运行以下命令：

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

创建 GKE 集群可能需要一些时间才能完成。如需验证 Google Cloud 是否已完成集群创建，请前往 Google Cloud 控制台中的 Kubernetes 集群。

为 Hugging Face 凭据创建 Kubernetes Secret

如需为 Hugging Face 凭据创建 Kubernetes Secret，请按以下步骤操作：

  1. 配置 kubectl 以与您的 GKE 集群通信：

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

  2. 创建一个 Kubernetes Secret 来存储您的 Hugging Face 令牌：

    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. 创建 cloudbuild.yaml 文件以使用 Google Cloud Build。此文件会创建工作负载容器并将其存储在 Artifact Registry 中：

    steps:
- name: 'gcr.io/cloud-builders/docker'
  args: [ 'build', '-t', '$_ARTIFACT_REPO_LOCATION-docker.pkg.dev/$PROJECT_ID/gemma/finetune-gemma-gpu:1.0.0', '.' ]
images:
- '$_ARTIFACT_REPO_LOCATION-docker.pkg.dev/$PROJECT_ID/gemma/finetune-gemma-gpu:1.0.0'

  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
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_gemma3_config.yaml /accel_fsdp_gemma3_config.yaml
CMD accelerate launch --config_file accel_fsdp_gemma3_config.yaml finetune.py

  4. 创建 accel_fsdp_gemma3_config.yaml 文件。此配置文件指示 Hugging Face Accelerate 将调节作业拆分到多个 GPU 上。

    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: true
  fsdp_offload_params: false
  fsdp_reshard_after_forward: true
  fsdp_state_dict_type: FULL_STATE_DICT
  fsdp_transformer_layer_cls_to_wrap: Gemma3DecoderLayer
  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_gemma3_config.yaml"
        - "finetune.py"
        - "--model_id"
        - "google/gemma-3-12b-pt"
        - "--output_dir"
        - "gemma-12b-text-to-sql"
        - "--per_device_train_batch_size"
        - "8"
        - "--gradient_accumulation_steps"
        - "8"
        - "--num_train_epochs"
        - "3"
        - "--learning_rate"
        - "1e-5"
        - "--save_strategy"
        - "steps"
        - "--save_steps"
        - "100"
        resources:
          limits:
            nvidia.com/gpu: "8"
        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
      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-3-12b-pt", 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-12b-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=16, 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_seq_length", type=int, default=512, 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=8, help="Batch size per device during training")
    parser.add_argument("--gradient_accumulation_steps", type=int, default=1, 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)
    tokenizer.pad_token = tokenizer.eos_token
    gemma_chat_template = (
        "{{ bos_token }}"
        "{% if messages[0]['role'] == 'system' %}{{ messages[0]['content'] }}{% endif %}"
        "{% for message in messages %}"
        "{% if message['role'] == 'user' %}<start_of_turn>user\n{{ message['content'] }}<end_of_turn>\n{% elif message['role'] == 'assistant' %}<start_of_turn>model\n{{ message['content'] }}<end_of_turn>\n{% endif %}"
        "{% endfor %}"
    )
    tokenizer.chat_template = gemma_chat_template
    # --- 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": "user", "content": user_prompt.format(question=example["sql_prompt"][0], context=example["sql_context"][0])},
            {"role": "assistant", "content": example["sql"][0]}
        ]
        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="eager",
        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=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
        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_seq_length=args.max_seq_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_ratio=0.03,
        lr_scheduler_type="constant",
        push_to_hub=True,
        report_to="tensorboard",
        dataset_kwargs={
            "add_special_tokens": False,
            "append_concat_token": True,
        }
    )
    # --- 7. Create Trainer and Start Training ---
    trainer = SFTTrainer(
        model=model,
        args=training_args,
        train_dataset=dataset["train"],
        eval_dataset=dataset["test"],
        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}")
    model.cpu()
    trainer.save_model(args.output_dir)
    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。

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

  3. 导出图片网址。您将在本教程的后续步骤中使用它：

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

启动微调工作负载

如需启动微调工作负载，请执行以下操作：

  1. 应用微调清单以创建微调作业：

    envsubst < finetune.yaml | kubectl apply -f -

    由于您使用的是 GKE Autopilot 模式下的集群，因此启动启用 GPU 的节点可能需要几分钟时间。

  2. 通过运行以下命令来监控作业：

    watch kubectl get pods

  3. 通过运行以下命令来检查作业的日志：

    kubectl logs job.batch/finetune-job -f

    作业资源会下载模型数据，然后在所有 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 的用量在整个训练过程中上升并稳定在高水平。
  • 作业时长：作业大约需要 10 分钟才能在指定的 A4 集群上完成。

清理

为避免因本教程中使用的资源导致您的 Google Cloud 账号产生费用，请删除包含这些资源的项目，或者保留项目但删除各个资源。

删除您的资源

  1. 如需删除微调作业，请运行以下命令：

    kubectl delete job finetune-job

  2. 如需删除 GKE 集群，请运行以下命令：

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

删除项目

删除 Google Cloud 项目：

gcloud projects delete PROJECT_ID

后续步骤