Menskalakan Reinforcement Learning dengan GKE dan Managed Lustre

1. Pengantar

Jika lebih suka menjalankan skrip yang dikemas secara langsung tanpa tutorial langkah demi langkah, Anda dapat menemukannya di repositori GoogleCloudPlatform/devrel-demos.

Dalam codelab ini, Anda akan mempelajari cara men-deploy pipeline pelatihan berperforma tinggi untuk Reinforcement Learning (RL) menggunakan Google Kubernetes Engine (GKE) dan Managed Lustre.

Beban kerja Reinforcement Learning, terutama yang menggunakan algoritma seperti Group Relative Policy Optimization (GRPO), menghasilkan data dalam jumlah besar selama "Pembuatan Pengalaman" dan memerlukan pembuatan titik pemeriksaan yang sering. Penyimpanan objek standar dapat menyebabkan hambatan selama lonjakan I/O ini, sehingga akselerator yang mahal menjadi tidak digunakan.

Anda akan menggunakan Managed Lustre, sistem file paralel, untuk menghilangkan hambatan ini dan mencapai throughput pelatihan yang lebih tinggi.

Yang akan Anda lakukan

• Konfigurasi variabel lingkungan untuk cluster Ray berbasis GPU.
• Menyediakan cluster GPU Spot di GKE dan instance Managed Lustre menggunakan Cluster Toolkit.
• Deploy cluster KubeRay dan pasang sistem file Lustre.
• Kirimkan workload pelatihan NeMo-RL.
• Amati throughput tinggi dan latensi titik pemeriksaan rendah menggunakan Cloud Monitoring.

Yang Anda butuhkan

• Browser web seperti Chrome.
• Project Google Cloud yang mengaktifkan penagihan.

Codelab ini menggunakan resource kelas atas, termasuk GPU Spot dan Managed Lustre. Meskipun kami menggunakan kapasitas Spot untuk mengurangi biaya, resource ini bisa mahal jika dibiarkan berjalan. Selalu ikuti langkah-langkah Pembersihan di akhir lab untuk menghindari biaya yang tidak terduga.

Codelab ini ditujukan bagi pengguna teknis tingkat lanjut, engineer platform, dan peneliti AI yang sudah memahami konsep GKE dan penyimpanan.

Perkiraan Durasi Total: 45 hingga 60 menit ditambah 2 jam waktu pelatihan

2. Sebelum memulai

Buat Project Google Cloud

1. Di Konsol Google Cloud, pilih atau buat project Google Cloud.
2. Pastikan penagihan diaktifkan untuk project Cloud Anda.

Mulai Cloud Shell

Cloud Shell adalah lingkungan command line yang berjalan di Google Cloud yang telah dilengkapi dengan alat yang diperlukan.

1. Klik Activate Cloud Shell di bagian atas Konsol Google Cloud.
2. Setelah terhubung ke Cloud Shell, verifikasi autentikasi Anda:

   gcloud auth list
   

3. Pastikan project Anda telah dikonfigurasi:

   gcloud config get project
   

4. Jika project Anda tidak ditetapkan seperti yang diharapkan, tetapkan project:

   export PROJECT_ID=<YOUR_PROJECT_ID>
   gcloud config set project $PROJECT_ID
   

Instal Cluster Toolkit

Codelab ini menggunakan Cluster Toolkit (gcluster) untuk men-deploy cluster GKE. Untuk mengetahui petunjuk cara menyiapkan Cluster Toolkit, lihat panduan penyiapan Cluster Toolkit.

Mengaktifkan API

Jalankan perintah ini di Cloud Shell untuk mengaktifkan semua API yang diperlukan:

gcloud services enable \
  container.googleapis.com \
  lustre.googleapis.com \
  compute.googleapis.com \
  servicenetworking.googleapis.com

3. Mengonfigurasi Variabel Lingkungan

Agar perintah dalam codelab ini tetap konsisten, siapkan beberapa variabel lingkungan.

Buat file bernama env.sh dan isi dengan konfigurasi Anda. Anda dapat menggunakan template berikut:

# Environment Variables for the RL Demo execution
export PROJECT_ID="{{'<var>'}}PROJECT_ID{{'</var>'}}"
export ZONE="us-east1-b"
export REGION="us-east1"
export CLUSTER_NAME="ray-a4-gpu-spot"
export HF_TOKEN="{{'<var>'}}YOUR_HF_TOKEN{{'</var>'}}" # Required for downloading models
export WANDB_API_KEY="{{'<var>'}}YOUR_WANDB_API_KEY{{'</var>'}}" # Optional

# Topology defaults
export NUM_NODES="8"
export GPUS_PER_NODE="8" # Fixed for A4/B200 architecture

HF_TOKEN adalah Token Akses Pengguna Hugging Face Anda, yang diperlukan untuk mendownload model. Anda bisa mendapatkannya dengan mengikuti petunjuk di Token Keamanan Hugging Face.

Ganti <YOUR_PROJECT_ID> dan <YOUR_HF_TOKEN> dengan nilai sebenarnya.

WANDB_API_KEY bersifat opsional. Digunakan untuk Weights & Biases, alat untuk memantau dan memvisualisasikan eksperimen machine learning. Jika Anda tidak memiliki akun atau tidak ingin menggunakannya, Anda dapat mengosongkannya atau membiarkannya seperti apa adanya.

Sumber file untuk memuat variabel ke sesi Anda saat ini:

source env.sh

4. Men-deploy Cluster GKE dan Managed Lustre menggunakan Cluster Toolkit

Pada langkah ini, Anda akan menggunakan Cluster Toolkit (gcluster) untuk men-deploy cluster GKE dengan GPU Spot dan secara otomatis menyediakan penyimpanan Managed Lustre dengan driver CSI Lustre dan PersistentVolumeClaim (lustre-pvc) yang telah dikonfigurasi sebelumnya.

Siapkan Blueprint

Sebelum men-deploy, tinjau cetak biru examples/gke-a4/gke-a4.yaml (untuk mengetahui detailnya, lihat Membuat cluster A4):

1. Aktifkan Managed Lustre: Hapus komentar bagian modul managed-lustre dan lustre-pvc di gke-a4.yaml.
2. Aktifkan Add-on RayOperator: Tetapkan enable_ray_operator: true di bagian setelan modul gke_cluster di gke-a4.yaml.

Men-deploy Infrastruktur

Tetapkan CIDR yang diizinkan untuk akses Cloud Shell dan deploy menggunakan gcluster deploy:

export AUTHORIZED_CIDR="$(curl -s ifconfig.me)/32"

gcluster deploy examples/gke-a4/gke-a4.yaml \
  --vars project_id=${PROJECT_ID},deployment_name=${CLUSTER_NAME},region=${REGION},zone=${ZONE},static_node_count=${NUM_NODES},authorized_cidr=${AUTHORIZED_CIDR},spot=true

Menetapkan spot=true di --vars secara dinamis menyediakan node pool GPU menggunakan kapasitas Spot tanpa memerlukan pengeditan manual pada reservation_affinity. Perhatikan bahwa kapasitas Spot dapat dihentikan. Untuk lingkungan produksi, hapus spot=true untuk menggunakan kapasitas sesuai permintaan atau mengonfigurasi reservasi.

Tunggu hingga deployment selesai. Cluster Toolkit menyediakan jaringan VPC, peering Private Service Access (PSA), sistem file Managed Lustre, driver CSI Lustre, add-on RayOperator, dan klaim penyimpanan Kubernetes (lustre-pvc) secara otomatis dalam satu deployment yang terkoordinasi.

5. Men-deploy Cluster Ray di GKE

Pada langkah ini, Anda akan men-deploy cluster KubeRay di node GKE dan memasang sistem file Lustre menggunakan PersistentVolumeClaim (lustre-pvc) yang disediakan secara otomatis oleh Cluster Toolkit.

Buat Konfigurasi RayCluster

Buat file bernama ray-cluster.yaml. Ini menentukan node head dan worker KubeRay, menggunakan jenis akselerator nvidia-b200 dan memasang volume Lustre di /lustre.

cat << EOF > ray-cluster.yaml
apiVersion: ray.io/v1
kind: RayCluster
metadata:
  name: ${CLUSTER_NAME}
  namespace: default
spec:
  rayVersion: '2.54.0'
  headGroupSpec:
    rayStartParams:
      dashboard-host: '0.0.0.0'
    template:
      spec:
        nodeSelector:
          cloud.google.com/gke-accelerator: nvidia-b200
        tolerations:
        - key: "nvidia.com/gpu"
          operator: "Exists"
          effect: "NoSchedule"
        containers:
        - name: ray-head
          image: nvcr.io/nvidia/nemo-rl:v0.4.0
          ports:
          - containerPort: 6379
            name: gcs-server
          - containerPort: 8265
            name: dashboard
          - containerPort: 10001
            name: client
          resources:
            limits:
              cpu: "32"
              memory: "1000Gi"
            requests:
              cpu: "8"
              memory: "64Gi"
          volumeMounts:
          - mountPath: /lustre
            name: lustre-storage
        volumes:
        - name: lustre-storage
          persistentVolumeClaim:
            claimName: lustre-pvc
  workerGroupSpecs:
  - groupName: gpu-worker-group
    replicas: ${NUM_NODES}
    minReplicas: ${NUM_NODES}
    maxReplicas: ${NUM_NODES}
    rayStartParams: {}
    template:
      spec:
        nodeSelector:
          cloud.google.com/gke-accelerator: nvidia-b200
        tolerations:
        - key: "nvidia.com/gpu"
          operator: "Exists"
          effect: "NoSchedule"
        containers:
        - name: ray-worker
          image: nvcr.io/nvidia/nemo-rl:v0.4.0
          resources:
            limits:
              nvidia.com/gpu: "8"
              cpu: "100"
              memory: "1000Gi"
            requests:
              nvidia.com/gpu: "8"
              cpu: "100"
              memory: "1000Gi"
          volumeMounts:
          - mountPath: /lustre
            name: lustre-storage
          - mountPath: /dev/shm
            name: dshm
        volumes:
        - name: lustre-storage
          persistentVolumeClaim:
            claimName: lustre-pvc
        - name: dshm
          emptyDir:
            medium: Memory
EOF

Hubungkan ke Cluster

Pastikan sesi Cloud Shell Anda diautentikasi ke cluster GKE Anda:

gcloud container clusters get-credentials ${CLUSTER_NAME} \
  --region ${REGION} \
  --project ${PROJECT_ID}

Terapkan Konfigurasi RayCluster

Terapkan konfigurasi cluster Ray:

kubectl apply -f ray-cluster.yaml

Memverifikasi Status Cluster

Pantau pembuatan pod:

kubectl get pods -w

Tunggu hingga pod head dan pekerja berstatus Running.

6. Mengirimkan Workload Reinforcement Learning

Pada langkah ini, Anda akan mengirimkan tugas pelatihan NeMo-RL GRPO ke cluster Ray Anda.

Menghubungkan ke Dasbor Ray

Untuk mengirimkan tugas dan melihat metrik, Anda harus terhubung ke Dasbor Ray. Karena dasbor berada di GKE, gunakan penerusan port untuk mengaksesnya dari Cloud Shell:

# Run this in a separate Cloud Shell tab or in the background
kubectl port-forward service/${CLUSTER_NAME}-head-svc 8265:8265 &

Membuat Skrip Eksekusi

Buat file bernama run_nemo_rl.sh. Skrip ini akan dieksekusi di pekerja cluster Ray. Kita menggunakan cat << EOF untuk mengisi variabel lingkungan yang Anda tetapkan sebelumnya.

cat << EOF > run_nemo_rl.sh
#!/bin/bash
set -ex

# Override job runtime conflicts (NeMo-RL passes os.environ to ray.init)
export RAY_OVERRIDE_JOB_RUNTIME_ENV=1

echo "--- Running on Ray Cluster ---"
cd /opt/nemo-rl

# Ensure directories exist on the high-speed Lustre drive
mkdir -p /lustre/huggingface_cache
mkdir -p /lustre/nemo_rl_qwen_72b_ds_cp

echo "Launching NeMo-RL GRPO training..."
uv run python examples/run_grpo_math.py \
  --config examples/configs/grpo_math_70B_megatron.yaml \
  policy.model_name='Qwen/Qwen2.5-72B-Instruct' \
  policy.megatron_cfg.converter_type='Qwen2ForCausalLM' \
  logger.wandb_enabled=False \
  cluster.num_nodes=${NUM_NODES} \
  cluster.gpus_per_node=${GPUS_PER_NODE} \
  logger.wandb.name='nemo-rl-grpo-test1' \
  grpo.max_num_steps=20 \
  grpo.num_generations_per_prompt=8 \
  grpo.num_prompts_per_step=32 \
  policy.train_global_batch_size=256 \
  checkpointing.enabled=True \
  checkpointing.save_period=2 \
  checkpointing.keep_top_k=2 \
  checkpointing.metric_name=null \
  checkpointing.checkpoint_dir=/lustre/nemo_rl_qwen_72b_ds_cp/nemo-rl-grpo-test1 \
  data.dataset_name='DeepScaler'
EOF
chmod +x run_nemo_rl.sh

Membuat File Pengabaian Ray

Buat file .rayignore untuk mencegah Ray mengupload direktori berukuran besar atau yang tidak perlu:

cat << EOF > .rayignore
cluster-toolkit/
.git/
*.sh.log
EOF

Membuat Konfigurasi Lingkungan Runtime

Buat file JSON untuk meneruskan variabel lingkungan ke tugas Ray:

cat << EOF > ray_runtime_env_nemo.json
{
  "env_vars": {
    "HF_TOKEN": "${HF_TOKEN}",
    "WANDB_API_KEY": "${WANDB_API_KEY}",
    "HF_HOME": "/lustre/huggingface_cache",
    "GLOO_SOCKET_IFNAME": "eth0",
    "NCCL_SOCKET_IFNAME": "eth0"
  }
}
EOF

Kirim Tugas

Gunakan Ray CLI untuk mengirimkan tugas ke endpoint dasbor. Jika perintah ray tidak ditemukan di Cloud Shell, Anda dapat menginstalnya dengan pip install ray:

ray job submit \
    --address="http://localhost:8265" \
    --working-dir . \
    --runtime-env ray_runtime_env_nemo.json \
    -- bash run_nemo_rl.sh

Anda akan melihat streaming log di terminal Cloud Shell. Tugas ini akan memuat model, menginisialisasi worker Ray, dan memulai loop pelatihan GRPO.

7. Memantau Performa Pelatihan

Pada langkah ini, Anda akan mengamati performa sistem file Lustre selama pelatihan dan pembuatan titik pemeriksaan.

Memeriksa Log Pelatihan

Seiring berjalannya pelatihan, Anda akan melihat log yang menunjukkan bahwa titik pemeriksaan sedang disimpan ke /lustre/nemo_rl_qwen_72b_ds_cp/nemo-rl-grpo-test1. Perhatikan bahwa pembuatan titik pemeriksaan terjadi secara asinkron dan tidak memblokir pekerja Ray terlalu lama.

Untuk melihat kecepatan pembuatan titik pemeriksaan, cari baris log yang menunjukkan titik pemeriksaan yang disimpan.

Melihat Metrik Lustre di Konsol Cloud

Untuk melihat metrik instance Lustre Anda:

1. Di Konsol Google Cloud, telusuri Managed Service for Lustre.
2. Klik nama instance Anda (misalnya, ${CLUSTER_NAME}-lustre atau rl-demo-gpu-lustre).
3. Klik tab Monitoring.

Di sini, Anda dapat mengamati:

• Throughput (Byte/dtk): Lihat lonjakan selama pembuatan titik pemeriksaan.
• Kapasitas: Memantau seberapa banyak ruang yang digunakan oleh titik pemeriksaan.

Lustre mampu menulis dengan kecepatan sangat tinggi, menulis checkpoint dalam waktu minimal

8. Membersihkan Resource

Langkah ini sangat penting untuk mencegah tagihan berkelanjutan. Jika dibiarkan berjalan, cluster GPU dan Managed Lustre dapat menimbulkan biaya yang signifikan.

Jalankan perintah berikut di Cloud Shell untuk menghancurkan semua infrastruktur yang disediakan (cluster GKE, node pool GPU, instance Managed Lustre, dan jaringan VPC) dalam satu langkah:

gcluster destroy "${CLUSTER_NAME}"

Perintah ini berjalan secara sinkron di latar depan, menghentikan semua infrastruktur yang dikelola oleh deployment dan menampilkan log progres ke terminal Anda. Tunggu hingga perintah selesai sepenuhnya sebelum menutup sesi Cloud Shell.

9. Selamat

Anda telah berhasil menyelesaikan codelab Menskalakan Reinforcement Learning dengan GKE dan Managed Lustre.

Yang telah Anda pelajari

• Cara menggunakan Cluster Toolkit untuk menyediakan cluster GPU GKE dengan instance Spot dan penyimpanan Lustre Terkelola.
• Cara men-deploy cluster KubeRay dan memasang penyimpanan Lustre.
• Cara mengirimkan workload pelatihan GRPO NeMo-RL.
• Cara mengamati performa penyimpanan selama pelatihan.

Langkah berikutnya

• Jelajahi fitur NVIDIA NeMo-RL lainnya.
• Pelajari lebih lanjut AI Hypercomputer Google Cloud.
• Tinjau dokumentasi Managed Service for Lustre.