在 GKE 上使用 Ray 微調 Gemma 3，執行視覺工作

本教學課程說明如何使用多節點 GKE 叢集上的 Ray 架構，微調 Gemma 3 模型。這個叢集使用兩個 A4 虛擬機器 (VM) 執行個體，每個執行個體都附加了八個 NVIDIA B200 GPU。

本教學課程的內容分為兩部分：

  1. 在 GKE Autopilot 叢集上準備 Ray 叢集。
  2. 執行分散式訓練工作，使用 2 個 A4 執行個體，每個執行個體有 8 個 B200 GPU。

本教學課程適用於機器學習 (ML) 工程師、研究人員、平台管理員和營運人員，以及對跨多個節點和 GPU 分散 AI 工作負載感興趣的資料和 AI 專家。

目標

  • 使用 Hugging Face 存取 Gemma 3 模型。

  • 準備環境。

  • 建立 GKE Autopilot 叢集，並在其中安裝 Ray 運算子。

  • 在 GKE 叢集上設定 Ray 叢集，接受 Ray Job。

  • 設定並執行 Ray 工作，根據視覺輸入內容微調 Gemma 3 模型。

  • 監控工作負載。

  • 清除所用資源。

費用

在本文件中，您會使用下列 Google Cloud的計費元件：

如要根據預測用量估算費用，請使用 Pricing Calculator。

初次使用 Google Cloud 的使用者可能符合免費試用期資格。

事前準備

  1. 登入 Google Cloud 帳戶。如果您是 Google Cloud新手，歡迎 建立帳戶，親自評估產品在實際工作環境中的成效。新客戶還能獲得價值 $300 美元的免費抵免額，可用於執行、測試及部署工作負載。

  2. 安裝 Google Cloud CLI。

  3. 若您採用的是外部識別資訊提供者 (IdP)，請先使用聯合身分登入 gcloud CLI。

  4. 執行下列指令，初始化 gcloud CLI：

    gcloud init

  5. 建立或選取 Google Cloud 專案。

    選取或建立專案所需的角色

    • 選取專案：選取專案時，不需要具備特定 IAM 角色，只要您在專案中獲派角色，即可選取該專案。
    • 建立專案：如要建立專案，您需要「專案建立者」角色 (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 專案名稱。

  6. 確認專案已啟用計費功能 Google Cloud 。

  7. 啟用必要的 API：

    啟用 API 時所需的角色

    如要啟用 API，您必須具備 serviceusage.services.enable 權限。如果您建立了專案，可能已透過「擁有者」角色 (roles/owner) 取得這項權限。否則，您可以透過「服務使用情形管理員」角色 (roles/serviceusage.serviceUsageAdmin) 取得這項權限。瞭解如何授予角色。 

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

  8. 安裝 Google Cloud CLI。

  9. 若您採用的是外部識別資訊提供者 (IdP)，請先使用聯合身分登入 gcloud CLI。

  10. 執行下列指令，初始化 gcloud CLI：

    gcloud init

  11. 建立或選取 Google Cloud 專案。

    選取或建立專案所需的角色

    • 選取專案：選取專案時，不需要具備特定 IAM 角色，只要您在專案中獲派角色，即可選取該專案。
    • 建立專案：如要建立專案，您需要「專案建立者」角色 (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 專案名稱。

  12. 確認專案已啟用計費功能 Google Cloud 。

  13. 啟用必要的 API：

    啟用 API 時所需的角色

    如要啟用 API，您必須具備 serviceusage.services.enable 權限。如果您建立了專案，可能已透過「擁有者」角色 (roles/owner) 取得這項權限。否則，您可以透過「服務使用情形管理員」角色 (roles/serviceusage.serviceUsageAdmin) 取得這項權限。瞭解如何授予角色。 

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

  14. 將角色授予使用者帳戶。針對下列每個 IAM 角色，執行一次下列指令： roles/compute.admin, roles/iam.serviceAccountUser, roles/file.editor, roles/storage.admin, roles/container.clusterAdmin, roles/serviceusage.serviceUsageAdmin 

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

    更改下列內容：

    • PROJECT_ID：專案 ID。
    • USER_IDENTIFIER：使用者帳戶的 ID。 例如：myemail@example.com。
    • ROLE：授予使用者帳戶的 IAM 角色。
  15. 為 Google Cloud 專案啟用預設服務帳戶： 
    gcloud iam service-accounts enable PROJECT_NUMBER-compute@developer.gserviceaccount.com \
    --project=PROJECT_ID

    將 PROJECT_NUMBER 替換為專案編號。如要查看專案編號，請參閱「 取得現有專案」。

  16. 將編輯者角色 (roles/editor) 授予預設服務帳戶：
    gcloud projects add-iam-policy-binding PROJECT_ID \
    --member="serviceAccount:PROJECT_NUMBER-compute@developer.gserviceaccount.com" \
    --role=roles/editor
  17. 為使用者帳戶建立本機驗證憑證： 
    gcloud auth application-default login
  18. 登入或建立 Hugging Face 帳戶。

使用 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="RESERVATION_ID"
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_ID：要用來建立叢集的預訂網址。根據保留項目所在的專案，指定下列其中一個值：
    • 專案中已有預留項目：指定預留項目名稱 (例如 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：儲存訓練檢查點結果的 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 clusters」(Kubernetes 叢集) 頁面。

建立 Hugging Face 憑證的 Kubernetes Secret

在 Cloud Shell 中，如要建立 Hugging Face 憑證的 Kubernetes Secret，請執行下列操作：

  1. 設定 kubectl 以連線至叢集：

    gcloud container clusters get-credentials $CLUSTER_NAME \
    --region=$REGION

  2. 建立 Kubernetes 密鑰，儲存您的 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 bucket

如要使用新的 bucket 儲存訓練構件，請執行下列指令：

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

如要使用現有 bucket，可以略過這個步驟。不過，您必須確保值區與叢集位於相同區域。

建立身分與存取權管理 (IAM) 服務帳戶

在 Cloud Shell 中建立 IAM 服務帳戶 (也稱為 Google 服務帳戶或 GSA)，並授予存取 Cloud Storage bucket 的權限，請按照下列步驟操作：

  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. 將 Storage 管理員角色 (roles/storage.admin) 授予 Cloud Storage 值區的 IAM 服務帳戶：

    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 服務帳戶： <x0A> 

    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 運算子會建立 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：

    envsubst < ray_job.yaml | kubectl apply -f -

  3. 確認叢集中有新的 Pod：

    kubectl get pods

    請記下輸出內容中顯示的 test-ray-job- Pod 完整名稱。這個名稱是您工作的專屬名稱。

  4. 檢查訓練進度。將 gemma-training-ray-job-UNIQUE_ID 換成您在上一步記下的專屬 Pod 名稱。

    kubectl logs -f <test-ray-job-UNIQUE_ID>

    您看到的輸出內容類似如下：

    2025-08-20 08:29:34,966 INFO cli.py:41 -- Job submission server address: http://gemma3-tuning-head-svc.default.svc.cluster.local:8265
2025-08-20 08:29:34,991 SUCC cli.py:65 -- -----------------------------------------------
2025-08-20 08:29:34,991 SUCC cli.py:66 -- Job 'test-ray-job-82mm7' submitted successfully
2025-08-20 08:29:34,991 SUCC cli.py:67 -- -----------------------------------------------
2025-08-20 08:29:34,992 INFO cli.py:291 -- Next steps
2025-08-20 08:29:34,992 INFO cli.py:292 -- Query the logs of the job:
2025-08-20 08:29:34,992 INFO cli.py:294 -- ray job logs test-ray-job-82mm7
2025-08-20 08:29:34,992 INFO cli.py:296 -- Query the status of the job:
2025-08-20 08:29:34,992 INFO cli.py:298 -- ray job status test-ray-job-82mm7
2025-08-20 08:29:34,992 INFO cli.py:300 -- Request the job to be stopped:
2025-08-20 08:29:34,992 INFO cli.py:302 -- ray job stop test-ray-job-82mm7
2025-08-20 08:29:35,003 INFO cli.py:312 -- Tailing logs until the job exits (disable with --no-wait):
2025-08-20 08:29:34,982 INFO job_manager.py:531 -- Runtime env is setting up.
Starting training task!
Commencing training!
2025-08-20 08:30:08,498 INFO worker.py:1606 -- Using address 10.76.0.17:6379 set in the environment variable RAY_ADDRESS
2025-08-20 08:30:08,506 INFO worker.py:1747 -- Connecting to existing Ray cluster at address: 10.76.0.17:6379...
2025-08-20 08:30:08,527 INFO worker.py:1918 -- Connected to Ray cluster. View the dashboard at 10.76.0.17:8265
2025-08-20 08:30:08,701 INFO tune.py:253 -- Initializing Ray automatically. For cluster usage or custom Ray initialization, call `ray.init(...)` before `<FrameworkTrainer>(...)`.
2025-08-20 08:30:08,951 WARNING tune_controller.py:2132 -- The maximum number of pending trials has been automatically set to the number of available cluster CPUs, which is high (519 CPUs/pending trials). If you're running an experiment with a large number of trials, this could lead to scheduling overhead. In this case, consider setting the `TUNE_MAX_PENDING_TRIALS_PG` environment variable to the desired maximum number of concurrent pending trials.
2025-08-20 08:30:08,953 WARNING tune_controller.py:2132 -- The maximum number of pending trials has been automatically set to the number of available cluster CPUs, which is high (519 CPUs/pending trials). If you're running an experiment with a large number of trials, this could lead to scheduling overhead. In this case, consider setting the `TUNE_MAX_PENDING_TRIALS_PG` environment variable to the desired maximum number of concurrent pending trials.

View detailed results here: YOUR_GCS_BUCKET/gemma_vision_train_2025_08_20_08_30_07
To visualize your results with TensorBoard, run: `tensorboard --logdir /tmp/ray/session_2025-08-20_04-43-14_215096_1/artifacts/2025-08-20_08-30-08/gemma_vision_train_2025_08_20_08_30_07/driver_artifacts`

Training started with configuration:
╭──────────────────────────────────────────────────────────────────────╮
│ Training config                                                      │
├──────────────────────────────────────────────────────────────────────┤
│ train_loop_config/dataset_name                  ...-descriptions-vlm │
│ train_loop_config/gcs_bucket                    ...-bucket-yooo-west │
│ train_loop_config/gradient_accumulation_steps                      4 │
│ train_loop_config/learning_rate                               0.0002 │
│ train_loop_config/logging_steps                                   10 │
│ train_loop_config/lora_alpha                                      16 │
│ train_loop_config/lora_dropout                                  0.05 │
│ train_loop_config/lora_r                                          16 │
│ train_loop_config/max_seq_length                                 512 │
│ train_loop_config/model_id                      google/gemma-3-4b-it │
│ train_loop_config/num_train_epochs                                 3 │
│ train_loop_config/output_dir                    ...-4b-seo-optimized │
│ train_loop_config/per_device_train_batch_size                      1 │
│ train_loop_config/push_to_hub                                  False │
│ train_loop_config/save_steps                                     100 │
│ train_loop_config/save_strategy                                epoch │
╰──────────────────────────────────────────────────────────────────────╯
(RayTrainWorker pid=45455, ip=10.76.0.71) Setting up process group for: env:// [rank=0, world_size=16]
(TorchTrainer pid=45197, ip=10.76.0.71) Started distributed worker processes:
(TorchTrainer pid=45197, ip=10.76.0.71) - (node_id=4c934ab2f646a578b03cc335586f30b943e811b645526a74c50bfca1, ip=10.76.0.71, pid=45455) world_rank=0, local_rank=0, node_rank=0
(TorchTrainer pid=45197, ip=10.76.0.71) - (node_id=4c934ab2f646a578b03cc335586f30b943e811b645526a74c50bfca1, ip=10.76.0.71, pid=45450) world_rank=1, local_rank=1, node_rank=0
(TorchTrainer pid=45197, ip=10.76.0.71) - (node_id=4c934ab2f646a578b03cc335586f30b943e811b645526a74c50bfca1, ip=10.76.0.71, pid=45454) world_rank=2, local_rank=2, node_rank=0
(TorchTrainer pid=45197, ip=10.76.0.71) - (node_id=4c934ab2f646a578b03cc335586f30b943e811b645526a74c50bfca1, ip=10.76.0.71, pid=45448) world_rank=3, local_rank=3, node_rank=0
(TorchTrainer pid=45197, ip=10.76.0.71) - (node_id=4c934ab2f646a578b03cc335586f30b943e811b645526a74c50bfca1, ip=10.76.0.71, pid=45453) world_rank=4, local_rank=4, node_rank=0
(TorchTrainer pid=45197, ip=10.76.0.71) - (node_id=4c934ab2f646a578b03cc335586f30b943e811b645526a74c50bfca1, ip=10.76.0.71, pid=45452) world_rank=5, local_rank=5, node_rank=0
(TorchTrainer pid=45197, ip=10.76.0.71) - (node_id=4c934ab2f646a578b03cc335586f30b943e811b645526a74c50bfca1, ip=10.76.0.71, pid=45451) world_rank=6, local_rank=6, node_rank=0
(TorchTrainer pid=45197, ip=10.76.0.71) - (node_id=4c934ab2f646a578b03cc335586f30b943e811b645526a74c50bfca1, ip=10.76.0.71, pid=45449) world_rank=7, local_rank=7, node_rank=0
(TorchTrainer pid=45197, ip=10.76.0.71) - (node_id=c0db52b44f891f3d6a1cedcbea4c6beb2c8434c66ef414dc15e65743, ip=10.76.0.135, pid=45729) world_rank=8, local_rank=0, node_rank=1
(TorchTrainer pid=45197, ip=10.76.0.71) - (node_id=c0db52b44f891f3d6a1cedcbea4c6beb2c8434c66ef414dc15e65743, ip=10.76.0.135, pid=45726) world_rank=9, local_rank=1, node_rank=1
(TorchTrainer pid=45197, ip=10.76.0.71) - (node_id=c0db52b44f891f3d6a1cedcbea4c6beb2c8434c66ef414dc15e65743, ip=10.76.0.135, pid=45728) world_rank=10, local_rank=2, node_rank=1
(TorchTrainer pid=45197, ip=10.76.0.71) - (node_id=c0db52b44f891f3d6a1cedcbea4c6beb2c8434c66ef414dc15e65743, ip=10.76.0.135, pid=45727) world_rank=11, local_rank=3, node_rank=1
(TorchTrainer pid=45197, ip=10.76.0.71) - (node_id=c0db52b44f891f3d6a1cedcbea4c6beb2c8434c66ef414dc15e65743, ip=10.76.0.135, pid=45725) world_rank=12, local_rank=4, node_rank=1
(TorchTrainer pid=45197, ip=10.76.0.71) - (node_id=c0db52b44f891f3d6a1cedcbea4c6beb2c8434c66ef414dc15e65743, ip=10.76.0.135, pid=45724) world_rank=13, local_rank=5, node_rank=1
(TorchTrainer pid=45197, ip=10.76.0.71) - (node_id=c0db52b44f891f3d6a1cedcbea4c6beb2c8434c66ef414dc15e65743, ip=10.76.0.135, pid=45723) world_rank=14, local_rank=6, node_rank=1
(TorchTrainer pid=45197, ip=10.76.0.71) - (node_id=c0db52b44f891f3d6a1cedcbea4c6beb2c8434c66ef414dc15e65743, ip=10.76.0.135, pid=45722) world_rank=15, local_rank=7, node_rank=1

...

Training finished iteration 3 at 2025-08-20 08:40:43. Total running time: 10min 34s
╭─────────────────────────────────────────╮
│ Training result                         │
├─────────────────────────────────────────┤
│ checkpoint_dir_name   checkpoint_000002 │
│ time_this_iter_s               152.6374 │
│ time_total_s                  525.88585 │
│ training_iteration                    3 │
│ epoch                           2.75294 │
│ grad_norm                      47.27161 │
│ learning_rate                    0.0002 │
│ loss                            22.5275 │
│ mean_token_accuracy             0.90325 │
│ num_tokens                     1583017. │
│ step                                 60 │
╰─────────────────────────────────────────╯

...

Training completed after 3 iterations at 2025-08-20 08:40:52. Total running time: 10min 43s
2025-08-20 08:40:53,113 INFO tune.py:1009 -- Wrote the latest version of all result files and experiment state to 'YOUR_GCS_BUCKET/gemma_vision_train_2025_08_20_08_30_07' in 0.1663s.

2025-08-20 08:40:58,304 SUCC cli.py:65 -- ----------------------------------
2025-08-20 08:40:58,305 SUCC cli.py:66 -- Job 'test-ray-job-82mm7' succeeded
2025-08-20 08:40:58,305 SUCC cli.py:67 -- ----------------------------------

    監控工作負載

您可以使用 Ray 的資訊主頁，監控叢集中排定的工作負載。

如要存取這個資訊主頁，請在新的終端機視窗中執行下列指令，設定叢集的通訊埠轉送功能：

kubectl port-forward service/gemma3-tuning-head-svc 8265:8265 > fwd.log 2>&1 &

  1. 在瀏覽器中開啟下列連結：http://localhost:8265。

  2. 如果您使用 Cloud Shell，執行上一個步驟的指令後，可以點選「網頁預覽」按鈕。

    選取「變更通訊埠」選項，輸入 8265，然後按一下「變更並預覽」。 Ray 資訊主頁會在新分頁中開啟。

清除所用資源

為避免因為本教學課程所用資源，導致系統向 Google Cloud 帳戶收取費用，請刪除含有相關資源的專案，或者保留專案但刪除個別資源。

刪除資源

  1. 如要刪除 Ray 叢集並釋出 GPU 節點，請執行下列指令：

    kubectl delete -f ray_cluster.yaml
    GKE 會自動縮減叢集規模，並釋出 Ray 使用的 A4 機器。

  2. 如要刪除整個 GKE 叢集，請執行下列指令：

    gcloud container clusters delete $CLUSTER_NAME \
    --region=$REGION

  3. 如要刪除整個 Cloud Storage bucket 及其所有內容，請執行下列指令：

    gcloud storage rm -r gs://$GCS_BUCKET

刪除專案

刪除 Google Cloud 專案：

gcloud projects delete PROJECT_ID

後續步驟