在多主机 A4 GKE 集群上微调 Gemma 4

使用 Hugging Face 访问 Gemma 4

如需使用 Hugging Face 访问 Gemma 4，请完成以下步骤：

1. 登录 Hugging Face 并 接受 Gemma 4 许可协议。
2. 创建 Hugging Face write 访问令牌。
   依次点击您的个人资料 > 设置 > 访问令牌 > +创建新令牌。
3. 复制并保存 write 访问令牌值。您可以使用此令牌下载基础模型，并在 GKE 缩容 GPU 节点之前将微调的适配器检查点推送到 Hugging Face Hub。

准备环境

如需准备环境，请设置以下环境变量：

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 NUM_NODES="YOUR_NUMBER_OF_NODES"

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 集群 的名称。此名称不得超过 40 个字符，并且只能包含小写字母数字字符和连字符。

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

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

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

• YOUR_ARTIFACT_REGISTRY_LOCATION：您要在其中创建 Artifact Registry 代码库的 Google Cloud 区域 （例如 us-central1）。为最大限度地缩短映像提取延迟时间，请使用您为 YOUR_REGION 指定的同一区域。

• YOUR_NUMBER_OF_NODES：微调作业中的 A4 虚拟机节点数。对于此多主机教程，其中包含两个 a4-highgpu-8g 实例中的 16 个 NVIDIA B200 GPU，请将此值设置为 2。

在 Autopilot 模式下创建多主机 GKE 集群

在 Autopilot 模式下创建多主机 GKE 集群：

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

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

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

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

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

为 Hugging Face 凭据创建 Kubernetes Secret

创建一个 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 文件，以使用 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-multihost-gpu:2.0.0'
     - '.'
   images:
   - '$_ARTIFACT_REPO_LOCATION-docker.pkg.dev/$PROJECT_ID/gemma/finetune-gemma-multihost-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>=5.5.0" \
       trl==0.29.1 \
       peft==0.18.1 \
       accelerate==1.13.0 \
       bitsandbytes==0.49.2 \
       datasets==4.8.4 \
       evaluate==0.4.5 \
       tensorboard==2.20.0 \
       protobuf==6.31.1 \
       sentencepiece==0.2.0
   WORKDIR /workspace
   COPY finetune.py /workspace/finetune.py
   COPY accel_fsdp_gemma4_config.yaml /workspace/accel_fsdp_gemma4_config.yaml
   CMD ["accelerate", "launch", "--config_file", "/workspace/accel_fsdp_gemma4_config.yaml", "/workspace/finetune.py"]

4. 创建 accel_fsdp_gemma4_config.yaml 文件。此配置指示 Hugging Face Accelerate 使用 FSDP v2 将 Gemma4TextDecoderLayer 分片到两个主机上的 16 个 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: true
     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: 2
   num_processes: 16
   rdzv_backend: static
   same_network: true
   tpu_env: []
   tpu_use_cluster: false
   tpu_use_sudo: false
   use_cpu: false

5. 创建 finetune.yaml Kubernetes JobSet 清单：

   apiVersion: resource.k8s.io/v1
   kind: ResourceClaimTemplate
   metadata:
     name: mrdma
   spec:
     spec:
       devices:
         requests:
         - name: mrdma
           exactly:
             deviceClassName: mrdma.google.com
   ---
   apiVersion: jobset.x-k8s.io/v1alpha2
   kind: JobSet
   metadata:
     name: finetune-jobset
     namespace: default
   spec:
     failurePolicy:
       maxRestarts: 2
     replicatedJobs:
     - name: workers
       replicas: 1
       template:
         spec:
           parallelism: ${NUM_NODES}
           completions: ${NUM_NODES}
           backoffLimit: 0
           template:
             metadata:
               annotations:
                 kubectl.kubernetes.io/default-container: finetuner
             spec:
               terminationGracePeriodSeconds: 600
               restartPolicy: OnFailure
               nodeSelector:
                 cloud.google.com/compute-class: "Accelerator"
                 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
               containers:
               - name: finetuner
                 image: $IMAGE_URL
                 command: ["bash", "-c"]
                 args:
                 - |
                   NUM_PROCESSES=$(( ${NUM_NODES} * 8 ))
                   accelerate launch \
                     --config_file /workspace/accel_fsdp_gemma4_config.yaml \
                     --num_machines ${NUM_NODES} \
                     --num_processes ${NUM_PROCESSES} \
                     --machine_rank ${JOB_COMPLETION_INDEX} \
                     --main_process_ip finetune-jobset-workers-0-0.finetune-jobset.default.svc.cluster.local \
                     --main_process_port 29500 \
                     /workspace/finetune.py \
                     --base_model google/gemma-4-31b-it \
                     --new_model gemma-31b-text-to-sql \
                     --per_device_train_batch_size 4 \
                     --gradient_accumulation_steps 4 \
                     --num_train_epochs 3 \
                     --learning_rate 1e-5 \
                     --save_strategy steps \
                     --save_steps 15 \
                     --push_to_hub
                 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
                 - name: NUM_NODES
                   value: "${NUM_NODES}"
                 volumeMounts:
                 - mountPath: /dev/shm
                   name: dshm
               volumes:
               - name: dshm
                 emptyDir:
                   medium: Memory
                   sizeLimit: 64Gi

6. 创建 finetune.py 监督式微调脚本：

   import argparse
   import torch
   from datasets import load_dataset
   from huggingface_hub import login
   from peft import LoraConfig
   from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
   from trl import SFTConfig, SFTTrainer
   
   
   def get_args():
       def str2bool(v):
           if isinstance(v, bool):
               return v
           return v.lower() in ("yes", "true", "t", "1")
   
       parser = argparse.ArgumentParser()
       parser.add_argument(
           "--base_model",
           "--model_id",
           dest="base_model",
           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 gated models and Hub uploads",
       )
       parser.add_argument(
           "--trust_remote",
           type=str2bool,
           default=False,
           help="Trust remote code when loading tokenizer",
       )
       parser.add_argument(
           "--use_fast",
           type=str2bool,
           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(
           "--new_model",
           "--output_dir",
           dest="new_model",
           type=str,
           default="gemma-31b-text-to-sql",
           help="Directory and repository name 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=4,
           help="Batch size per device during training",
       )
       parser.add_argument(
           "--gradient_accumulation_steps",
           type=int,
           default=4,
           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=15,
           help="Save checkpoint every X steps",
       )
       parser.add_argument(
           "--push_to_hub",
           action="store_true",
           help="Push model back up to Hugging Face Hub",
       )
       parser.add_argument(
           "--hub_private_repo",
           type=str2bool,
           default=True,
           help="Push to a private repository on Hugging Face Hub",
       )
       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
       else:
           torch_dtype_obj = torch.float16
       tokenizer = AutoTokenizer.from_pretrained(
           args.base_model,
           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 Configure LoRA ---
       config = AutoConfig.from_pretrained(args.base_model)
       config.use_cache = False
       print("Loading base model...")
       model = AutoModelForCausalLM.from_pretrained(
           args.base_model,
           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",
           target_modules=[
               "q_proj",
               "k_proj",
               "v_proj",
               "o_proj",
               "gate_proj",
               "up_proj",
               "down_proj",
           ],
           exclude_modules=r".*(vision_tower|embed_vision|audio_tower|embed_audio).*",
           task_type="CAUSAL_LM",
       )
       # --- 6. Configure Training Arguments ---
       training_args = SFTConfig(
           output_dir=args.new_model,
           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=torch_dtype_obj == torch.float16,
           bf16=torch_dtype_obj == torch.bfloat16,
           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.new_model}")
       if trainer.is_fsdp_enabled:
           trainer.accelerator.state.fsdp_plugin.set_state_dict_type(
               "FULL_STATE_DICT"
           )
       trainer.save_model(args.new_model)
       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. 安装编排多主机工作负载所需的 JobSet 自定义资源定义 (CRD)：

   kubectl apply --server-side \
       -f https://github.com/kubernetes-sigs/jobset/releases/download/v0.12.0/manifests.yaml
   
   # Wait for JobSet controller webhook to become ready to prevent intermittent failures during deploy
   kubectl rollout status deployment/jobset-controller-manager -n jobset-system --timeout=600s

3. 在您在上一步中创建的 llm-finetuning-gemma 目录中，将容器构建提交到 Cloud Build：

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

4. 导出多主机容器映像网址。在本教程的后续步骤中部署 JobSet 清单时，您将使用该网址：

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

启动微调工作负载

如需部署和监控分布式微调工作负载，请完成以下步骤：

1. 将环境变量替换为微调清单，以创建微调作业：

   envsubst '${RESERVATION} ${IMAGE_URL} ${NUM_NODES}' < finetune.yaml \
       | kubectl apply -f -

   由于您的集群在 GKE Autopilot 模式下运行，因此预配两个启用 GPU 的 A4 节点并提取容器映像可能需要几分钟时间。

2. 观察工作器 Pod，直到两个 Pod 都转换为 Running 状态：

   watch kubectl get pods

3. 工作器 Pod 转换为 Running 后，流式传输训练日志：

   kubectl logs -l "job-name=finetune-jobset-workers-0" -f

监控工作负载

您可以监控 GKE 集群中的 GPU 利用率，以验证两个 A4 主机中的所有 16 个 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 利用率：对于运行正常的分布式微调作业，您应该会在训练步骤中看到所有 16 个 NVIDIA B200 GPU 的 GPU 利用率上升并稳定在 95%–100% 附近。
• 作业时长：在两个 a4-highgpu-8g 节点（16 个 B200 GPU）上， 3 个周期的微调作业大约需要 2 个半小时才能完成。

查看微调的适配器权重

训练完成后，在 Hugging Face Hub 上查看微调的 LoRA 适配器权重和 检查点，网址为 https://huggingface.co/YOUR_HF_USERNAME/gemma-31b-text-to-sql。