Ottimizza e scala l'apprendimento per rinforzo con NVIDIA NeMo RL su GKE

Questo tutorial mostra come orchestrare un ambiente di addestramento distribuito per il reinforcement learning (RL) su Google Kubernetes Engine (GKE). Utilizzi Ray e il framework NVIDIA NeMo RL per configurare un ambiente di addestramento distribuito per ottimizzare un modello.

Questo tutorial si concentra sulla pipeline di addestramento dell'ottimizzazione relativa dei criteri di gruppo (GRPO) su GKE con Ray e NeMo RL. GRPO è un algoritmo di apprendimento per rinforzo progettato per migliorare la capacità di ragionamento di un modello. Questo algoritmo efficiente in termini di memoria semplifica il processo di apprendimento per rinforzo eliminando il Critic, o modello di valore, e utilizzando un calcolo relativo basato sul gruppo.

Prima di eseguire questo tutorial, ti consigliamo di completare il tutorial Perfezionare e scalare l'apprendimento per rinforzo con verl su GKE. Il seguente tutorial utilizza la stessa configurazione e lo stesso setup del cluster del tutorial sull'ottimizzazione e lo scaling del RL con verl.

Sfondo

Le sezioni seguenti forniscono una breve panoramica dei concetti utilizzati in questo tutorial.

Apprendimento per rinforzo (RL)

L'RL insegna ai modelli attraverso l'esperienza, l'esplorazione e il feedback, anziché l'imitazione statica. Anche se il pre-addestramento insegna a un modello cosa dire, l'apprendimento per rinforzo con feedback umano (RLHF) gli insegna a essere utile, sicuro e logico. L'RL funge da ponte tra un modello di base e un modello ottimizzato per un caso d'uso specializzato.

Per saperne di più, consulta Che cos'è l'apprendimento per rinforzo?

Ottimizzazione delle policy relative al gruppo (GRPO)

GRPO, un algoritmo reso popolare da DeepSeek, offre un'alternativa efficiente in termini di memoria all'ottimizzazione delle norme prossimali (PPO) per l'allineamento degli LLM rimuovendo il modello Critic. Anziché una rete di critici, GRPO genera un gruppo di risposte per lo stesso prompt e utilizza la ricompensa media di quel gruppo come baseline.

Per saperne di più, consulta GRPO.

NVIDIA NeMo RL

NeMo RL è la libreria open source di NVIDIA per il post-training progettata per l'RL scalabile. Parte dell'ecosistema più ampio del framework NeMo, NeMo RL consente esperimenti su piccola scala su una singola GPU e deployment multimodali su migliaia di GPU.

Per ulteriori informazioni, consulta NVIDIA NeMo RL.

Set di dati GSM8k

In questo tutorial utilizzi il set di dati GSM8k, che contiene 8500 problemi di matematica di alta qualità e linguisticamente diversi per la scuola elementare.

Utilizzando GSM8k e GRPO, il modello genera un gruppo di n risposte diverse per lo stesso problema. GRPO confronta queste risposte con la media del gruppo. Il modello viene premiato maggiormente per i percorsi che sono costantemente corretti e logicamente validi rispetto al resto del gruppo. Nel tempo, il modello impara che articolare chiaramente i suoi passaggi è il modo più affidabile per massimizzare la ricompensa, riducendo di fatto la ricompensa per le risposte con un rendimento scarso.

Per saperne di più, consulta GSM8k.

Obiettivi

Questo tutorial mostra come configurare RL su GKE con NeMo RL completando i seguenti passaggi:

  1. Prepara l'ambiente.
  2. Configura un cluster GKE con GPU B200 o H200.
  3. Configura KubeRay per gestire un cluster Ray distribuito.
  4. Utilizza Managed Lustre per l'archiviazione ad alte prestazioni.
  5. Esegui un job di addestramento GRPO che utilizza NeMo RL.

Prima di iniziare

  • Accedi al tuo account Google Cloud . Se non conosci Google Cloud, crea un account per valutare le prestazioni dei nostri prodotti in scenari reali. I nuovi clienti ricevono anche 300 $di crediti senza costi per l'esecuzione, il test e il deployment dei workload.
  • Installa Google Cloud CLI.

  • Se utilizzi un provider di identità (IdP) esterno, devi prima accedere a gcloud CLI con la tua identità federata.

  • Per inizializzare gcloud CLI, esegui questo comando:

    gcloud init
  • Crea o seleziona un Google Cloud progetto.

    Ruoli richiesti per selezionare o creare un progetto

    • Seleziona un progetto: la selezione di un progetto non richiede un ruolo IAM specifico. Puoi selezionare qualsiasi progetto per il quale ti è stato concesso un ruolo.
    • Crea un progetto: per creare un progetto, devi disporre del ruolo Autore progetto (roles/resourcemanager.projectCreator), che contiene l'autorizzazione resourcemanager.projects.create. Scopri come concedere i ruoli.
    • Creare un progetto Google Cloud :

      gcloud projects create PROJECT_ID

      Sostituisci PROJECT_ID con un nome per il progetto Google Cloud che stai creando.

    • Seleziona il progetto Google Cloud che hai creato:

      gcloud config set project PROJECT_ID

      Sostituisci PROJECT_ID con il nome del progetto Google Cloud .

  • Verifica che la fatturazione sia abilitata per il tuo progetto Google Cloud .

  • Abilita le API richieste:

    Ruoli richiesti per abilitare le API

    Per abilitare le API, devi disporre del ruolo IAM Amministratore utilizzo dei servizi (roles/serviceusage.serviceUsageAdmin), che include l'autorizzazione serviceusage.services.enable. Scopri come concedere i ruoli.

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

  • Se utilizzi un provider di identità (IdP) esterno, devi prima accedere a gcloud CLI con la tua identità federata.

  • Per inizializzare gcloud CLI, esegui questo comando:

    gcloud init
  • Crea o seleziona un Google Cloud progetto.

    Ruoli richiesti per selezionare o creare un progetto

    • Seleziona un progetto: la selezione di un progetto non richiede un ruolo IAM specifico. Puoi selezionare qualsiasi progetto per il quale ti è stato concesso un ruolo.
    • Crea un progetto: per creare un progetto, devi disporre del ruolo Autore progetto (roles/resourcemanager.projectCreator), che contiene l'autorizzazione resourcemanager.projects.create. Scopri come concedere i ruoli.
    • Creare un progetto Google Cloud :

      gcloud projects create PROJECT_ID

      Sostituisci PROJECT_ID con un nome per il progetto Google Cloud che stai creando.

    • Seleziona il progetto Google Cloud che hai creato:

      gcloud config set project PROJECT_ID

      Sostituisci PROJECT_ID con il nome del progetto Google Cloud .

  • Verifica che la fatturazione sia abilitata per il tuo progetto Google Cloud .

  • Abilita le API richieste:

    Ruoli richiesti per abilitare le API

    Per abilitare le API, devi disporre del ruolo IAM Amministratore utilizzo dei servizi (roles/serviceusage.serviceUsageAdmin), che include l'autorizzazione serviceusage.services.enable. Scopri come concedere i ruoli.

    gcloud services enable container.googleapis.com storage.googleapis.com compute.googleapis.com
  • Concedi ruoli al tuo account utente. Esegui il seguente comando una volta per ciascuno dei seguenti ruoli IAM: roles/container.admin, roles/iam.serviceAccountAdmin, roles/storage.admin

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

    Sostituisci quanto segue:

    • PROJECT_ID: il tuo ID progetto.
    • USER_IDENTIFIER: l'identificatore del tuo account utente . Ad esempio: myemail@example.com.
    • ROLE: il ruolo IAM che concedi al tuo account utente.

prepara l'ambiente

In questo tutorial utilizzi Cloud Shell.

  1. Vai alla consoleGoogle Cloud .

  2. Nella parte superiore della finestra della console Google Cloud , fai clic sul pulsante Attiva Cloud Shell.

  3. Imposta le seguenti variabili di ambiente:

    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
    

    Sostituisci i seguenti valori:

    • CLUSTER_NAME: il nome del cluster GKE.
    • CONTROL_PLANE_LOCATION: la regione Compute Engine per il control plane del cluster GKE.
    • NODE_LOCATION: la posizione dei nodi. Seleziona una zona in cui sono disponibili le GPU NVIDIA B200 o H200.
    • GPU_TYPE: l'acceleratore che hai prenotato nella prenotazione di capacità di Compute Engine. Deve essere uno dei seguenti valori:
      • nvidia-b200: NVIDIA B200 (180 GB)
      • nvidia-h200-141gb: NVIDIA H200 (141 GB)
    • MACHINE_TYPE: il tipo di macchina da utilizzare:
      • Per le GPU NVIDIA B200 (180 GB), utilizza a4-highgpu-8g o versioni successive.
      • Per le GPU NVIDIA H200 (141 GB), utilizza a3-ultragpu-8g o versioni successive.
    • GKE_VERSION: la versione di GKE da utilizzare:

      • Per le GPU NVIDIA B200 (180 GB), utilizza 1.32.2-gke.1422000 o versioni successive.
      • Per le GPU NVIDIA H200 (141 GB), utilizza 1.31.4-gke.1183000 o versioni successive.
    • BUCKET_NAME: il nome di base del bucket Cloud Storage.

    • YOUR_HUGGING_FACE_TOKEN: il tuo token Hugging Face.

  4. Crea le seguenti variabili di ambiente per la rete:

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

    Sostituisci i seguenti valori:

    • GVNIC-NAME: il prefisso per il nome della rete gVNIC. Puoi utilizzare qualsiasi prefisso.
    • RDMA-NAME: il prefisso per la rete di accesso diretto alla memoria (RDMA) remota. Puoi utilizzare qualsiasi prefisso.

Configurazione dell'infrastruttura

In questa sezione crei reti VPC e un cluster GKE.

Crea una rete VPC

  1. Crea una rete VPC per l'interfaccia 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. Crea una rete VPC e subnet per RDMA che includa otto subnet per otto 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
    

Crea il cluster GKE

Puoi impostare NeMo RL in un cluster GKE Autopilot o Standard. Ti consigliamo di utilizzare un cluster Autopilot per un'esperienza Kubernetes completamente gestita. Per scegliere la modalità operativa GKE più adatta ai tuoi workload, consulta Informazioni sulle modalità operative di GKE.

Autopilot

  1. Crea un cluster Autopilot:

    gcloud container clusters create-auto ${CLUSTER_NAME} \
        --location=${CONTROL_PLANE_LOCATION} \
        --enable-multi-networking  \
        --enable-ray-operator
    
  2. Recupera le credenziali per il tuo cluster:

    gcloud container clusters get-credentials ${CLUSTER_NAME} \
        --location=${CONTROL_PLANE_LOCATION}
    
  3. Installa il programma di installazione NCCL RDMA per Autopilot:

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

Standard

  1. Crea un 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. Recupera le credenziali per il tuo cluster:

    gcloud container clusters get-credentials ${CLUSTER_NAME} \
        --location=${CONTROL_PLANE_LOCATION}
    
  3. Crea il pool di nodi 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. Installa il programma di installazione NCCL RDMA:

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

Configurare i mapping di rete

  1. Salva il seguente manifest come 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. Applica il manifest:

    kubectl apply -f network-mapping.yaml
    

Preparare lo spazio di archiviazione

In questa sezione, creerai bucket Cloud Storage e un'istanza Managed Lustre, che esegue il provisioning dello spazio di archiviazione ad alte prestazioni richiesto per il carico di lavoro RL.

  1. Crea un bucket Cloud Storage:

    gcloud storage buckets create gs://${GS_BUCKET} \
        --location=${CONTROL_PLANE_LOCATION} \
        --enable-hierarchical-namespace \
        --uniform-bucket-level-access
    
  2. Crea un service account Kubernetes (KSA) e associalo al 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. Configura Managed Lustre completando i seguenti passaggi:

    1. Crea un'istanza Managed Lustre seguendo i passaggi descritti nella sezione Crea un'istanza Managed Lustre. Assicurati che l'istanza utilizzi la stessa rete del tuo cluster GKE.
    2. Accedi all'istanza Managed Lustre seguendo i passaggi descritti in Accedere a un'istanza Managed Lustre esistente.

Esegui il deployment di RayCluster

In questa sezione clonerai il repository di esempio, preparerai i manifest e eseguirai uno script launcher.sh:

  1. Clona il repository di esempio:

    git clone https://github.com/GoogleCloudPlatform/kubernetes-engine-samples.git
    cd kubernetes-engine-samples
    
  2. Vai alla directory di lavoro:

    cd ai-ml/nemo-rl-on-gke/nemoRL
    
  3. Ispeziona il manifest 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
    

    Sostituisci NCCL_TUNER_CONFIG_PATH con uno dei seguenti valori, in base all'acceleratore che utilizzi in questo tutorial:

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

    In questo manifest, il nodo head gestisce il job e ospita la dashboard Ray. I nodi worker eseguono i job di addestramento.

  4. Installa il cluster Ray:

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

    Per questo tutorial, utilizzi due nodi worker. Se vuoi modificare il numero di nodi worker, modifica il valore di REPLICA_COUNT.

  5. Per eseguire il deployment del cluster Ray, esegui lo script launcher.sh:

    bash launcher.sh
    
  6. Verifica che i nodi worker e head siano in esecuzione:

    kubectl get pods
    

    L'output è simile al seguente:

    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. Verifica che il cluster Ray sia in esecuzione:

    kubectl ray get clusters
    

    L'output è simile al seguente:

    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
    

Avvia il job GRPO

Quando il cluster Ray è pronto, puoi inviare un job Ray al cluster Ray in esecuzione su GKE. NeMo RL scarica automaticamente il modello durante l'esecuzione del job di addestramento RL.

Per inviare un job Ray, avvia una sessione interattiva per eseguire il job.

  1. Per stabilire una connessione locale al cluster Ray, esegui questo comando:

      kubectl ray session ray-cluster-kuberay
    

    Questo comando avvia l'inoltro delle porte tra la tua macchina locale e il nodo head Ray nel tuo cluster GKE. Tieni presente che il terminale sarà occupato mentre questa sessione è attiva; per procedere, apri un'istanza del terminale separata.

  2. Modifica il 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."
    

    Sostituisci i seguenti valori nel file gemma3-27b-gsm8k.sh:

    • YOUR_WANDB_API_KEY: la tua chiave API WandB.
    • YOUR_HF_TOKEN: il tuo token Hugging Face.

    In questo file puoi vedere la configurazione per eseguire un job con il modello gemma3-27b-it sul set di dati GSM8k. Per completare la pipeline di addestramento GRPO, questo script definisce i seguenti parametri:

    • num_prompts_per_step: 16 e num_generations_per_prompt: 64: il modello Gemma3-27b-it genera un ampio gruppo di risposte per ogni prompt. In questa configurazione, il modello produce 1024 risposte totali (16 × 64 = 1024).
    • policy.generation.colocated.enabled=False: questo parametro disattiva la funzionalità di generazione colocalizzata, il che significa che il modello non genera risposte nello stesso nodo del processo di addestramento. Nell'RL standard, le stesse GPU gestiscono sia l'addestramento che la generazione. In questa configurazione di NeMo RL, dedichi nodi specifici (gestiti con il parametro policy.generation.colocated.resources) esclusivamente all'inferenza vLLM, mentre il resto del cluster si concentra sui calcoli di addestramento più complessi. Separando questi carichi di lavoro, eviti la contesa delle risorse tra i buffer di addestramento ad alta intensità di memoria e i carichi di lavoro di inferenza ad alta intensità di calcolo.
  3. Per inviare il Job, esegui questo comando:

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

    Quando il job è in esecuzione, l'output mostra i risultati dell'addestramento, la tempistica e le metriche di rendimento.

Monitorare lo stato del job GRPO

Al termine del job, NeMo RL archivia i checkpoint nel percorso configurato.

  1. Installa l'utilità apt tree:

    apt install tree
    
  2. Per monitorare lo stato del job GRPO, controlla i log del nodo principale Ray:

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

    L'output è simile al seguente:

    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
    

Esegui la pulizia

Per evitare che ti vengano addebitati dei costi, elimina le risorse:

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

Passaggi successivi