在 A4 GKE 集群上微调 Gemma 4

本教程介绍如何在 Gemma 4 Gemma 4 大语言模型 (LLM) Google Cloud。此集群使用单个 A4 虚拟机 (VM) 实例，并挂接了 8 个 NVIDIA B200 GPU。

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

  1. 使用 GKE Autopilot 部署高性能 GKE 集群。在此部署过程中，您将使用 Cloud Build 构建自定义容器映像，并预安装必要的软件。
  2. 部署集群后，您将运行一个微调作业，该作业使用完全分片数据并行 (FSDP) 在单个 A4 (a4-highgpu-8g) 节点的所有 8 个 GPU 上对模型进行分片。该作业使用本教程附带的一组脚本，并利用 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 用户可能有资格申请免费试用。

准备工作

  1. 启用所需的 API（如果尚未启用）：

    启用 API 所需的角色

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

    gcloud services enable compute.googleapis.com container.googleapis.com artifactregistry.googleapis.com cloudbuild.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}-compute@developer.gserviceaccount.com" \
    --project=YOUR_PROJECT_ID

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

    ROLES=("roles/artifactregistry.writer"
       "roles/cloudbuild.builds.builder"
       "roles/logging.logWriter"
       "roles/monitoring.metricWriter"
       "roles/monitoring.viewer"
       "roles/stackdriver.resourceMetadata.writer"
       "roles/storage.objectViewer")
for role in "${ROLES[@]}"; do
  gcloud projects add-iam-policy-binding YOUR_PROJECT_ID \
      --member="serviceAccount:${PROJECT_NUMBER}-compute@developer.gserviceaccount.com" \
      --role="${role}" 1>/dev/null
done
unset ROLES

    如需验证是否已授予角色，请运行以下命令：

    gcloud projects get-iam-policy YOUR_PROJECT_ID \
    --flatten="bindings[].members" \
    --filter="bindings.members:serviceAccount:${PROJECT_NUMBER}-compute@developer.gserviceaccount.com" \
    --format="table(bindings.role)"

  4. 为您的用户账号创建本地身份验证凭证：

    gcloud auth application-default login

  5. 为您的项目启用 OS Login：

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

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

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

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

使用 Hugging Face 访问 Gemma 4

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

  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="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，请按以下步骤操作：

  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: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 文件。此配置文件指示 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

  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。

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

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

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

启动微调工作负载

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

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

    envsubst < finetune.yaml | kubectl apply -f -

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

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

    watch kubectl get pods

  3. Pod 运行后，检查作业的日志：

    kubectl logs job.batch/finetune-job -f

    作业资源会下载模型数据，然后在所有 8 个 GPU 上对模型进行微调。下载大约需要 5 分钟才能完成。 下载完成后，微调过程大约需要 45 分钟才能完成。该作业每 20 步保存一个检查点，并将微调的低秩适应 (LoRA) 适配器推送到名为 gemma-31b-text-to-sql 的私有 Hugging Face 仓库。

  4. 作业完成后，获取 Hugging Face 上微调模型的网址，然后在浏览器中打开该网址：

    HF_USERNAME=$(curl -s -H "Authorization: Bearer ${HF_TOKEN}" \
    https://huggingface.co/api/whoami-v2 | jq -r .name)
echo "https://huggingface.co/${HF_USERNAME}/gemma-31b-text-to-sql"

监控工作负载

您可以监控 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 的使用率上升并稳定在高水平 。
  • 作业时长：该作业应在大约 45 分钟内在指定的 A4 集群上 完成。

清理

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

删除您的资源

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

    kubectl delete job finetune-job

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

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

  3. 如需删除 Artifact Registry 代码库，请运行以下命令：

    gcloud artifacts repositories delete gemma \
    --location="${ARTIFACT_REPO_LOCATION}"

后续步骤