使用 Ray 在 GKE 上微调 Gemma 3 以执行视觉任务

使用 Hugging Face 访问 Gemma 3

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

1. 签署同意协议，以使用 Gemma 3。

2. 创建 Hugging Face read access 令牌。

3. 复制并保存 read access 令牌值。您将在本教程的后面部分用到它。

准备环境

通过配置必要的设置和设置环境变量来准备环境。

运行以下命令：

export PROJECT_ID="YOUR_PROJECT_ID"
export RESERVATION="YOUR_RESERVATION_NAME"
export REGION="YOUR_REGION"
export CLUSTER_NAME="YOUR_CLUSTER_NAME"
export HF_TOKEN="YOUR_HF_TOKEN"
export NETWORK="default"
export RAY_SA="YOUR_RAY_SA"
export GSA_NAME="YOUR_GSA_NAME"
export GCS_BUCKET="YOUR_GCS_BUCKET"

gcloud config set project $PROJECT_ID
gcloud config set billing/quota_project $PROJECT_ID

替换以下内容：

• YOUR_PROJECT_ID：您要在其中创建 GKE 集群的 Google Cloud 项目 的名称。
• YOUR_RESERVATION_NAME：您要用于创建集群的预留的网址。根据预留所在的项目，指定以下值之一：
  • 预留存在于您的项目中：指定预留名称（例如 my-reservation）。
  • 预留存在于其他项目中：请指定完整路径，格式为 projects/RESERVATION_PROJECT_ID/reservations/RESERVATION_NAME。
• YOUR_REGION：您要在其中创建 GKE 集群的区域。 您只能在预留所在的区域中创建集群。
• YOUR_CLUSTER_NAME：要创建的 GKE 集群的名称。
• YOUR_HF_TOKEN：您在之前的步骤中创建的 Hugging Face 令牌。
• YOUR_RAY_SA：Kubernetes 集群内的服务账号名称。
• YOUR_GSA_NAME：Google 服务账号的名称。
• YOUR_GCS_BUCKET：用于存储训练检查点结果的存储桶的名称。

在 Autopilot 模式下创建 GKE 集群

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

gcloud container clusters create-auto $CLUSTER_NAME \
    --enable-ray-operator \
    --enable-ray-cluster-monitoring \
    --enable-ray-cluster-logging \
    --location=$REGION

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

为 Hugging Face 凭据创建 Kubernetes Secret

在 Cloud Shell 中，如需为 Hugging Face 凭据创建 Kubernetes Secret，请执行以下操作：

1. 配置 kubectl 以连接到您的集群：

   gcloud container clusters get-credentials $CLUSTER_NAME \
       --region=$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 -

创建 Cloud Storage 存储桶

如果您想使用新存储桶来存储训练工件，请运行以下命令：

gcloud storage buckets create gs://$GCS_BUCKET --location=$REGION

如果您想使用现有存储桶，可以跳过此步骤。不过，您必须确保存储桶与集群位于同一区域。

创建 IAM 服务账号

在 Cloud Shell 中，如需创建 IAM 服务账号（也称为 Google 服务账号或 GSA）并授予其访问 Cloud Storage 存储桶的权限，请执行以下操作：

1. 创建 IAM 服务账号

   gcloud iam service-accounts create $GSA_NAME \
       --project=$PROJECT_ID \
       --description="Account for Ray to save Gemma training results" \
       --display-name="Ray Gemma Training SA"

2. 向您的 IAM 服务账号授予对 Cloud Storage 存储桶的 Storage Admin 角色 (roles/storage.admin)：

   gcloud storage buckets add-iam-policy-binding gs://${GCS_BUCKET#gs://} \
       --member="serviceAccount:$GSA_NAME@$PROJECT_ID.iam.gserviceaccount.com" \
       --role="roles/storage.admin"

创建 Kubernetes 服务账号

在 Cloud Shell 中，创建 Kubernetes 服务账号并配置 Workload Identity，以向 Ray pod 授予对 Google Cloud 资源的访问权限：

1. 创建 Kubernetes 服务账号：

   kubectl create serviceaccount $RAY_SA \
       --dry-run=client -o yaml | kubectl apply -f -

2. 将您的 Kubernetes 服务账号绑定到您之前创建的 IAM 服务账号，以启用 Workload Identity：

   gcloud iam service-accounts add-iam-policy-binding $GSA_NAME@$PROJECT_ID.iam.gserviceaccount.com \
       --project=$PROJECT_ID \
       --role="roles/iam.workloadIdentityUser" \
       --member="serviceAccount:$PROJECT_ID.svc.id.goog[default/$RAY_SA]"

3. 使用 IAM 服务账号的电子邮件地址为 Kubernetes 服务账号添加注解：

   kubectl annotate serviceaccount $RAY_SA iam.gke.io/gcp-service-account=$GSA_NAME@$PROJECT_ID.iam.gserviceaccount.com --overwrite

将训练代码保存为 ConfigMap

为避免将训练脚本嵌入到容器映像中，您可以将其存储为集群中的 ConfigMap。此 ConfigMap 会装载到 Pod 文件系统，这样一来，您无需重新创建整个 Ray 集群即可更新训练脚本。

如需将训练脚本存储为集群中的 ConfigMap，请完成以下步骤：

1. 创建一个名为 code 的目录，然后在该目录中创建一个名为 vision_train.py 的文件。

2. 将以下代码复制到 vision_train.py 文件中：

   import argparse
   import datetime
   import ray
   import ray.train.huggingface.transformers
   import torch
   from PIL import Image
   from datasets import load_dataset
   from peft import LoraConfig
   from ray.train import ScalingConfig, RunConfig
   from ray.train.torch import TorchTrainer
   from transformers import AutoProcessor, AutoModelForImageTextToText, BitsAndBytesConfig
   from trl import SFTConfig
   from trl import SFTTrainer
   
   # System message for the assistant
   system_message = "You are an expert product description writer for Amazon."
   
   # User prompt that combines the user query and the schema
   user_prompt = """Create a Short Product description based on the provided <PRODUCT> and <CATEGORY> and image.
   Only return description. The description should be SEO optimized and for a better mobile search experience.
   
   <PRODUCT>
   {product}
   </PRODUCT>
   
   <CATEGORY>
   {category}
   </CATEGORY>
   """
   
   def get_args():
       parser = argparse.ArgumentParser()
       parser.add_argument("--model_id", type=str, default="google/gemma-3-4b-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("--dataset_name", type=str, default="philschmid/amazon-product-descriptions-vlm", help="Hugging Face dataset name")
       parser.add_argument("--output_dir", type=str, default="gemma-3-4b-seo-optimized", help="Directory to save model checkpoints")
       parser.add_argument("--gcs_bucket", type=str, required=True, help="storage bucket name used to synchronize tasks and save checkpoints")
       parser.add_argument("--push_to_hub", help="Push model to Hugging Face hub", action="store_true")
   
       # 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=1, 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=2e-4, 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="epoch", help="Checkpoint save strategy")
       parser.add_argument("--save_steps", type=int, default=100, help="Save checkpoint every X steps")
   
       return parser.parse_args()
   
   # Convert dataset to OAI messages
   def format_data(sample):
       return {
           "messages": [
               {
                   "role": "system",
                   "content": [{"type": "text", "text": system_message}],
               },
               {
                   "role": "user",
                   "content": [
                       {
                           "type": "text",
                           "text": user_prompt.format(
                               product=sample["Product Name"],
                               category=sample["Category"],
                           ),
                       },
                       {
                           "type": "image",
                           "image": sample["image"],
                       },
                   ],
               },
               {
                   "role": "assistant",
                   "content": [{"type": "text", "text": sample["description"]}],
               },
           ],
       }
   
   def process_vision_info(messages: list[dict]) -> list[Image.Image]:
       image_inputs = []
       # Iterate through each conversation
       for msg in messages:
           # Get content (ensure it's a list)
           content = msg.get("content", [])
           if not isinstance(content, list):
               content = [content]
   
           # Check each content element for images
           for element in content:
               if isinstance(element, dict) and ("image" in element or element.get("type") == "image"):
                   # Get the image and convert to RGB
                   if "image" in element:
                       image = element["image"]
                   else:
                       image = element
                   image_inputs.append(image.convert("RGB"))
       return image_inputs
   
   def train(args):
       # Load dataset from the hub
       dataset = load_dataset(args.dataset_name, split="train")
   
       # Convert dataset to OAI messages
       # need to use list comprehension to keep Pil.Image type, .mape convert image to bytes
       dataset = [format_data(sample) for sample in dataset]
   
       # Hugging Face model id
       model_id = args.model_id
   
       # Check if GPU benefits from bfloat16
       if torch.cuda.get_device_capability()[0] < 8:
           raise ValueError("GPU does not support bfloat16, please use a GPU that supports bfloat16.")
   
       # Define model init arguments
       model_kwargs = dict(
           attn_implementation="eager",  # Use "flash_attention_2" when running on Ampere or newer GPU
           torch_dtype=torch.bfloat16,  # What torch dtype to use, defaults to auto
           # device_map="auto",  # Let torch decide how to load the model
       )
   
       # BitsAndBytesConfig int-4 config
       model_kwargs["quantization_config"] = BitsAndBytesConfig(
           load_in_4bit=True,
           bnb_4bit_use_double_quant=True,
           bnb_4bit_quant_type="nf4",
           bnb_4bit_compute_dtype=model_kwargs["torch_dtype"],
           bnb_4bit_quant_storage=model_kwargs["torch_dtype"],
       )
   
       # Load model and tokenizer
       model = AutoModelForImageTextToText.from_pretrained(model_id, **model_kwargs)
       processor = AutoProcessor.from_pretrained(model_id, use_fast=True)
   
       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",
           modules_to_save=[
               "lm_head",
               "embed_tokens",
           ],
       )
   
       args = SFTConfig(
           output_dir=args.output_dir,  # directory to save and repository id
           num_train_epochs=args.num_train_epochs,  # number of training epochs
           per_device_train_batch_size=args.per_device_train_batch_size,  # batch size per device during training
           gradient_accumulation_steps=args.gradient_accumulation_steps,  # number of steps before performing a backward/update pass
           gradient_checkpointing=True,  # use gradient checkpointing to save memory
           optim="adamw_torch_fused",  # use fused adamw optimizer
           logging_steps=args.logging_steps,  # log every N steps
           save_strategy=args.save_strategy,  # save checkpoint every epoch
           learning_rate=args.learning_rate,  # learning rate, based on QLoRA paper
           bf16=True,  # use bfloat16 precision
           max_grad_norm=0.3,  # max gradient norm based on QLoRA paper
           warmup_ratio=0.03,  # warmup ratio based on QLoRA paper
           lr_scheduler_type="constant",  # use constant learning rate scheduler
           push_to_hub=args.push_to_hub,  # push model to hub
           report_to="tensorboard",  # report metrics to tensorboard
           gradient_checkpointing_kwargs={
               "use_reentrant": False
           },  # use reentrant checkpointing
           dataset_text_field="",  # need a dummy field for collator
           dataset_kwargs={"skip_prepare_dataset": True},  # important for collator
       )
       args.remove_unused_columns = False  # important for collator
   
       # Create a data collator to encode text and image pairs
       def collate_fn(examples):
           texts = []
           images = []
           for example in examples:
               image_inputs = process_vision_info(example["messages"])
               text = processor.apply_chat_template(
                   example["messages"], add_generation_prompt=False, tokenize=False
               )
               texts.append(text.strip())
               images.append(image_inputs)
   
           # Tokenize the texts and process the images
           batch = processor(text=texts, images=images, return_tensors="pt", padding=True)
   
           # The labels are the input_ids, and we mask the padding tokens and image tokens in the loss computation
           labels = batch["input_ids"].clone()
   
           # Mask image tokens
           boi_token = processor.tokenizer.special_tokens_map.get("boi_token")
           image_token_id = processor.tokenizer.convert_tokens_to_ids(boi_token) if boi_token else None
   
           # Mask tokens for not being used in the loss computation
           labels[labels == processor.tokenizer.pad_token_id] = -100
           if image_token_id is not None:
               labels[labels == image_token_id] = -100
           labels[labels == 262144] = -100
   
           batch["labels"] = labels
           return batch
   
       trainer = SFTTrainer(
           model=model,
           args=args,
           train_dataset=dataset,
           peft_config=peft_config,
           processing_class=processor,
           data_collator=collate_fn,
       )
   
       callback = ray.train.huggingface.transformers.RayTrainReportCallback()
       trainer.add_callback(callback)
       trainer = ray.train.huggingface.transformers.prepare_trainer(trainer)
   
       # Start training, the model will be automatically saved to the Hub and the output directory
       trainer.train()
   
       # Save the final model again to the Hugging Face Hub
       trainer.save_model()
   
   if __name__ == "__main__":
       args = get_args()
       print("Starting training task!")
       training_name = f"gemma_vision_train_{datetime.datetime.now().strftime('%Y_%m_%d_%H_%M_%S')}"
   
       gcs_bucket = args.gcs_bucket
       if not gcs_bucket.startswith("gs://"):
           gcs_bucket = "gs://" + gcs_bucket
   
       run_config = RunConfig(
           storage_path=gcs_bucket,
           name=training_name,
       )
       scaling_config = ScalingConfig(num_workers=16, use_gpu=True, accelerator_type="B200")
       ray_trainer = TorchTrainer(train, train_loop_config=args, scaling_config=scaling_config, run_config=run_config)
       print("Commencing training!")
       result = ray_trainer.fit()
   

3. 保存文件。

4. 在集群中创建 ConfigMap 对象：

   kubectl create cm ray-job-cm --from-file=code -o yaml --dry-run=client | kubectl apply -f -

   如需更新训练脚本，请重新运行上述命令。更改可能需要一分钟才能传播到所有 pod。

配置 Ray 集群

1. 如需在 GKE 集群中创建 Ray 集群，请将以下 YAML 保存为名为 ray_cluster.yaml 的文件。

   apiVersion: ray.io/v1
   kind: RayCluster
   metadata:
     name: gemma3-tuning
   spec:
     rayVersion: '2.48.0'
     headGroupSpec:
       rayStartParams:
         dashboard-host: '0.0.0.0'
       template:
         metadata:
         spec:
           serviceAccountName: $RAY_SA
           containers:
           - name: ray-head
             image: rayproject/ray:2.48.0
             ports:
             - containerPort: 6379
               name: gcs
             - containerPort: 8265
               name: dashboard
             - containerPort: 10001
               name: client
             resources:
               limits:
                 cpu: "24"
                 ephemeral-storage: "9Gi"
                 memory: "64Gi"
               requests:
                 cpu: "24"
                 ephemeral-storage: "9Gi"
                 memory: "64Gi"
             env:
               - name: HF_TOKEN
                 valueFrom:
                   secretKeyRef:
                     name: hf-secret
                     key: hf_api_token
             volumeMounts:
               - name: job-code
                 mountPath: /code/
               - mountPath: /mnt/local-ssd/
                 name: local-storage
           volumes:
             - name: job-code
               configMap:
                 name: ray-job-cm
             - name: local-storage
               emptyDir: { }
     workerGroupSpecs:
     - replicas: 2
       minReplicas: 1
       maxReplicas: 5
       groupName: gpu-group
       rayStartParams: {}
       template:
         spec:
           serviceAccountName: $RAY_SA
           containers:
           - name: ray-worker
             image: rayproject/ray:2.48.0-gpu
             resources:
               limits:
                 nvidia.com/gpu: "8"
                 cpu: "80"
                 memory: "320Gi"
               requests:
                 nvidia.com/gpu: "8"
                 cpu: "80"
                 memory: "320Gi"
             env:
               - name: HF_TOKEN
                 valueFrom:
                   secretKeyRef:
                     name: hf-secret
                     key: hf_api_token
             volumeMounts:
               - name: job-code
                 mountPath: /code/
               - mountPath: /mnt/local-ssd/
                 name: local-storage
           volumes:
             - name: job-code
               configMap:
                 name: ray-job-cm
             - name: local-storage
               emptyDir: { }
           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

2. 使用以下命令将此 YAML 定义应用于您的集群：

   envsubst < ray_cluster.yaml | kubectl apply -f -

   $RESERVATION 标志会自动替换为您配置为环境变量的名称。

   Ray Operator 会创建 raylet Pod，这会触发集群自动扩缩，以便为这些 Pod 提供合适的节点。系统会在集群中创建三个 pod：一个头节点和两个工作器节点。工作器节点配备了 B200 GPU。

3. 如需验证所有三个 pod 是否已准备就绪，请运行以下命令：

   kubectl get pods

   就绪的 Ray 集群的 pod 列表类似于以下内容：

   NAME                                   READY   STATUS    RESTARTS   AGE
   gemma3-tuning-gpu-group-worker-s4h8f   2/2     Running   0          16m
   gemma3-tuning-gpu-group-worker-stg5f   2/2     Running   0          5m34s
   gemma3-tuning-head-zbdvp               2/2     Running   0          16m
   

安排训练作业

1. 将以下内容保存为 ray_job.yaml 文件：

   apiVersion: ray.io/v1
   kind: RayJob
   metadata:
     name: test-ray-job
   spec:
     entrypoint: python /code/vision_train.py --gcs_bucket $GCS_BUCKET
     runtimeEnvYAML: |
       pip:
         - torch==2.8.0
         - torchvision==0.23.0
         - ray==2.48.0
         - transformers==4.55.2
         - datasets==4.0.0
         - evaluate==0.4.5
         - accelerate==1.10.0
         - pillow==11.3.0
         - bitsandbytes==0.47.0
         - trl==0.21.0
         - peft==0.17.0
     clusterSelector:
       ray.io/cluster: gemma3-tuning
   

2. 将 RayJob 定义提交到 RayCluster：