Google 会使用 AI 技术将内容翻译成您偏好的语言。AI 翻译可能包含错误。

在 A4 GKE 集群上微调 Gemma 4

集群 GPU

本教程介绍如何在 Google Cloud上的单主机、多 GPU Google Kubernetes Engine (GKE) 集群中微调 Gemma 4 大语言模型 (LLM)。此集群使用单个 A4 虚拟机 (VM) 实例，该实例挂接了 8 个 NVIDIA B200 GPU。

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

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

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

目标

使用 Hugging Face 访问 Gemma 4 模型。
准备环境。
创建并部署 A4 GKE 集群。
使用 Hugging Face Accelerate 库和完全分片数据并行处理 (FSDP) 来微调 Gemma 4 模型。
监控作业。
清理。

费用

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

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

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

准备工作

启用所需的 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
```

为您的 Google Cloud 项目启用默认服务账号：

export PROJECT_NUMBER="$(gcloud projects describe "YOUR_PROJECT_ID" --format "value(project_number)")"
gcloud iam service-accounts enable "${PROJECT_NUMBER}-compute@developer.gserviceaccount.com" \
    --project=YOUR_PROJECT_ID

授予默认服务账号运行微调工作负载所需的角色：

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}-compute@developer.gserviceaccount.com" \
      --role="${role}"
done
unset ROLES

为您的用户账号创建本地身份验证凭据：
```
gcloud auth application-default login
```
注意：如果您使用的是本地 shell 和外部身份提供方 (IdP)，并且在运行上述命令后遇到身份验证错误，请使用联合身份登录 gcloud CLI。

为项目启用 OS Login：

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

如需获得完成本教程所需的权限，请让您的管理员为您授予项目的以下 IAM 角色：

Kubernetes Engine Admin (roles/container.admin)
Compute Admin (roles/compute.admin)
Storage Admin (roles/storage.admin)
Artifact Registry Administrator (roles/artifactregistry.admin)
Cloud Build Editor (roles/cloudbuild.builds.editor)
Service Account User (roles/iam.serviceAccountUser)
Service Account Admin (roles/iam.serviceAccountAdmin)
Project IAM Admin (roles/resourcemanager.projectIamAdmin)
Service Usage Admin (roles/serviceusage.serviceUsageAdmin)

如需详细了解如何授予角色，请参阅管理对项目、文件夹和组织的访问权限。

您也可以通过自定义角色或其他预定义角色来获取所需的权限。

使用 Hugging Face 访问 Gemma 4

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

登录 Hugging Face
创建 Hugging Face write 访问令牌。
依次点击您的个人资料 > 设置 > 访问令牌 > +创建新令牌
复制并保存 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 区域。为最大限度缩短映像拉取延迟时间，建议您使用与 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，请按以下步骤操作：

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

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

创建一个 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 -

准备工作负载

如需准备工作负载，请执行以下操作：

创建工作负载脚本。
使用 Docker 和 Cloud Build 创建微调容器。

创建工作负载脚本

如需创建微调工作负载使用的脚本，请执行以下操作：

为工作负载脚本创建目录。将此目录用作工作目录。
```
mkdir llm-finetuning-gemma
cd llm-finetuning-gemma
```

创建 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:2.0.0'
  - '.'
images:
- '$_ARTIFACT_REPO_LOCATION-docker.pkg.dev/$PROJECT_ID/gemma/finetune-gemma-gpu:2.0.0'
options:
  logging: CLOUD_LOGGING_ONLY

创建一个 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

创建 accel_fsdp_gemma4_config.yaml 文件。此配置文件指示 Hugging Face Accelerate 使用 FSDP 将调优作业拆分到单个宿主机上的八个本地 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: 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

创建 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

创建 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 创建微调容器

创建 Artifact Registry Docker 制品库：

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

在您之前步骤中创建的 llm-finetuning-gemma 目录中，运行以下命令以创建微调映像并将其推送到 Artifact Registry。
```
gcloud builds submit . \
    --substitutions=_ARTIFACT_REPO_LOCATION="${ARTIFACT_REPO_LOCATION}"
```

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

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

启动微调工作负载

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

应用微调清单以创建微调作业：
```
envsubst < finetune.yaml | kubectl apply -f -
```
由于您使用的是 GKE Autopilot 模式下的集群，因此启动启用 GPU 的节点可能需要几分钟时间。
运行以下命令来监控作业：
```
watch kubectl get pods
```
Pod 运行后，检查作业的日志：
```
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 的用量在整个训练过程中上升并稳定在高水平。
作业时长：作业在指定的 A4 集群上大约需要 2 小时 30 分钟才能完成。

清理

为避免产生额外费用，请删除在本教程中创建的资源。

删除您的资源

如需删除微调作业，请运行以下命令：
```
kubectl delete job finetune-job
```

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

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