Menyesuaikan dan menskalakan reinforcement learning dengan NVIDIA NeMo RL di GKE

Tutorial ini menunjukkan cara mengatur lingkungan pelatihan terdistribusi untuk reinforcement learning (RL) di Google Kubernetes Engine (GKE). Anda menggunakan Ray dan framework NVIDIA NeMo RL untuk menyiapkan lingkungan pelatihan terdistribusi guna melakukan fine-tuning model.

Tutorial ini berfokus pada pipeline pelatihan Group Relative Policy Optimization (GRPO) di GKE dengan Ray dan NeMo RL. GRPO adalah algoritma pembelajaran penguatan yang dirancang untuk meningkatkan kemampuan penalaran model. Algoritma yang hemat memori ini menyederhanakan proses RL dengan menghilangkan Critic, atau model nilai, dan menggunakan perhitungan berbasis grup relatif.

Sebelum menjalankan tutorial ini, sebaiknya selesaikan tutorial Menyesuaikan dan menskalakan pembelajaran reinforcement dengan Verl di GKE. Tutorial berikut menggunakan penyiapan dan konfigurasi cluster yang sama dengan tutorial penyempurnaan dan penskalaan RL dengan verl.

Latar belakang

Bagian berikut memberikan ringkasan singkat tentang konsep yang digunakan dalam tutorial ini.

Reinforcement learning (RL)

RL mengajari model melalui pengalaman, eksplorasi, dan masukan, bukan imitasi statis. Meskipun pra-pelatihan mengajarkan model apa yang harus dikatakan, Reinforcement Learning from Human Feedback (RLHF) mengajarkannya cara menjadi bermanfaat, aman, dan logis. RL berfungsi sebagai jembatan antara model dasar dan model yang di-fine-tune untuk kasus penggunaan khusus.

Untuk mengetahui informasi selengkapnya, lihat Apa yang dimaksud dengan reinforcement learning?

Pengoptimalan Kebijakan Relatif Grup (GRPO)

GRPO, sebuah algoritma yang dipopulerkan oleh DeepSeek, menawarkan alternatif yang hemat memori untuk penyelarasan LLM dengan Proximal Policy Optimization (PPO) dengan menghapus model Critic. Daripada jaringan Critic, GRPO menghasilkan sekelompok respons untuk perintah yang sama dan menggunakan reward rata-rata grup tersebut sebagai dasar.

Untuk mengetahui informasi selengkapnya, lihat GRPO.

NVIDIA NeMo RL

NeMo RL adalah library pasca-pelatihan open source NVIDIA yang dirancang untuk RL yang dapat diskalakan. Sebagai bagian dari ekosistem framework NeMo yang lebih luas, NeMo RL memungkinkan eksperimen skala kecil pada satu GPU dan deployment multi-node di ribuan GPU.

Untuk mengetahui informasi selengkapnya, lihat NVIDIA NeMo RL.

Set data GSM8k

Dalam tutorial ini, Anda akan menggunakan set data GSM8k, yang berisi 8.500 soal matematika sekolah dasar berkualitas tinggi dan beragam secara linguistik.

Dengan menggunakan GSM8k dan GRPO, model menghasilkan sekelompok n respons berbeda untuk masalah yang sama. GRPO membandingkan respons ini dengan rata-rata grup. Model akan mendapatkan lebih banyak reward untuk jalur yang secara konsisten benar dan logis dibandingkan dengan jalur lainnya dalam grup. Seiring waktu, model akan mempelajari bahwa menjelaskan langkah-langkahnya dengan jelas adalah cara paling andal untuk memaksimalkan reward, sehingga secara efektif mengurangi reward untuk jawaban dengan performa rendah.

Untuk mengetahui informasi selengkapnya, lihat GSM8k.

Tujuan

Tutorial ini menunjukkan cara menyiapkan RL di GKE dengan NeMo RL dengan menyelesaikan langkah-langkah berikut:

  1. Siapkan lingkungan Anda.
  2. Siapkan cluster GKE dengan GPU B200 atau H200.
  3. Konfigurasi KubeRay untuk mengelola cluster Ray terdistribusi.
  4. Gunakan Managed Lustre untuk penyimpanan berperforma tinggi.
  5. Jalankan tugas pelatihan GRPO yang menggunakan NeMo RL.

Sebelum memulai

  • Login ke akun Google Cloud Anda. Jika Anda baru menggunakan Google Cloud, buat akun untuk mengevaluasi performa produk kami dalam skenario dunia nyata. Pelanggan baru juga mendapatkan kredit gratis senilai $300 untuk menjalankan, menguji, dan men-deploy workload.
  • Instal Google Cloud CLI.

  • Jika Anda menggunakan penyedia identitas (IdP) eksternal, Anda harus login ke gcloud CLI dengan identitas gabungan Anda terlebih dahulu.

  • Untuk melakukan inisialisasi gcloud CLI, jalankan perintah berikut:

    gcloud init
  • Buat atau pilih Google Cloud project.

    Peran yang diperlukan untuk memilih atau membuat project

    • Pilih project: Memilih project tidak memerlukan peran IAM tertentu—Anda dapat memilih project mana pun yang telah diberi peran.
    • Membuat project: Untuk membuat project, Anda memerlukan peran Pembuat Project (roles/resourcemanager.projectCreator), yang berisi izin resourcemanager.projects.create. Pelajari cara memberikan peran.
    • Buat Google Cloud project:

      gcloud projects create PROJECT_ID

      Ganti PROJECT_ID dengan nama untuk Google Cloud project yang Anda buat.

    • Pilih project Google Cloud yang Anda buat:

      gcloud config set project PROJECT_ID

      Ganti PROJECT_ID dengan nama project Google Cloud Anda.

  • Verifikasi bahwa penagihan diaktifkan untuk project Google Cloud Anda.

  • Aktifkan API yang diperlukan:

    Peran yang diperlukan untuk mengaktifkan API

    Untuk mengaktifkan API, Anda memerlukan peran IAM Service Usage Admin (roles/serviceusage.serviceUsageAdmin), yang berisi izin serviceusage.services.enable. Pelajari cara memberikan peran.

    gcloud services enable container.googleapis.com storage.googleapis.com compute.googleapis.com
  • Instal Google Cloud CLI.

  • Jika Anda menggunakan penyedia identitas (IdP) eksternal, Anda harus login ke gcloud CLI dengan identitas gabungan Anda terlebih dahulu.

  • Untuk melakukan inisialisasi gcloud CLI, jalankan perintah berikut:

    gcloud init
  • Buat atau pilih Google Cloud project.

    Peran yang diperlukan untuk memilih atau membuat project

    • Pilih project: Memilih project tidak memerlukan peran IAM tertentu—Anda dapat memilih project mana pun yang telah diberi peran.
    • Membuat project: Untuk membuat project, Anda memerlukan peran Pembuat Project (roles/resourcemanager.projectCreator), yang berisi izin resourcemanager.projects.create. Pelajari cara memberikan peran.
    • Buat Google Cloud project:

      gcloud projects create PROJECT_ID

      Ganti PROJECT_ID dengan nama untuk Google Cloud project yang Anda buat.

    • Pilih project Google Cloud yang Anda buat:

      gcloud config set project PROJECT_ID

      Ganti PROJECT_ID dengan nama project Google Cloud Anda.

  • Verifikasi bahwa penagihan diaktifkan untuk project Google Cloud Anda.

  • Aktifkan API yang diperlukan:

    Peran yang diperlukan untuk mengaktifkan API

    Untuk mengaktifkan API, Anda memerlukan peran IAM Service Usage Admin (roles/serviceusage.serviceUsageAdmin), yang berisi izin serviceusage.services.enable. Pelajari cara memberikan peran.

    gcloud services enable container.googleapis.com storage.googleapis.com compute.googleapis.com
  • Memberikan peran ke akun pengguna Anda. Jalankan perintah berikut satu kali untuk setiap peran IAM berikut: roles/container.admin, roles/iam.serviceAccountAdmin, roles/storage.admin

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

    Ganti kode berikut:

    • PROJECT_ID: Project ID Anda.
    • USER_IDENTIFIER: ID untuk akun pengguna Anda. Misalnya, myemail@example.com.
    • ROLE: Peran IAM yang Anda berikan ke akun pengguna Anda.

Menyiapkan lingkungan Anda

Dalam tutorial ini, Anda akan menggunakan Cloud Shell.

  1. Buka Google Cloud console.

  2. Di bagian atas jendela konsol Google Cloud , klik tombol Activate Cloud Shell.

  3. Tetapkan variabel lingkungan berikut:

    export PROJECT_ID=$(gcloud config get project)
    export PROJECT_NUMBER=$(gcloud projects describe ${PROJECT_ID} --format="value(projectNumber)")
    export CONTROL_PLANE_LOCATION=CONTROL_PLANE_LOCATION
    export NODE_LOCATION=NODE_LOCATION
    export CLUSTER_NAME=CLUSTER_NAME
    export GPU_TYPE=GPU_TYPE
    export MACHINE_TYPE=MACHINE_TYPE
    export GKE_VERSION=GKE_VERSION
    export KSA_NAME=generic-ksa
    export NAMESPACE=default
    export GS_BUCKET=BUCKET_NAME-${PROJECT_ID}
    export HF_TOKEN=YOUR_HUGGING_FACE_TOKEN
    

    Ganti nilai berikut:

    • CLUSTER_NAME: nama cluster GKE Anda.
    • CONTROL_PLANE_LOCATION: region Compute Engine untuk bidang kontrol cluster GKE.
    • NODE_LOCATION: lokasi untuk node Anda. Pilih zona tempat GPU NVIDIA B200 atau H200 tersedia.
    • GPU_TYPE: akselerator yang Anda pesan dalam reservasi kapasitas Compute Engine. Harus berupa salah satu nilai berikut:
      • nvidia-b200: NVIDIA B200 (180 GB)
      • nvidia-h200-141gb: NVIDIA H200 (141 GB)
    • MACHINE_TYPE: jenis mesin yang akan digunakan:
      • Untuk GPU NVIDIA B200 (180 GB), gunakan a4-highgpu-8g atau yang lebih baru.
      • Untuk GPU NVIDIA H200 (141 GB), gunakan a3-ultragpu-8g atau yang lebih baru.
    • GKE_VERSION: versi GKE yang akan digunakan:

      • Untuk GPU NVIDIA B200 (180 GB), gunakan 1.32.2-gke.1422000 atau yang lebih baru.
      • Untuk GPU NVIDIA H200 (141 GB), gunakan 1.31.4-gke.1183000 atau yang lebih baru.
    • BUCKET_NAME: nama dasar untuk bucket Cloud Storage Anda.

    • YOUR_HUGGING_FACE_TOKEN: token Hugging Face Anda.

  4. Buat variabel lingkungan berikut untuk jaringan:

    export GVNIC_NETWORK_PREFIX="GVNIC-NAME"
    export RDMA_NETWORK_PREFIX="RDMA-NAME"
    

    Ganti nilai berikut:

    • GVNIC-NAME: awalan untuk nama jaringan gVNIC. Anda dapat menggunakan awalan apa pun yang Anda inginkan.
    • RDMA-NAME: awalan untuk jaringan akses memori langsung (RDMA) jarak jauh. Anda dapat menggunakan awalan apa pun yang Anda inginkan.

Menyiapkan infrastruktur

Di bagian ini, Anda akan membuat jaringan VPC dan cluster GKE.

Create a VPC network

  1. Buat jaringan VPC untuk antarmuka gVNIC:

    gcloud compute networks create ${GVNIC_NETWORK_PREFIX}-net \
        --project=${PROJECT_ID} \
        --subnet-mode=custom
    gcloud compute networks subnets create ${GVNIC_NETWORK_PREFIX}-sub \
        --network=${GVNIC_NETWORK_PREFIX}-net \
        --location=${CONTROL_PLANE_LOCATION} \
        --range=192.168.0.0/24
    gcloud compute firewall-rules create ${GVNIC_NETWORK_PREFIX}-internal \
        --network=${GVNIC_NETWORK_PREFIX}-net \
        --action=ALLOW \
        --rules=tcp:0-65535,udp:0-65535,icmp \
        --source-ranges=192.168.0.0/16
    
  2. Buat jaringan VPC dan subnet untuk RDMA yang mencakup delapan subnet untuk delapan GPU:

    gcloud compute networks create ${RDMA_NETWORK_PREFIX}-net \
        --network-profile=${CONTROL_PLANE_LOCATION}-vpc-roce \
        --subnet-mode=custom
    
    for N in $(seq 0 7); do
      gcloud compute networks subnets create ${RDMA_NETWORK_PREFIX}-sub-$N \
        --network=${RDMA_NETWORK_PREFIX}-net \
        --location=${CONTROL_PLANE_LOCATION} \
        --range=192.168.$((N+1)).0/24 &
    done
    wait
    

Membuat cluster GKE

Anda dapat menyetel NeMo RL di cluster GKE Autopilot atau Standard. Sebaiknya gunakan cluster Autopilot untuk pengalaman Kubernetes yang terkelola sepenuhnya. Untuk memilih mode operasi GKE yang paling sesuai untuk workload Anda, lihat Tentang mode operasi GKE.

Autopilot

  1. Buat cluster Autopilot:

    gcloud container clusters create-auto ${CLUSTER_NAME} \
        --location=${CONTROL_PLANE_LOCATION} \
        --enable-multi-networking  \
        --enable-ray-operator
    
  2. Dapatkan kredensial untuk cluster Anda:

    gcloud container clusters get-credentials ${CLUSTER_NAME} \
        --location=${CONTROL_PLANE_LOCATION}
    
  3. Instal penginstal NCCL RDMA untuk Autopilot:

    kubectl apply -f https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/refs/heads/master/gpudirect-rdma/nccl-rdma-installer-autopilot.yaml
    

Standar

  1. Buat cluster Standard:

    gcloud container clusters create ${CLUSTER_NAME} \
        --location=${CONTROL_PLANE_LOCATION} \
        --enable-dataplane-v2 \
        --enable-ip-alias \
        --enable-multi-networking \
        --addons=RayOperator \
        --num-nodes=1
    
  2. Dapatkan kredensial untuk cluster Anda:

    gcloud container clusters get-credentials ${CLUSTER_NAME} \
        --location=${CONTROL_PLANE_LOCATION}
    
  3. Buat node pool GPU:

    gcloud container node-pools create gpu-pool \
        --cluster=${CLUSTER_NAME} \
        --node-locations=${NODE_LOCATION} \
        --machine-type=${MACHINE_TYPE} \
        --accelerator=type=${GPU_TYPE},count=8 \
        --spot \
        --additional-node-network=network=${GVNIC_NETWORK_PREFIX}-net,subnetwork=${GVNIC_NETWORK_PREFIX}-sub \
        --additional-node-network=network=${RDMA_NETWORK_PREFIX}-net,subnetwork=${RDMA_NETWORK_PREFIX}-sub-0 \
        --additional-node-network=network=${RDMA_NETWORK_PREFIX}-net,subnetwork=${RDMA_NETWORK_PREFIX}-sub-1 \
        --additional-node-network=network=${RDMA_NETWORK_PREFIX}-net,subnetwork=${RDMA_NETWORK_PREFIX}-sub-2 \
        --additional-node-network=network=${RDMA_NETWORK_PREFIX}-net,subnetwork=${RDMA_NETWORK_PREFIX}-sub-3 \
        --additional-node-network=network=${RDMA_NETWORK_PREFIX}-net,subnetwork=${RDMA_NETWORK_PREFIX}-sub-4 \
        --additional-node-network=network=${RDMA_NETWORK_PREFIX}-net,subnetwork=${RDMA_NETWORK_PREFIX}-sub-5 \
        --additional-node-network=network=${RDMA_NETWORK_PREFIX}-net,subnetwork=${RDMA_NETWORK_PREFIX}-sub-6 \
        --additional-node-network=network=${RDMA_NETWORK_PREFIX}-net,subnetwork=${RDMA_NETWORK_PREFIX}-sub-7
    
  4. Instal penginstal NCCL RDMA:

    kubectl apply -f https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/refs/heads/master/gpudirect-rdma/nccl-rdma-installer.yaml
    

Mengonfigurasi pemetaan jaringan

  1. Simpan manifes berikut sebagai network-mapping.yaml:

    # Copyright 2026 Google LLC. All rights reserved.
    #
    # Licensed under the Apache License, Version 2.0 (the "License");
    # you may not use this file except in compliance with the License.
    # You may obtain a copy of the License at
    #
    #     http://www.apache.org/licenses/LICENSE-2.0
    #
    # Unless required by applicable law or agreed to in writing, software
    # distributed under the License is distributed on an "AS IS" BASIS,
    # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    # See the License for the specific language governing permissions and
    # limitations under the License.
    apiVersion: networking.gke.io/v1
    kind: GKENetworkParamSet
    metadata:
      name: gvnic-1
    spec:
      vpc: ${GVNIC_NETWORK_PREFIX}-net
      vpcSubnet: ${GVNIC_NETWORK_PREFIX}-sub
      deviceMode: NetDevice
    ---
    apiVersion: networking.gke.io/v1
    kind: Network
    metadata:
      name: gvnic-1
    spec:
      type: "Device"
      parametersRef:
        group: networking.gke.io
        kind: GKENetworkParamSet
        name: gvnic-1
    ---
    apiVersion: networking.gke.io/v1
    kind: GKENetworkParamSet
    metadata:
      name: rdma-0
    spec:
      vpc: ${RDMA_NETWORK_PREFIX}-net
      vpcSubnet: ${RDMA_NETWORK_PREFIX}-sub-0
      deviceMode: RDMA
    ---
    apiVersion: networking.gke.io/v1
    kind: Network
    metadata:
      name: rdma-0
    spec:
      type: "Device"
      parametersRef:
        group: networking.gke.io
        kind: GKENetworkParamSet
        name: rdma-0
    ---
    apiVersion: networking.gke.io/v1
    kind: GKENetworkParamSet
    metadata:
      name: rdma-1
    spec:
      vpc: ${RDMA_NETWORK_PREFIX}-net
      vpcSubnet: ${RDMA_NETWORK_PREFIX}-sub-1
      deviceMode: RDMA
    ---
    apiVersion: networking.gke.io/v1
    kind: Network
    metadata:
      name: rdma-1
    spec:
      type: "Device"
      parametersRef:
        group: networking.gke.io
        kind: GKENetworkParamSet
        name: rdma-1
    ---
    apiVersion: networking.gke.io/v1
    kind: GKENetworkParamSet
    metadata:
      name: rdma-2
    spec:
      vpc: ${RDMA_NETWORK_PREFIX}-net
      vpcSubnet: ${RDMA_NETWORK_PREFIX}-sub-2
      deviceMode: RDMA
    ---
    apiVersion: networking.gke.io/v1
    kind: Network
    metadata:
      name: rdma-2
    spec:
      type: "Device"
      parametersRef:
        group: networking.gke.io
        kind: GKENetworkParamSet
        name: rdma-2
    ---
    apiVersion: networking.gke.io/v1
    kind: GKENetworkParamSet
    metadata:
      name: rdma-3
    spec:
      vpc: ${RDMA_NETWORK_PREFIX}-net
      vpcSubnet: ${RDMA_NETWORK_PREFIX}-sub-3
      deviceMode: RDMA
    ---
    apiVersion: networking.gke.io/v1
    kind: Network
    metadata:
      name: rdma-3
    spec:
      type: "Device"
      parametersRef:
        group: networking.gke.io
        kind: GKENetworkParamSet
        name: rdma-3
    ---
    apiVersion: networking.gke.io/v1
    kind: GKENetworkParamSet
    metadata:
      name: rdma-4
    spec:
      vpc: ${RDMA_NETWORK_PREFIX}-net
      vpcSubnet: ${RDMA_NETWORK_PREFIX}-sub-4
      deviceMode: RDMA
    ---
    apiVersion: networking.gke.io/v1
    kind: Network
    metadata:
      name: rdma-4
    spec:
      type: "Device"
      parametersRef:
        group: networking.gke.io
        kind: GKENetworkParamSet
        name: rdma-4
    ---
    apiVersion: networking.gke.io/v1
    kind: GKENetworkParamSet
    metadata:
      name: rdma-5
    spec:
      vpc: ${RDMA_NETWORK_PREFIX}-net
      vpcSubnet: ${RDMA_NETWORK_PREFIX}-sub-5
      deviceMode: RDMA
    ---
    apiVersion: networking.gke.io/v1
    kind: Network
    metadata:
      name: rdma-5
    spec:
      type: "Device"
      parametersRef:
        group: networking.gke.io
        kind: GKENetworkParamSet
        name: rdma-5
    ---
    apiVersion: networking.gke.io/v1
    kind: GKENetworkParamSet
    metadata:
      name: rdma-6
    spec:
      vpc: ${RDMA_NETWORK_PREFIX}-net
      vpcSubnet: ${RDMA_NETWORK_PREFIX}-sub-6
      deviceMode: RDMA
    ---
    apiVersion: networking.gke.io/v1
    kind: Network
    metadata:
      name: rdma-6
    spec:
      type: "Device"
      parametersRef:
        group: networking.gke.io
        kind: GKENetworkParamSet
        name: rdma-6
    ---
    apiVersion: networking.gke.io/v1
    kind: GKENetworkParamSet
    metadata:
      name: rdma-7
    spec:
      vpc: ${RDMA_NETWORK_PREFIX}-net
      vpcSubnet: ${RDMA_NETWORK_PREFIX}-sub-7
      deviceMode: RDMA
    ---
    apiVersion: networking.gke.io/v1
    kind: Network
    metadata:
      name: rdma-7
    spec:
      type: "Device"
      parametersRef:
        group: networking.gke.io
        kind: GKENetworkParamSet
        name: rdma-7
    
  2. Terapkan manifes:

    kubectl apply -f network-mapping.yaml
    

Menyiapkan penyimpanan

Di bagian ini, Anda akan membuat bucket Cloud Storage dan instance Managed Lustre, yang menyediakan penyimpanan berperforma tinggi yang diperlukan untuk workload RL Anda.

  1. Membuat bucket Cloud Storage:

    gcloud storage buckets create gs://${GS_BUCKET} \
        --location=${CONTROL_PLANE_LOCATION} \
        --enable-hierarchical-namespace \
        --uniform-bucket-level-access
    
  2. Buat Akun Layanan Kubernetes (KSA) dan ikat ke bucket:

    kubectl create serviceaccount ${KSA_NAME} --namespace ${NAMESPACE}
    
    gcloud storage buckets add-iam-policy-binding gs://${GS_BUCKET} \
        --member "principal://iam.googleapis.com/projects/${PROJECT_NUMBER}/locations/global/workloadIdentityPools/${PROJECT_ID}.svc.id.goog/subject/ns/${NAMESPACE}/sa/${KSA_NAME}" \
        --role "roles/storage.objectUser"
    
  3. Siapkan Managed Lustre dengan menyelesaikan langkah-langkah berikut:

    1. Buat instance Managed Lustre dengan mengikuti langkah-langkah di Membuat instance Managed Lustre. Pastikan instance menggunakan jaringan yang sama dengan cluster GKE Anda.
    2. Akses instance Managed Lustre dengan mengikuti langkah-langkah di Mengakses instance Managed Lustre yang ada.

Men-deploy RayCluster

Di bagian ini, Anda akan meng-clone repositori contoh, menyiapkan manifes, dan menjalankan skrip launcher.sh:

  1. Gandakan repositori sampel

    git clone https://github.com/GoogleCloudPlatform/kubernetes-engine-samples.git
    cd kubernetes-engine-samples
    
  2. Buka direktori kerja:

    cd ai-ml/nemo-rl-on-gke/nemoRL
    
  3. Periksa manifes values.yaml:

    # Copyright 2026 Google LLC
    #
    # Licensed under the Apache License, Version 2.0 (the "License");
    # you may not use this file except in compliance with the License.
    # You may obtain a copy of the License at
    #
    #     http://www.apache.org/licenses/LICENSE-2.0
    #
    # Unless required by applicable law or agreed to in writing, software
    # distributed under the License is distributed on an "AS IS" BASIS,
    # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    # See the License for the specific language governing permissions and
    # limitations under the License.
    
    image:
      repository: "nvcr.io/nvidia/nemo-rl"
      tag: "v0.5.0" 
      pullPolicy: Always
    
    nameOverride: "kuberay"
    fullnameOverride: ""
    
    common:
      containerEnv: {}
    
    configMap:
      fluentbit:
        data:
          fluent-bit.conf: |
            [INPUT]
                Name              tail
                Path              /tmp/ray/session_latest/logs/worker-*
                Tag               ray-worker
            [INPUT]
                Name              tail
                Path              /tmp/ray/session_latest/logs/raylet*
                Tag               raylet
            [INPUT]
                Name              tail
                Path              /tmp/ray/session_latest/logs/*
                Exclude_Path      /tmp/ray/session_latest/logs/debug_state.txt,/tmp/ray/session_latest/logs/raylet*,/tmp/ray/session_latest/logs/worker-*
                Tag               ray-misc
            [OUTPUT]
                Name              stackdriver
                Match             *
                resource          gce_instance
                labels_key        labels
    
    # --- Head Node Configuration ---
    head:
      enableInTreeAutoscaling: false
      serviceAccountName: ""
      rayStartParams:
        dashboard-host: '0.0.0.0'
      template:
        metadata:
          annotations:
            gke-gcsfuse/volumes: "true"
            networking.gke.io/default-interface: 'eth0'
      containerEnv:
      - name: RAY_GROUP
        value: "head"
      resources:
        limits:
          cpu: "64"
          memory: "500G"
          nvidia.com/gpu: 0
        requests:
          cpu: "64"
          memory: "500G"
          nvidia.com/gpu: 0
      tolerations:
        # - operator: "Exists"
        #   key: "components.gke.io/gke-managed-components"
        # - key: "nvidia.com/gpu"
        #   operator: "Exists"
        #   effect: "NoSchedule"
      volumeMounts:
        - mountPath: /data
          name: lustre-data
    
      volumes:
        - name: log-volume
          emptyDir: {}
        - name: fluentbit-config-volume
          configMap:
            name: "ray-cluster-kuberay-fluentbit-config"
        - name: lustre-data
          persistentVolumeClaim:
            claimName: lustre-pvc
      sidecarContainers:
        - name: fluent-bit
          image: fluent/fluent-bit:latest
          env:
          - name: RAY_GROUP
            value: "head"
          volumeMounts:
            - name: fluentbit-config-volume
              mountPath: /fluent-bit/etc/
            - mountPath: /tmp/ray
              name: log-volume
    
      # --- HEAD POD STARTUP SCRIPT ---
      command:
        - "bash"
        - "-c"
        - |
          set -ex
          echo "--- Head Pod Setup ---"
          apt-get update
          apt-get install -y sudo netcat-openbsd pciutils
          cd /opt/nemo-rl
          /usr/bin/python -m pip install uv
          /usr/bin/python -m uv venv
          echo "Head pod setup complete. Starting Ray..."
    
          exec ${KUBERAY_GEN_RAY_START_CMD}
    
      args: []
      headService: {}
      # nodeSelector:
      #   cloud.google.com/gke-accelerator: nvidia-b200 #cloud.google.com/gke-nodepool: cpu-node-pool-llama #cpu-node-pool
    
    # --- Default Worker (Disabled) ---
    worker:
      disabled: true
    
    # --- A4 GPU Worker Groups ---
    additionalWorkerGroups:
      worker-grp-0:
        disabled: false
        replicas: 4
        annotations:
          networking.gke.io/default-interface: 'eth0'
          networking.gke.io/interfaces: |
            [
              {"interfaceName":"eth0","network":"default"},
              {"interfaceName":"eth1","network":"gvnic-1"},
              {"interfaceName":"eth2","network":"rdma-0"},
              {"interfaceName":"eth3","network":"rdma-1"},
              {"interfaceName":"eth4","network":"rdma-2"},
              {"interfaceName":"eth5","network":"rdma-3"},
              {"interfaceName":"eth6","network":"rdma-4"},
              {"interfaceName":"eth7","network":"rdma-5"},
              {"interfaceName":"eth8","network":"rdma-6"},
              {"interfaceName":"eth9","network":"rdma-7"}
            ]
        containerEnv:
          - name: RAY_GROUP
            valueFrom:
              fieldRef:
                fieldPath: metadata.labels['ray.io/group']
          - name: NCCL_NET  
            value: "gIB"
          - name: NCCL_IB_GID_INDEX
            value: "3"   
          - name: GLOO_SOCKET_IFNAME
            value: "eth0"
          - name: NCCL_CROSS_NIC
            value: "0"
          - name: NCCL_SOCKET_IFNAME
            value: "eth0"
          - name: TP_SOCKET_IFNAME # Specific to DTensor/PyTorch Distributed
            value: "eth0"
          - name: NCCL_TUNER_CONFIG_PATH
            value: "/usr/local/gib/configs/tuner_config_a4.txtpb"
          - name: NCCL_NET_GDR_LEVEL
            value: "PIX"
          - name: LD_LIBRARY_PATH
            value: /usr/local/nvidia/lib64
        resources:
          limits:
            nvidia.com/gpu: 8
            cpu: "206"
            memory: "2400Gi"
          requests:
            nvidia.com/gpu: 8
            cpu: "206"
            memory: "2400Gi"
    
        nodeSelector:
          cloud.google.com/gke-accelerator: nvidia-b200
        tolerations:
          - operator: "Exists"
            key: "nvidia.com/gpu"
          - operator: "Exists"
            key: "cloud.google.com/impending-node-termination"
          - operator: "Exists"
            key: "user-workload"
        securityContext:
          privileged: true
        volumes:
          - name: log-volume
            emptyDir: {}
          - name: shared-memory
            emptyDir:
              medium: "Memory"
              sizeLimit: 240Gi
          - name: ray-tmp
            emptyDir:
              medium: "Memory"
          - name: fluentbit-config-volume
            configMap:
              name: "ray-cluster-kuberay-fluentbit-config"
          - name: nvidia-install-dir-host
            hostPath:
              path: /home/kubernetes/bin/nvidia
          - name: gib-nccl-plugin-volume
            hostPath: 
              path: /home/kubernetes/bin/gib
          - name: lustre-data
            persistentVolumeClaim:
              claimName: lustre-pvc
        volumeMounts:
          - mountPath: /tmp/ray
            name: log-volume
          - name: shared-memory
            mountPath: /dev/shm
          - name: nvidia-install-dir-host
            mountPath: /usr/local/nvidia
          - name: gib-nccl-plugin-volume
            mountPath: /usr/local/gib
          - mountPath: /data
            name: lustre-data   
        # --- WORKER POD STARTUP SCRIPT ---
        command:
          - "bash"
          - "-c"
          - |
            set -ex
    
            echo "--- Worker Pod Setup ---"
            apt-get update
            apt-get install -y sudo netcat-openbsd pciutils
            cd /opt/nemo-rl
            /usr/bin/python -m pip install uv
            /usr/bin/python -m uv venv
    
            ldconfig /usr/local/nvidia/lib64/
            ldconfig -p | grep libcuda | sed 's/^/  /'
            export LD_LIBRARY_PATH="/usr/local/gib/lib64:$LD_LIBRARY_PATH"
            source /usr/local/gib/scripts/set_nccl_env.sh
    
            echo "Worker pod setup complete. Starting Ray..."
    
            exec ${KUBERAY_GEN_RAY_START_CMD}
    
    
        sidecarContainers:
          - name: fluent-bit
            env:
              - name: RAY_GROUP
                valueFrom:
                  fieldRef:
                    fieldPath: metadata.labels['ray.io/group']
            image: fluent/fluent-bit:latest
            volumeMounts:
              - name: fluentbit-config-volume
                mountPath: /fluent-bit/etc/
              - mountPath: /tmp/ray
                name: log-volume
    
    # --- Service Config ---
    service:
      type: ClusterIP
    

    Ganti NCCL_TUNER_CONFIG_PATH dengan salah satu nilai berikut, berdasarkan akselerator yang Anda gunakan dalam tutorial ini:

    • NVIDIA B200 (180 GB): /usr/local/gib/configs/tuner_config_a4.txtpb
    • NVIDIA H200 (141 GB): /usr/local/gib/configs/tuner_config_a3u.txtpb

    Dalam manifes ini, node head mengelola Job dan menghosting Dasbor Ray. Node pekerja menjalankan Tugas pelatihan.

  4. Instal cluster Ray:

    export REPLICA_COUNT=2
    helm install ray-cluster . \
      --set values.additionalWorkerGroups.worker-grp-0.replicas=$REPLICA_COUNT
    

    Untuk tutorial ini, Anda akan menggunakan dua node pekerja. Jika Anda ingin mengubah jumlah node pekerja, ubah nilai REPLICA_COUNT.

  5. Untuk men-deploy cluster Ray, jalankan skrip launcher.sh:

    bash launcher.sh
    
  6. Pastikan node pekerja dan head berjalan:

    kubectl get pods
    

    Outputnya mirip dengan hal berikut ini:

    NAME                                          READY STATUS RESTARTS AGE
    ray-cluster-kuberay-head-sw7dp                3/3   Running 0      33h
    ray-cluster-kuberay-worker-grp-0-worker-gkbxw 3/3   Running 0      33h
    ray-cluster-kuberay-worker-grp-0-worker-kdg62 3/3   Running 0      33h
    
  7. Pastikan cluster Ray sedang berjalan:

    kubectl ray get clusters
    

    Outputnya mirip dengan hal berikut ini:

    NAME                 NAMESPACE DESIRED WORKERS AVAILABLE WORKERS CPUS GPUS TPUS MEMORY CONDITION STATUS AGE
    ray-cluster-kuberay  default   2       2           618     17   0    1573741824k RayClusterProvisioned ready 33h
    

Luncurkan Tugas GRPO

Setelah cluster Ray siap, Anda dapat mengirimkan Ray Job ke cluster Ray yang berjalan di GKE. NeMo RL otomatis mendownload model selama eksekusi tugas pelatihan RL.

Untuk mengirimkan Ray Job, mulai sesi interaktif untuk menjalankan Job.

  1. Untuk membuat koneksi lokal ke cluster Ray Anda, jalankan perintah ini:

      kubectl ray session ray-cluster-kuberay
    

    Perintah ini memulai penerusan port antara mesin lokal Anda dan node head Ray di cluster GKE Anda. Perhatikan bahwa terminal Anda akan digunakan saat sesi ini aktif; untuk melanjutkan, buka instance terminal terpisah.

  2. Edit file gemma3-27b-gsm8k.sh:

    # Copyright 2026 Google LLC
    #
    # Licensed under the Apache License, Version 2.0 (the "License");
    # you may not use this file except in compliance with the License.
    # You may obtain a copy of the License at
    #
    #     http://www.apache.org/licenses/LICENSE-2.0
    #
    # Unless required by applicable law or agreed to in writing, software
    # distributed under the License is distributed on an "AS IS" BASIS,
    # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    # See the License for the specific language governing permissions and
    # limitations under the License.
    
    #!/bin/bash
    WANDB_API_KEY='YOUR_WANDB_API_KEY' # Update this with your WANDB API key
    HF_TOKEN='YOUR_HF_TOKEN' # Update this with your HF token
    WORLD_SIZE=16
    
    # --- Step 1: Find the Ray Head Pod ---
    echo "Finding Ray head pod..."
    export HEAD_POD_NAME=$(kubectl get pods --selector=ray.io/node-type=head -o jsonpath='{.items[0].metadata.name}')
    if [ -z "$HEAD_POD_NAME" ]; then
        echo "Error: No running Ray head pod found. Please check your cluster."
        exit 1
    fi
    echo "Found head pod: $HEAD_POD_NAME"
    echo ""
    
    # --- Step 2: Define the Job Script to Run ---
    # This is the script that will be executed *inside* the head pod.
    # It assumes the 'uv venv' setup from the values.yaml is already done.
    JOB_SCRIPT=$(cat <<EOF
    set -ex
    
    echo "--- Running on Ray Head Pod ($HOSTNAME) ---"
    cd /opt/nemo-rl
    
    git pull && git checkout main
    
    sed -i 's/subset: Optional\[str\] = None/subset: Optional[str] = "main"/' /opt/nemo-rl/nemo_rl/data/datasets/response_datasets/response_dataset.py
    sed -i 's/raw_dataset = load_dataset(data_path)/raw_dataset = load_dataset(data_path, "main")/' /opt/nemo-rl/nemo_rl/data/datasets/utils.py
    
    echo "Setting environment variables..."
    export WANDB_API_KEY=$WANDB_API_KEY
    export HF_TOKEN=$HF_TOKEN
    export HF_HOME=/opt/nemo-rl/
    
    ###-----Example to launch Gemma3-27B on 3 nodes (24 GPUs)----------
    uv run python examples/run_grpo_math.py \
      --config examples/configs/recipes/llm/grpo-gemma3-27b-it-8n4g-fsdp2tp4-actckpt-long.yaml \
      cluster.num_nodes=2 \
      cluster.gpus_per_node=8 \
      grpo.max_num_steps=300 \
      checkpointing.checkpoint_dir=/data/nemo_rl_gemma3_27b_3_17 \
      data.dataset_name=ResponseDataset \
      +data.train_data_path=openai/gsm8k \
      +data.val_data_path=openai/gsm8k \
      +data.val_split=test \
      +data.train_split=train \
      +data.subset="main" \
      +data.input_key="question" \
      +data.output_key="answer" \
      logger.tensorboard_enabled=False \
      logger.wandb_enabled=True \
      logger.wandb.name='nemo_rl_gemma3_27b_3_17' \
      grpo.num_prompts_per_step=16 \
      grpo.num_generations_per_prompt=64 \
      policy.generation.colocated.enabled=False \
      policy.generation.colocated.resources.num_nodes=1 \
      policy.generation.colocated.resources.gpus_per_node=8 \
      policy.generation.vllm_cfg.tensor_parallel_size=8 \
      policy.generation.vllm_cfg.gpu_memory_utilization=0.9 \
      policy.dtensor_cfg.tensor_parallel_size=8
    
    echo "--- Job Finished ---"
    EOF
    )
    
    # --- Step 3: Execute the Job ---
    echo "Submitting job to $HEAD_POD_NAME..."
    echo "$JOB_SCRIPT" | tr -d '\r' | kubectl exec -i $HEAD_POD_NAME -c ray-head -- /bin/bash
    
    echo ""
    echo "Job submission complete."
    

    Ganti nilai berikut dalam file gemma3-27b-gsm8k.sh:

    • YOUR_WANDB_API_KEY: kunci API WandB Anda.
    • YOUR_HF_TOKEN: token Hugging Face Anda.

    Dalam file ini, Anda dapat melihat konfigurasi untuk menjalankan Tugas dengan model gemma3-27b-it pada set data GSM8k. Untuk menyelesaikan pipeline pelatihan GRPO, skrip ini menentukan parameter berikut:

    • num_prompts_per_step: 16 dan num_generations_per_prompt: 64: model Gemma3-27b-it menghasilkan sekelompok besar respons untuk setiap perintah. Dalam konfigurasi ini, model menghasilkan total 1.024 respons (16 × 64 = 1.024).
    • policy.generation.colocated.enabled=False: parameter ini menonaktifkan fitur pembuatan yang ditempatkan bersama, yang berarti model tidak membuat respons di node yang sama dengan proses pelatihan. Dalam RL standar, GPU yang sama menangani pelatihan dan pembuatan. Dalam penyiapan NeMo RL ini, Anda mendedikasikan node tertentu (dikelola dengan parameter policy.generation.colocated.resources) hanya untuk inferensi vLLM, sementara cluster lainnya berfokus pada matematika pelatihan yang berat. Dengan memisahkan workload ini, Anda mencegah perebutan resource antara buffer pelatihan yang intensif memori dan workload inferensi yang intensif komputasi.
  3. Untuk mengirimkan Job, jalankan perintah berikut:

    bash gemma3-27b-it/gemma3-27b-gsm8k.sh
    

    Saat Tugas berjalan, output akan menampilkan hasil pelatihan, waktu, dan metrik performa.

Memantau kondisi Tugas GRPO

Setelah Ray menyelesaikan Tugas, NeMo RL akan menyimpan titik pemeriksaan di jalur yang dikonfigurasi.

  1. Instal utilitas hierarki apt:

    apt install tree
    
  2. Untuk memantau kualitas Tugas GRPO, periksa log node head Ray:

    kubectl exec -it $(kubectl get pods -l ray.io/node-type=head -o name) -c ray-head -- bash
    

    Outputnya mirip dengan hal berikut ini:

    root@ray-cluster-kuberay-worker-grp-0-worker-gkbxw:/opt/nemo-rl# tree /data/nemo_rl_gemma3_27b_3_17/
    /data/nemo_rl_gemma3_27b_3_17/
    `-- step_10
        |-- config.yaml
        |-- policy
        |   |-- optimizer
        |   |   |-- __0_0.distcp
        |   |   |-- __10_0.distcp
        |   |   |-- __11_0.distcp
        |   |   |-- __12_0.distcp
        |   |   |-- __13_0.distcp
        |   |   |-- __14_0.distcp
        |   |   |-- __15_0.distcp
        |   |   |-- __1_0.distcp
        |   |   |-- __2_0.distcp
        |   |   |-- __3_0.distcp
        |   |   |-- __4_0.distcp
        |   |   |-- __5_0.distcp
        |   |   |-- __6_0.distcp
        |   |   |-- __7_0.distcp
        |   |   |-- __8_0.distcp
        |   |   `-- __9_0.distcp
        |   |-- tokenizer
        |   |   |-- chat_template.jinja
        |   |   |-- special_tokens_map.json
        |   |   |-- tokenizer.json
        |   |   `-- tokenizer_config.json
        |   `-- weights
        |       |-- __0_0.distcp
        |       |-- __10_0.distcp
        |       |-- __11_0.distcp
        |       |-- __12_0.distcp
        |       |-- __13_0.distcp
        |       |-- __14_0.distcp
        |       |-- __15_0.distcp
        |       |-- __1_0.distcp
        |       |-- __2_0.distcp
        |       |-- __3_0.distcp
        |       |-- __4_0.distcp
        |       |-- __5_0.distcp
        |       |-- __6_0.distcp
        |       |-- __7_0.distcp
        |       |-- __8_0.distcp
        |       `-- __9_0.distcp
        |-- train_dataloader.pt
        `-- training_info.json
    
    6 directories, 39 files
    

Pembersihan

Untuk menghindari biaya, hapus resource:

helm delete ray-cluster
gcloud container clusters delete ${CLUSTER_NAME} --location=${CONTROL_PLANE_LOCATION}
gcloud storage rm -r gs://${GS_BUCKET}

Langkah berikutnya