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

Autopilot 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). Usarás 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, en su lugar, usa un cálculo relativo basado en grupos.

Antes de ejecutar este instructivo, te recomendamos que completes el instructivo Ajusta y escala el aprendizaje por refuerzo con Verl en GKE. En el siguiente instructivo, se usan la misma configuración y el mismo 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, el 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 de matemáticas 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 correctas y lógicas de forma constante 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:

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

Antes de comenzar

Accede a tu cuenta de Google Cloud . Si eres nuevo en Google Cloud, crea una cuenta para evaluar el rendimiento de nuestros productos en situaciones reales. Los clientes nuevos también obtienen $300 en créditos gratuitos para ejecutar, probar y, además, implementar cargas de trabajo.

Instala Google Cloud CLI.

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

Si usas un proveedor de identidad externo (IdP), primero debes acceder a la gcloud CLI con tu identidad federada.

Para inicializar gcloud CLI, ejecuta el siguiente comando:

gcloud init

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.

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 Google Cloud proyecto.

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

Habilita las APIs necesarias:

Roles necesarios para habilitar las APIs

Para habilitar las APIs, necesitas el rol de IAM de administrador de Service Usage (roles/serviceusage.serviceUsageAdmin), que contiene el permiso serviceusage.services.enable. Obtén más información para otorgar roles.

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

Instala Google Cloud CLI.

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

Si usas un proveedor de identidad externo (IdP), primero debes acceder a la gcloud CLI con tu identidad federada.

Para inicializar gcloud CLI, ejecuta el siguiente comando:

gcloud init

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.

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 Google Cloud proyecto.

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

Habilita las APIs necesarias:

Roles necesarios para habilitar las APIs

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

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 . Por ejemplo, myemail@example.com.
- ROLE: Es el rol de IAM que otorgas a tu cuenta de usuario.

Crea una cuenta de Hugging Face, si todavía no la tienes.
Asegúrate de tener un token de Hugging Face.
Asegúrate de que tu proyecto tenga la cuota suficiente 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.

Ve a la consola deGoogle Cloud .
En la parte superior de la Google Cloud ventana de la consola, haz clic en el botón Activar Cloud Shell.
Configura las siguientes variables de entorno:
```
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
```
Reemplaza los siguientes valores:
- CLUSTER_NAME: Es el nombre del clúster de GKE.
- CONTROL_PLANE_LOCATION: Es la región de Compute Engine para el plano de control del clúster de GKE.
- NODE_LOCATION: Es la ubicación de tus nodos. Selecciona una zona en la que estén disponibles las GPUs NVIDIA B200 o H200.
- 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)
- 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.
- GKE_VERSION: Es la versión de GKE que se usará:
  - Para las GPU NVIDIA B200 (180 GB), usa 1.32.2-gke.1422000 o una versión posterior.
  - Para las GPU NVIDIA H200 (141 GB), usa 1.31.4-gke.1183000 o una versión posterior.
- BUCKET_NAME: Es el nombre base de tu bucket de Cloud Storage.
- YOUR_HUGGING_FACE_TOKEN: Tu token de Hugging Face.
Crea las siguientes variables de entorno para la red:
```
export GVNIC_NETWORK_PREFIX="GVNIC-NAME"
export RDMA_NETWORK_PREFIX="RDMA-NAME"
```
Reemplaza los siguientes valores:
- 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

Crea una red de VPC para la interfaz de 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

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=${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 el clúster de GKE

Puedes configurar NeMo RL en un clúster de GKE Autopilot o Standard. Te recomendamos que uses un clúster de Autopilot para una experiencia de Kubernetes completamente administrada. Para elegir el modo de operación de GKE que se adapte mejor a tus cargas de trabajo, consulta Acerca de los modos de operación de GKE.

Autopilot

Crea un clúster de Autopilot:

gcloud container clusters create-auto ${CLUSTER_NAME} \
    --location=${CONTROL_PLANE_LOCATION} \
    --enable-multi-networking  \
    --enable-ray-operator

Obtén credenciales para el clúster:

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

Instala el instalador de RDMA de NCCL para Autopilot:

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

Estándar

Crea un clúster estándar:

gcloud container clusters create ${CLUSTER_NAME} \
    --location=${CONTROL_PLANE_LOCATION} \
    --enable-dataplane-v2 \
    --enable-ip-alias \
    --enable-multi-networking \
    --addons=RayOperator \
    --num-nodes=1

Obtén credenciales para el clúster:

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

Crea el grupo de nodos de 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

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

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

Aplica el manifiesto
```
kubectl apply -f network-mapping.yaml
```

Prepara el almacenamiento

En esta sección, crearás buckets de Cloud Storage y una instancia de Lustre administrada, que aprovisiona el almacenamiento de alto rendimiento necesario para tu carga de trabajo de AA.

Crea un bucket de Cloud Storage:

gcloud storage buckets create gs://${GS_BUCKET} \
    --location=${CONTROL_PLANE_LOCATION} \
    --enable-hierarchical-namespace \
    --uniform-bucket-level-access

Crea una cuenta de servicio de Kubernetes (KSA) y vincúlala 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"

Para configurar Managed Lustre, completa los siguientes pasos:
1. Crea una instancia de Managed Lustre siguiendo los pasos que se indican en Crea una instancia de Managed Lustre. Asegúrate de que la instancia use la misma red que tu clúster de GKE.
2. Sigue los pasos que se indican en Cómo acceder a una instancia de Managed Lustre existente para acceder a la instancia de Managed Lustre.

Implementa RayCluster

En esta sección, clonarás el repositorio de muestra, prepararás los manifiestos y ejecutarás una secuencia de comandos launcher.sh:

Clona el repositorio de ejemplo:

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

Navega hasta el directorio de trabajo:
```
cd ai-ml/nemo-rl-on-gke/nemoRL
```

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

Instala el clúster de Ray:
```
export REPLICA_COUNT=2
helm install ray-cluster . \
  --set values.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.
Para implementar el clúster de Ray, ejecuta la secuencia de comandos launcher.sh:
```
bash launcher.sh
```

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

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

kubectl ray get clusters

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.

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.

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

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: 64: El modelo Gemma3-27b-it genera un grupo grande de respuestas para cada instrucción. En esta configuración, el modelo produce 1,024 respuestas en total (16 × 64 = 1,024).
policy.generation.colocated.enabled=False: Este parámetro inhabilita la función de generación de ubicación conjunta, 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. Si separas estas cargas de trabajo, evitarás 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.

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.

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.

Instala la utilidad de árbol apt:
```
apt install tree
```

Para supervisar el estado del trabajo de GRPO, verifica los registros del nodo principal de Ray:

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

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

Realiza una limpieza

Para evitar que se generen cobros, borra los recursos:

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

Ajusta y escala el aprendizaje por refuerzo con NVIDIA NeMo RL en GKE Organiza tus páginas con colecciones Guarda y categoriza el contenido según tus preferencias.

Fondo

Aprendizaje por refuerzo (RL)

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

NVIDIA NeMo RL

Conjunto de datos de GSM8k

Objetivos

Antes de comenzar

Prepara el entorno

Configura la infraestructura

Crear red de VPC

Crea el clúster de GKE

Autopilot

Estándar

Configura las asignaciones de red

Prepara el almacenamiento

Implementa RayCluster

Inicia el trabajo de GRPO

Supervisa el estado del trabajo de GRPO

Realiza una limpieza

¿Qué sigue?

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