Ajusta y escala el aprendizaje por refuerzo con NVIDIA NeMo RL en GKE

Estándar

En este instructivo, se muestra cómo organizar un entorno de entrenamiento distribuido para el aprendizaje por refuerzo (RL) en Google Kubernetes Engine (GKE). Usas Ray y el framework de NVIDIA NeMo RL para configurar un entorno de entrenamiento distribuido y ajustar un modelo.

En este instructivo, se explica la canalización de entrenamiento de la optimización de políticas relativas al grupo (GRPO) en GKE con Ray y NeMo RL. El GRPO es un algoritmo de aprendizaje por refuerzo diseñado para mejorar la capacidad de razonamiento de un modelo. Este algoritmo eficiente en cuanto a la memoria simplifica el proceso de RL, ya que elimina el Critic, o modelo de valor, y usa un cálculo relativo basado en grupos.

Antes de ejecutar este instructivo, completa el instructivo Ajusta y escala el aprendizaje por refuerzo con Verl en GKE. En este instructivo, se usa la misma configuración del clúster que en el instructivo sobre el ajuste y el escalamiento del RL con Verl.

Fondo

En las siguientes secciones, se proporciona una breve descripción general de los conceptos que se usan en este instructivo.

Aprendizaje por refuerzo (RL)

El RL enseña a los modelos a través de la experiencia, la exploración y la retroalimentación, en lugar de la imitación estática. Si bien el entrenamiento previo le enseña a un modelo qué decir, el aprendizaje por refuerzo con retroalimentación humana (RLHF) le enseña a ser útil, seguro y lógico. El RL sirve como puente entre un modelo base y un modelo ajustado para un caso de uso especializado.

Para obtener más información, consulta ¿Qué es el aprendizaje por refuerzo?

Optimización de políticas relativas al grupo (GRPO)

GRPO, un algoritmo popularizado por DeepSeek, ofrece una alternativa eficiente en cuanto a la memoria a la optimización de políticas proximales (PPO) para la alineación de LLM, ya que quita el modelo de Critic. En lugar de una red de críticos, GRPO genera un grupo de respuestas para la misma instrucción y usa la recompensa promedio de ese grupo como referencia.

Para obtener más información, consulta GRPO.

NVIDIA NeMo RL

NeMo RL es la biblioteca de código abierto posterior al entrenamiento de NVIDIA diseñada para el RL escalable. Como parte del ecosistema más amplio del framework de NeMo, NeMo RL permite realizar experimentos a pequeña escala en una sola GPU y también implementaciones de varios nodos en miles de GPUs.

Para obtener más información, consulta NVIDIA NeMo RL.

Conjunto de datos de GSM8k

En este instructivo, usarás el conjunto de datos GSM8k, que contiene 8,500 problemas matemáticos de alta calidad y lingüísticamente diversos para estudiantes de primaria.

Con GSM8k y GRPO, el modelo genera un grupo de n respuestas diferentes para el mismo problema. El GRPO compara estas respuestas con el promedio del grupo. El modelo recibe más recompensas por las rutas que son coherentes, correctas y lógicas en comparación con el resto del grupo. Con el tiempo, el modelo aprende que articular sus pasos con claridad es la forma más confiable de maximizar la recompensa, lo que reduce de manera efectiva la recompensa por las respuestas de bajo rendimiento.

Para obtener más información, consulta GSM8k.

Objetivos

En este instructivo, se muestra cómo configurar RL en GKE con NeMo RL. Para ello, completa los siguientes pasos:

1. Prepara tu entorno.
2. Configura un clúster de GKE con GPU B200 o H200.
3. Configura KubeRay para administrar un clúster de Ray distribuido.
4. Usa Managed Lustre para el almacenamiento de alto rendimiento.
5. Ejecuta un trabajo de entrenamiento de GRPO que use NeMo RL.

Nota: En este instructivo, se usan la misma configuración y el mismo clúster que en el instructivo Ajusta y escala el aprendizaje por refuerzo con Verl en GKE. Si ya completaste ese instructivo, puedes omitir las secciones Prepara tu entorno y Configura la infraestructura. En tu infraestructura ya configurada, puedes continuar con la sección Implementa RayCluster para implementar un clúster de Ray.

Antes de comenzar

1. Instala Google Cloud CLI.

   Nota: Si ya instalaste gcloud CLI, ejecuta gcloud components update para asegurarte de tener la versión más reciente.

2. Configura gcloud CLI para usar tu identidad federada.

   Para obtener más información, consulta Accede a la gcloud CLI con tu identidad federada.

3. Para inicializar gcloud CLI, ejecuta el siguiente comando:

   gcloud init

4. Crea o selecciona un Google Cloud proyecto.

   Roles necesarios para seleccionar o crear un proyecto

   • Selecciona un proyecto: Para seleccionar un proyecto, no se requiere un rol de IAM específico. Puedes seleccionar cualquier proyecto en el que se te haya otorgado un rol.
   • Crear un proyecto: Para crear un proyecto, necesitas el rol de Creador de proyectos (roles/resourcemanager.projectCreator), que contiene el permiso resourcemanager.projects.create. Obtén más información para otorgar roles.
   Nota: Si no planeas conservar los recursos que creaste durante este procedimiento, crea un proyecto en lugar de seleccionar uno existente. Cuando termines, puedes borrar el proyecto y quitar todos los recursos asociados con él.

   • Crea un proyecto de Google Cloud :

     gcloud projects create PROJECT_ID

     Reemplaza PROJECT_ID por un nombre para el proyecto Google Cloud que estás creando.

   • Selecciona el proyecto Google Cloud que creaste:

     gcloud config set project PROJECT_ID

     Reemplaza PROJECT_ID por el nombre de tu proyecto de Google Cloud .

5. Verifica que la facturación esté habilitada para tu proyecto de Google Cloud .

6. Habilita las APIs necesarias:

   Roles necesarios para habilitar las APIs

   Para habilitar APIs, necesitas el permiso serviceusage.services.enable. Si creaste el proyecto, es probable que ya tengas este permiso a través del rol de propietario (roles/owner). De lo contrario, puedes obtener este permiso a través del rol de administrador de Service Usage (roles/serviceusage.serviceUsageAdmin). Obtén más información para otorgar roles.

   gcloud services enable container.googleapis.com storage.googleapis.com compute.googleapis.com

7. Otorga roles a tu cuenta de usuario. Ejecuta el siguiente comando una vez para cada uno de los siguientes roles de IAM: roles/container.admin, roles/iam.serviceAccountAdmin, roles/storage.admin

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

   Reemplaza lo siguiente:

   • PROJECT_ID: ID del proyecto
   • USER_IDENTIFIER: Es el identificador de tu cuenta de usuario de cuenta. Para obtener ejemplos, consulta Representa a los usuarios del grupo de personal en políticas de IAM.
   • ROLE: Es el rol de IAM que otorgas a tu cuenta de usuario.
8. Crea una cuenta de Hugging Face, si todavía no la tienes.
9. Asegúrate de tener un token de Hugging Face con read access.
10. Crea una cuenta de Weights & Biases (Wandb) si no tienes una.
11. Crea una clave de API de Wandb.
12. Asegúrate de que tu proyecto Google Cloud tenga suficiente cuota para las GPU B200 y H200. Para obtener más información, consulta Planifica la cuota de GPU y Cuota de GPU.

Prepara el entorno

En este instructivo, usarás Cloud Shell.

1. Ve a la consola deGoogle Cloud .

2. En la parte superior de la Google Cloud ventana de la consola, haz clic en el botón Activar Cloud Shell.

3. Configura las siguientes variables de entorno:

   export CONTROL_PLANE_REGION="YOUR_REGION"
   export NODE_ZONE="YOUR_ZONE"
   export CLUSTER_NAME="YOUR_CLUSTER_NAME"
   export GPU_TYPE="YOUR_GPU_TYPE"
   export MACHINE_TYPE="YOUR_MACHINE_TYPE"
   export KSA_NAME="generic-ksa"
   export NAMESPACE="default"
   export RESERVATION="RESERVATION_NAME"
   export LUSTRE_NAME="CHOSEN_LUSTRE_NAME"
   export HF_TOKEN="YOUR_HF_TOKEN"
   export WANDB_API_KEY="YOUR_WANDB_API_KEY"
   
   export PROJECT_ID=$(gcloud config get project)
   export PROJECT_NUMBER=$(gcloud projects describe "${PROJECT_ID}" --format="value(projectNumber)")

   Reemplaza los siguientes valores:

   • YOUR_REGION: Es la región de Compute Engine para el plano de control del clúster de GKE.
   • YOUR_NODE_ZONE: Es la zona de tus nodos. Selecciona una zona en la que estén disponibles las GPUs NVIDIA B200 o H200.
   • YOUR_CLUSTER_NAME: Es el nombre del clúster de GKE.
   • YOUR_GPU_TYPE: es el acelerador que reservaste en la reserva de capacidad de Compute Engine. Debe ser uno de los siguientes valores:
     • nvidia-b200: NVIDIA B200 (180 GB)
     • nvidia-h200-141gb: NVIDIA H200 (141 GB)
   • YOUR_MACHINE_TYPE: Es el tipo de máquina que se usará:
     • Para las GPU NVIDIA B200 (180 GB), usa a4-highgpu-8g o una versión posterior.
     • Para las GPU NVIDIA H200 (141 GB), usa a3-ultragpu-8g o una versión posterior.
   • YOUR_RESERVATION_NAME: Es el nombre de tu reserva de GPU.
   • CHOSEN_LUSTRE_NAME: Es el nombre de tu instancia de Lustre.
   • YOUR_HF_TOKEN: Tu token de Hugging Face.
   • YOUR_WANDB_API_KEY: Es tu clave de API de Wandb.

4. Crea las siguientes variables de entorno para la red:

   export NETWORK="YOUR_NETWORK_NAME"
   export GVNIC_NETWORK_PREFIX="GVNIC_NAME"
   export RDMA_NETWORK_PREFIX="RDMA_NAME"

   Reemplaza los siguientes valores:

   • NETWORK_NAME: Es el nombre de la red para GKE.
   • GVNIC_NAME: Es el prefijo del nombre de la red de gVNIC. Puedes usar el prefijo que quieras.
   • RDMA_NAME: Es el prefijo de la red de acceso directo a la memoria (RDMA) remota. Puedes usar el prefijo que quieras.

Configura la infraestructura

En esta sección, crearás redes de VPC y un clúster de GKE.

Crear red de VPC

1. Crea una red de VPC para la interfaz de gVNIC:

   gcloud compute networks create ${NETWORK} --subnet-mode=auto
   
   gcloud compute networks create ${GVNIC_NETWORK_PREFIX}-net \
       --subnet-mode=custom
   
   gcloud compute networks subnets create ${GVNIC_NETWORK_PREFIX}-sub \
       --network=${GVNIC_NETWORK_PREFIX}-net \
       --region=${CONTROL_PLANE_REGION} \
       --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 red de VPC y subredes para RDMA que incluya ocho subredes para ocho GPUs:

   gcloud compute networks create ${RDMA_NETWORK_PREFIX}-net \
       --network-profile=${NODE_ZONE}-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 \
       --region=${CONTROL_PLANE_REGION} \
       --range=192.168.$((N+1)).0/24 &
   done
   wait

Crea el clúster de GKE

Puedes configurar NeMo RL en un clúster de GKE Standard.

1. Crea un clúster estándar:

   gcloud container clusters create ${CLUSTER_NAME} \
       --location=${CONTROL_PLANE_REGION} \
       --workload-pool=${PROJECT_ID}.svc.id.goog \
       --enable-dataplane-v2 \
       --enable-ip-alias \
       --enable-multi-networking \
       --addons=RayOperator,LustreCsiDriver \
       --enable-legacy-lustre-port \
       --machine-type=n2-highmem-80 \
       --num-nodes=1 \
       --min-nodes=1 \
       --max-nodes=5 \
       --enable-autoscaling \
       --network=${NETWORK}

2. Obtén credenciales para el clúster:

   gcloud container clusters get-credentials $CLUSTER_NAME \
       --location=$CONTROL_PLANE_REGION

3. Crea el grupo de nodos de GPU:

   gcloud container node-pools create gpu-pool \
       --cluster=${CLUSTER_NAME} \
       --location=${CONTROL_PLANE_REGION} \
       --node-locations=${NODE_ZONE} \
       --machine-type=${MACHINE_TYPE} \
       --accelerator=type=${GPU_TYPE},count=8,gpu-driver-version=DEFAULT \
       --reservation-affinity=specific \
       --reservation=${RESERVATION} \
       --enable-autoscaling \
       --num-nodes=0 \
       --total-max-nodes=2 \
       --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. Instala el instalador de RDMA de NCCL:

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

Configura las asignaciones de red

1. Guarda el siguiente manifiesto como 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. Aplica el manifiesto

   envsubst < network-mapping.yaml | kubectl apply -f -

Prepara el almacenamiento

En esta sección, crearás una instancia de Managed Lustre, que aprovisiona el almacenamiento de alto rendimiento necesario para tu carga de trabajo de RL.

1. Asigna un rango de direcciones IP para el acceso privado a servicios:

   gcloud compute addresses create ${LUSTRE_NAME}-range \
       --global --purpose=VPC_PEERING \
       --prefix-length=20 --network=${NETWORK}

2. Conecta el intercambio de tráfico:

   gcloud services vpc-peerings connect \
       --service=servicenetworking.googleapis.com \
       --ranges=${LUSTRE_NAME}-range \
       --network=${NETWORK}

3. Crea una instancia de Managed Lustre:

   gcloud lustre instances create ${LUSTRE_NAME} \
       --per-unit-storage-throughput=500 \
       --capacity-gib=18000 \
       --filesystem=lustrefs \
       --location=${NODE_ZONE} \
       --network=projects/${PROJECT_ID}/global/networks/${NETWORK} \
       --gke-support-enabled

4. Accede a una instancia de Managed Lustre existente con el controlador CSI de Managed Lustre.

   1. Extrae la dirección IP de la instancia de Managed Lustre.

      export LUSTRE_IP=$(gcloud lustre instances describe ${LUSTRE_NAME} \
          --location=$NODE_ZONE --format="value(mountPoint)" | awk -F'@' '{print $1}')

   2. Inspecciona el manifiesto de lustre-pv.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.
      
      apiVersion: v1
      kind: PersistentVolume
      metadata:
        name: lustre-pv
      spec:
        storageClassName: lustre-rwx-500mbps-per-tib
        capacity:
          storage: 18000Gi
        accessModes:
          - ReadWriteMany
        persistentVolumeReclaimPolicy: Retain
        volumeMode: Filesystem
        claimRef:
          namespace: default
          name: lustre-pvc
        csi:
          driver: lustre.csi.storage.gke.io
          volumeHandle: "${PROJECT_ID}/${NODE_ZONE}/${LUSTRE_NAME}"
          volumeAttributes:
            ip: ${LUSTRE_IP}
            filesystem: lustrefs

   3. Aplica el manifiesto

      envsubst < lustre-pv.yaml | kubectl apply -f -

   4. Inspecciona el manifiesto de lustre-pvc.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.
      
      apiVersion: v1
      kind: PersistentVolumeClaim
      metadata:
        name: lustre-pvc
      spec:
        accessModes:
          - ReadWriteMany
        storageClassName: lustre-rwx-500mbps-per-tib
        volumeName: lustre-pv
        resources:
          requests:
            storage: 18000Gi

   5. Aplica el manifiesto

      kubectl apply -f lustre-pvc.yaml

Implementa RayCluster

En esta sección, clonarás el repositorio de muestra, prepararás los manifiestos y, luego, implementarás el clúster de Ray:

1. Clona el repositorio de ejemplo:

   git clone https://github.com/GoogleCloudPlatform/kubernetes-engine-samples.git

2. Navega hasta el directorio de trabajo:

   cd kubernetes-engine-samples/ai-ml/nemo-rl-on-gke/nemoRL

3. Inspecciona el manifiesto 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"
     nodeSelector:
       cloud.google.com/gke-nodepool: default-pool
     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
   

   Reemplaza NCCL_TUNER_CONFIG_PATH con cualquiera de los siguientes valores, según el acelerador que uses en este instructivo:

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

   En este manifiesto, el nodo principal administra el trabajo y aloja el panel de Ray. Los nodos trabajadores ejecutan los trabajos de entrenamiento.

4. Implementa el clúster de Ray:

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

   En este instructivo, usarás dos nodos de trabajo. Si quieres cambiar la cantidad de nodos trabajadores, cambia el valor de REPLICA_COUNT.

5. Verifica que los nodos trabajadores y principales se estén ejecutando:

   kubectl get pods

   El resultado es similar a lo siguiente:

   NAME                                          READY STATUS RESTARTS AGE
   ray-cluster-kuberay-head-sw7dp                2/2   Running 0      33h
   ray-cluster-kuberay-worker-grp-0-worker-gkbxw 2/2   Running 0      33h
   ray-cluster-kuberay-worker-grp-0-worker-kdg62 2/2   Running 0      33h
   

6. Verifica que el clúster de Ray esté en ejecución:

   kubectl ray get cluster

   El resultado es similar a lo siguiente:

   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
   

Inicia el trabajo de GRPO

Una vez que tu clúster de Ray esté listo, puedes enviar un trabajo de Ray a tu clúster de Ray en ejecución en GKE. NeMo RL descarga automáticamente el modelo durante la ejecución del trabajo de entrenamiento de RL.

Para enviar un trabajo de Ray, inicia una sesión interactiva para ejecutar el trabajo.

1. Para establecer una conexión local con tu clúster de Ray, ejecuta este comando:

   kubectl ray session ray-cluster-kuberay

   Este comando inicia el reenvío de puertos entre tu máquina local y el nodo principal de Ray en tu clúster de GKE. Ten en cuenta que tu terminal estará ocupada mientras esta sesión esté activa. Para continuar, abre otra instancia de terminal.

2. En una terminal independiente, navega a kubernetes-engine-samples/ai-ml/nemo-rl-on-gke/nemoRL/gemma3-27b-it y edita el archivo 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 2 nodes (16 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=10 \
     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=32 \
     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."
   

   Reemplaza los siguientes valores en el archivo gemma3-27b-gsm8k.sh:

   • YOUR_WANDB_API_KEY: Es tu clave de API de WandB.
   • YOUR_HF_TOKEN: Tu token de Hugging Face.

   En este archivo, puedes ver la configuración para ejecutar un trabajo con el modelo gemma3-27b-it en el conjunto de datos GSM8k. Para completar el canal de entrenamiento del GRPO, este script define los siguientes parámetros:

   • num_prompts_per_step: 16 y num_generations_per_prompt: 32: El modelo Gemma3-27b-it genera un grupo grande de respuestas para cada instrucción. En esta configuración, el modelo produce 512 respuestas en total (16 × 32 = 512).
   • policy.generation.colocated.enabled=False: Este parámetro inhabilita la función de generación colocada, lo que significa que el modelo no genera respuestas en el mismo nodo que el proceso de entrenamiento. En el RL estándar, las mismas GPUs controlan el entrenamiento y la generación. En esta configuración de NeMo RL, dedicas nodos específicos (administrados con el parámetro policy.generation.colocated.resources) exclusivamente a la inferencia de vLLM, mientras que el resto del clúster se enfoca en las matemáticas de entrenamiento de alta exigencia. Al separar estas cargas de trabajo, evitas la contención de recursos entre los búferes de entrenamiento que requieren mucha memoria y las cargas de trabajo de inferencia que requieren mucho procesamiento.

3. Para enviar el trabajo, ejecuta el siguiente comando:

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

   Cuando el trabajo se está ejecutando, el resultado muestra los resultados del entrenamiento, la sincronización y las métricas de rendimiento.

Completado correctamente: NeMo RL coordina un ciclo de vida distribuido en las VMs con GPU del clúster.

Supervisa el estado del trabajo de GRPO

Una vez que Ray finaliza el trabajo, NeMo RL almacena los puntos de control en la ruta de acceso configurada.

1. Para verificar el resultado del trabajo de GRPO, crea una sesión de SSH en el contenedor ray-head:

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

2. Instala la utilidad de árbol apt en la terminal del contenedor ray-head:

   apt update && apt install -y tree

3. Enumera la estructura de directorios del contenedor ray-head:

   tree /data/nemo_rl_gemma3_27b_3_17/

   El resultado es similar a lo siguiente:

   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
   

Nota: Los puntos de control que genera NeMo RL están en formato de punto de control distribuido de PyTorch. Para evaluar o publicar el modelo, debes convertir el punto de control al formato de ponderaciones de Hugging Face. Si deseas obtener instrucciones para convertir el punto de control, consulta Cómo convertir el formato DCP de PyTorch al formato de HF.

Realiza una limpieza

Para evitar que se apliquen cargos a tu Google Cloud cuenta por los recursos usados en este instructivo, borra los recursos individuales o el proyecto que los contiene.

Borra recursos

1. Borra tu clúster de Slurm:

   helm delete ray-cluster

2. Borra tu clúster de GKE:

   gcloud container clusters delete ${CLUSTER_NAME} \
       --location=${CONTROL_PLANE_REGION} \
       --quiet

3. Borra tu sistema de archivos Lustre:

   gcloud lustre instances delete ${LUSTRE_NAME} --location=${NODE_ZONE} --quiet

4. Borra el intercambio de tráfico entre VPC:

   gcloud services vpc-peerings delete \
       --service=servicenetworking.googleapis.com \
       --network=${NETWORK}

5. Borra el rango de direcciones IP privadas de Lustre:

   gcloud compute addresses delete ${LUSTRE_NAME}-range --global --quiet

6. Borra las subredes de RDMA y gVNIC:

   gcloud compute networks subnets delete ${GVNIC_NETWORK_PREFIX}-sub \
       --region=${CONTROL_PLANE_REGION} --quiet
   
   for N in $(seq 0 7); do
     gcloud compute networks subnets delete ${RDMA_NETWORK_PREFIX}-sub-$N \
       --region=${CONTROL_PLANE_REGION} --quiet &
   done
   wait

7. Borra las reglas de firewall y las redes:

   echo "[$(date)] ========== Deleting firewall rules and networks =========="
   
   NETWORKS=(
       "${RDMA_NETWORK_PREFIX}-net"
       "${GVNIC_NETWORK_PREFIX}-net"
       "${NETWORK}"
   )
   
   for NW in "${NETWORKS[@]}"; do
   
     echo "========== Deleting firewall rules for ${NW} =========="
     while true; do
         rules=$(gcloud compute firewall-rules list \
           --filter="network:${NW}" \
           --format="value(name)" \
           --project="${PROJECT_ID}")
   
           if [[ -z "${rules}" ]]; then
             echo "No firewall rules remain for ${NW}"
             break
           fi
   
           for rule in ${rules}; do
             echo "Deleting firewall rule ${rule}..."
             gcloud compute firewall-rules delete "${rule}" --project="${PROJECT_ID}" --quiet || true
           done
   
           sleep 3
         done
   
         echo "[$(date)] ========== Deleting network ${NW} =========="
         gcloud compute networks delete ${NW} --quiet || true
   done

Borra tu proyecto

Precaución: Borrar un proyecto tiene las siguientes consecuencias:

• Se borra todo su contenido. Si usaste un proyecto existente para las tareas de este documento, cuando lo borres, también se borrará cualquier otro trabajo que hayas realizado en el proyecto.
• Se pierden los IDs personalizados de los proyectos. Cuando creaste este proyecto, es posible que hayas creado un ID personalizado del proyecto que desees utilizar en el futuro. Para conservar las URLs que usan el ID del proyecto, como una URL appspot.com, borra los recursos seleccionados dentro del proyecto, en lugar de todo el proyecto.

Si planeas explorar varios instructivos, arquitecturas o guías de inicio rápido, reutilizar proyectos puede ser útil para no exceder los límites de cuota de los proyectos.

Borra un Google Cloud proyecto:

gcloud projects delete PROJECT_ID

¿Qué sigue?

• Obtén más información sobre Ray en GKE
• Explora la documentación de NVIDIA NeMo RL