Google menggunakan teknologi AI untuk menerjemahkan konten ke dalam bahasa pilihan Anda. Terjemahan AI mungkin mengandung kesalahan.

Menyesuaikan Gemma 4 di cluster GKE A4

GPU yang dikelompokkan

Tutorial ini menunjukkan cara menyesuaikan model bahasa besar (LLM) Gemma 4 di cluster Google Kubernetes Engine (GKE) multi-GPU satu host di Google Cloud. Cluster ini menggunakan satu instance virtual machine (VM) A4 dengan 8 GPU NVIDIA B200 terpasang.

Dua proses utama yang dijelaskan dalam tutorial ini adalah sebagai berikut:

Deploy cluster GKE berperforma tinggi menggunakan GKE Autopilot. Sebagai bagian dari deployment ini, Anda akan membuat image VM kustom dengan software yang diperlukan sudah diinstal sebelumnya.
Setelah cluster di-deploy, Anda menjalankan tugas penyesuaian terdistribusi dengan menggunakan kumpulan skrip yang menyertai tutorial ini. Tugas ini memanfaatkan library Hugging Face Accelerate.

Tutorial ini ditujukan untuk engineer, peneliti, administrator, dan operator platform machine learning (ML), serta spesialis data dan AI yang tertarik untuk men-deploy cluster GKE di Google Cloud untuk melatih LLM.

Tujuan

Akses model Gemma 4 menggunakan Hugging Face.
Siapkan lingkungan Anda.
Buat dan deploy cluster GKE A4.
Sesuaikan model Gemma 4 menggunakan library Hugging Face Accelerate dengan paralelisme data yang sepenuhnya di-shard (FSDP).
Pantau tugas Anda.
Jalankan pembersihan.

Biaya

Dalam dokumen ini, Anda akan menggunakan komponen Google Cloudyang dapat ditagih berikut:

Untuk membuat perkiraan biaya berdasarkan proyeksi penggunaan Anda, gunakan kalkulator harga.

Pengguna Google Cloud baru mungkin memenuhi syarat untuk mendapatkan uji coba gratis.

Sebelum memulai

Aktifkan API yang diperlukan, jika ada yang belum diaktifkan:

Peran yang diperlukan untuk mengaktifkan API
Untuk mengaktifkan API, Anda memerlukan izin serviceusage.services.enable. Jika Anda membuat project, kemungkinan Anda sudah memiliki izin ini melalui peran Pemilik (roles/owner). Jika tidak, Anda bisa mendapatkan izin ini melalui peran Admin Penggunaan Layanan (roles/serviceusage.serviceUsageAdmin). Pelajari cara memberikan peran.
```
gcloud services enable compute.googleapis.com container.googleapis.com file.googleapis.com logging.googleapis.com cloudresourcemanager.googleapis.com servicenetworking.googleapis.com
```

Aktifkan akun layanan default untuk project Google Cloud Anda:

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

Berikan peran yang diperlukan akun layanan default untuk menjalankan beban kerja penyesuaian:
```
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
```
Catatan: Langkah ini hanya memberikan peran bawaan yang diperlukan tutorial ini. Memberikan peran dasar seperti Editor (roles/editor) akan memberi akun layanan lebih banyak izin daripada yang diperlukan, dan kebijakan organisasi Anda mungkin memblokir pemberian tersebut. Untuk informasi selengkapnya, lihat Praktik terbaik untuk menggunakan akun layanan.
Buat kredensial autentikasi lokal untuk akun pengguna Anda:
```
gcloud auth application-default login
```
Catatan: Jika Anda menggunakan shell lokal dan penyedia identitas (IdP) eksternal, dan Anda mengalami error autentikasi setelah menjalankan perintah sebelumnya, maka login ke gcloud CLI dengan identitas gabungan Anda.

Aktifkan Login OS untuk project Anda:

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

Untuk mendapatkan izin yang Anda perlukan untuk menyelesaikan tutorial ini, minta administrator Anda untuk memberi Anda peran IAM berikut di project Anda:

Kubernetes Engine Admin (roles/container.admin)
Compute Admin (roles/compute.admin)
Storage Admin (roles/storage.admin)
Administrator Artifact Registry (roles/artifactregistry.admin)
Editor Cloud Build (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)

Untuk mengetahui informasi selengkapnya tentang pemberian peran, lihat Mengelola akses ke project, folder, dan organisasi.

Anda mungkin juga bisa mendapatkan izin yang diperlukan melalui peran khusus atau peran bawaan lainnya.

Mengakses Gemma 4 menggunakan Hugging Face

Untuk menggunakan Hugging Face guna mengakses Gemma 4, lakukan hal berikut:

Login ke Hugging Face
Buat token akses write Hugging Face.
Klik Profil Anda > Setelan > Token akses > +Buat token baru
Salin dan simpan nilai token write access. Anda akan menggunakannya nanti dalam tutorial ini.

Menyiapkan lingkungan Anda

Untuk menyiapkan lingkungan Anda, tetapkan hal berikut:

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}"

Ganti kode berikut:

YOUR_PROJECT_ID: ID Google Cloud project tempat Anda ingin membuat cluster GKE.
YOUR_CLUSTER_NAME: nama cluster GKE yang akan dibuat.
YOUR_REGION: region tempat Anda ingin membuat cluster GKE. Anda hanya dapat membuat cluster di region tempat reservasi Anda berada.
YOUR_RESERVATION_NAME: ID untuk kapasitas yang dipesan.
YOUR_HF_TOKEN: token akses Hugging Face yang Anda buat di bagian sebelumnya.
YOUR_ARTIFACT_REGISTRY_LOCATION: Google Cloud region tempat Anda ingin membuat repositori Artifact Registry. Untuk meminimalkan latensi penarikan image, sebaiknya gunakan region yang sama dengan yang Anda gunakan untuk YOUR_REGION.

Membuat cluster GKE dalam mode Autopilot

Untuk membuat cluster GKE dalam mode Autopilot, jalankan perintah berikut:

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

Mungkin perlu waktu beberapa saat untuk menyelesaikan pembuatan cluster GKE. Untuk memverifikasi bahwa Google Cloud telah selesai membuat cluster Anda, buka Kubernetes clusters di konsol Google Cloud .

Buat secret Kubernetes untuk kredensial Hugging Face

Untuk membuat secret Kubernetes untuk kredensial Hugging Face, ikuti langkah-langkah berikut:

Konfigurasi kubectl untuk berkomunikasi dengan cluster GKE Anda:

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

Buat secret Kubernetes untuk menyimpan token Hugging Face Anda:

kubectl create secret generic hf-secret \
    --from-literal=hf_api_token="${HF_TOKEN}" \
    --dry-run=client -o yaml | kubectl apply -f -

Menyiapkan workload Anda

Untuk menyiapkan workload, Anda perlu melakukan hal berikut:

Buat skrip beban kerja.
Gunakan Docker dan Cloud Build untuk membuat container penyesuaian.

Membuat skrip beban kerja

Untuk membuat skrip yang digunakan workload fine-tuning Anda, lakukan hal berikut:

Buat direktori untuk skrip beban kerja. Gunakan direktori ini sebagai direktori kerja Anda.
```
mkdir llm-finetuning-gemma
cd llm-finetuning-gemma
```

Buat file cloudbuild.yaml untuk menggunakan Google Cloud Build. File ini membuat container workload Anda dan menyimpannya di 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

Buat file Dockerfile untuk menjalankan tugas penyesuaian:

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

Buat file accel_fsdp_gemma4_config.yaml. File konfigurasi ini mengarahkan Hugging Face Accelerate untuk membagi tugas penyetelan di delapan GPU lokal pada satu host Anda dengan menggunakan FSDP:

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

Buat file 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

Buat file 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()

Menggunakan Docker dan Cloud Build untuk membuat container penyesuaian

Buat Repositori Docker Artifact Registry:

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

Di direktori llm-finetuning-gemma yang Anda buat di langkah sebelumnya, jalankan perintah berikut untuk membuat image penyesuaian dan mengirimkannya ke Artifact Registry.
```
gcloud builds submit . \
    --substitutions=_ARTIFACT_REPO_LOCATION="${ARTIFACT_REPO_LOCATION}"
```

Ekspor URL gambar. Anda akan menggunakannya di langkah selanjutnya dalam tutorial ini:

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

Mulai workload penyesuaian Anda

Untuk memulai workload penyesuaian, lakukan hal berikut:

Terapkan manifes penyesuaian untuk membuat tugas penyesuaian:
```
envsubst < finetune.yaml | kubectl apply -f -
```
Karena Anda menggunakan cluster dalam mode GKE Autopilot, mungkin perlu waktu beberapa menit untuk memulai node yang mendukung GPU.
Pantau tugas dengan menjalankan perintah berikut:
```
watch kubectl get pods
```
Setelah pod berjalan, periksa log tugas:
```
kubectl logs job.batch/finetune-job -f
```
Resource tugas mendownload data model, lalu menyetel model di semua delapan GPU. Proses download akan memerlukan waktu sekitar lima menit. Setelah download selesai, proses penyesuaian memerlukan waktu sekitar dua jam 30 menit untuk diselesaikan.

Memantau workload Anda

Anda dapat memantau penggunaan GPU di cluster GKE untuk memverifikasi bahwa tugas penyesuaian Anda berjalan secara efisien. Untuk melakukannya, buka link berikut di browser Anda:

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\"))"

Saat memantau workload, Anda dapat melihat hal berikut:

Penggunaan GPU: untuk tugas penyesuaian yang berjalan lancar, Anda dapat melihat penggunaan semua 8 GPU Anda meningkat dan stabil ke tingkat yang tinggi selama pelatihan.
Durasi tugas: tugas akan memerlukan waktu sekitar dua jam 30 menit untuk diselesaikan di cluster A4 yang ditentukan.

Pembersihan

Agar tidak menimbulkan biaya tambahan, hapus resource yang dibuat selama tutorial ini.

Menghapus resource

Untuk menghapus tugas penyesuaian, jalankan perintah berikut:
```
kubectl delete job finetune-job
```

Untuk menghapus cluster GKE, jalankan perintah berikut:

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