Escala el aprendizaje por refuerzo con GKE y Managed Lustre

1. Introducción

Si prefieres ejecutar los secuencias de comandos empaquetados directamente sin el instructivo paso a paso, puedes encontrarlos en el repositorio GoogleCloudPlatform/devrel-demos.

En este codelab, aprenderás a implementar una canalización de entrenamiento de alto rendimiento para el aprendizaje por refuerzo (RL) con Google Kubernetes Engine (GKE) y Lustre administrado.

Las cargas de trabajo de aprendizaje por refuerzo, en especial las que usan algoritmos como la optimización de políticas relativas al grupo (GRPO), generan grandes cantidades de datos durante la "generación de experiencias" y requieren la creación frecuente de puntos de control. El almacenamiento de objetos estándar puede causar cuellos de botella durante estas ráfagas de E/S, lo que deja inactivos los aceleradores costosos.

Usarás Managed Lustre, un sistema de archivos paralelo, para eliminar estos cuellos de botella y lograr una mayor capacidad de procesamiento del entrenamiento.

Actividades

• Configura las variables de entorno para un clúster de Ray basado en GPU.
• Aprovisiona un clúster de GPU Spot en GKE y una instancia de Managed Lustre con Cluster Toolkit.
• Implementa un clúster de KubeRay y activa el sistema de archivos Lustre.
• Envía una carga de trabajo de entrenamiento de NeMo-RL.
• Observa una alta capacidad de procesamiento y una baja latencia de puntos de control con Cloud Monitoring.

Requisitos

• Un navegador web, como Chrome
• Un proyecto de Google Cloud con facturación habilitada.

En este codelab, se usan recursos de alta gama, como GPUs Spot y Managed Lustre. Si bien usamos la capacidad de Spot para reducir los costos, estos recursos pueden ser costosos si se dejan en ejecución. Sigue siempre el paso de limpieza al final del lab para evitar cargos inesperados.

Este codelab está dirigido a usuarios técnicos avanzados, ingenieros de plataformas e investigadores de IA que estén familiarizados con los conceptos de GKE y almacenamiento.

Duración total estimada: 45 a 60 minutos más 2 horas de capacitación

2. Antes de comenzar

Crea un proyecto de Google Cloud

1. En la consola de Google Cloud, selecciona o crea un proyecto de Google Cloud.
2. Asegúrate de que la facturación esté habilitada para tu proyecto de Cloud.

Inicie Cloud Shell

Cloud Shell es un entorno de línea de comandos que se ejecuta en Google Cloud y que viene precargado con las herramientas necesarias.

1. Haz clic en Activar Cloud Shell en la parte superior de la consola de Google Cloud.
2. Una vez que te conectes a Cloud Shell, verifica tu autenticación:

   gcloud auth list
   

3. Confirma que tu proyecto esté configurado:

   gcloud config get project
   

4. Si tu proyecto no está configurado como se esperaba, configúralo:

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

Instala Cluster Toolkit

En este codelab, se usa Cluster Toolkit (gcluster) para implementar el clúster de GKE. Para obtener instrucciones sobre cómo configurar Cluster Toolkit, consulta la guía de configuración de Cluster Toolkit.

Habilita las APIs

Ejecuta este comando en Cloud Shell para habilitar todas las APIs requeridas:

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

3. Configura las variables de entorno

Para que los comandos de este codelab sean coherentes, configura algunas variables de entorno.

Crea un archivo llamado env.sh y complétalo con tu configuración. Puedes usar la siguiente plantilla:

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

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

HF_TOKEN es tu token de acceso de usuario de Hugging Face, que se requiere para descargar modelos. Para obtener uno, sigue las instrucciones en Tokens de seguridad de Hugging Face.

Reemplaza <YOUR_PROJECT_ID> y <YOUR_HF_TOKEN> por tus valores reales.

WANDB_API_KEY es opcional. Se usa para Weights & Biases, una herramienta para supervisar y visualizar experimentos de aprendizaje automático. Si no tienes una cuenta o no quieres usarla, puedes dejar el campo en blanco o tal como está.

Obtén el archivo para cargar las variables en tu sesión actual:

source env.sh

4. Implementa un clúster de GKE y Managed Lustre con Cluster Toolkit

En este paso, usarás Cluster Toolkit (gcluster) para implementar un clúster de GKE con GPUs Spot y aprovisionar automáticamente el almacenamiento de Managed Lustre con el controlador CSI de Lustre y PersistentVolumeClaim preconfigurado (lustre-pvc).

Prepara el esquema

Antes de implementar, revisa el blueprint de examples/gke-a4/gke-a4.yaml (para obtener más detalles, consulta Crea un clúster de A4):

1. Habilita Managed Lustre: Quita las marcas de comentario de las secciones de los módulos managed-lustre y lustre-pvc en gke-a4.yaml.
2. Habilita el complemento RayOperator: Establece enable_ray_operator: true en la configuración del módulo gke_cluster en gke-a4.yaml.

Implementa la infraestructura

Establece tu CIDR autorizado para el acceso a Cloud Shell y realiza la implementación con gcluster deploy:

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

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

El parámetro de configuración spot=true en --vars aprovisiona dinámicamente el grupo de nodos de GPU con capacidad de Spot sin necesidad de realizar ediciones manuales en reservation_affinity. Ten en cuenta que la capacidad de Spot se puede interrumpir. Para los entornos de producción, omite spot=true para usar la capacidad bajo demanda o configurar una reserva.

Espera a que se complete la implementación. Cluster Toolkit aprovisiona automáticamente la red de VPC, el intercambio de tráfico de acceso a servicios privados (PSA), el sistema de archivos Managed Lustre, el controlador CSI de Lustre, el complemento RayOperator y el reclamo de almacenamiento de Kubernetes (lustre-pvc) en una implementación coordinada.

5. Implementa el clúster de Ray en GKE

En este paso, implementarás un clúster de KubeRay en tus nodos de GKE y activarás el sistema de archivos Lustre con el PersistentVolumeClaim (lustre-pvc) que Cluster Toolkit aprovisionó automáticamente.

Crea la configuración de RayCluster

Crea un archivo llamado ray-cluster.yaml. Esto especifica los nodos principales y trabajadores de KubeRay, usando el tipo de acelerador nvidia-b200 y el montaje del volumen de Lustre en /lustre.

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

Conéctate al clúster

Asegúrate de que tu sesión de Cloud Shell esté autenticada en tu clúster de GKE:

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

Aplica la configuración de RayCluster

Aplica la configuración del clúster de Ray:

kubectl apply -f ray-cluster.yaml

Verifica el estado del clúster

Supervisa la creación de los Pods:

kubectl get pods -w

Espera hasta que los Pods principales y de trabajador estén en estado Running.

6. Envía la carga de trabajo de aprendizaje por refuerzo

En este paso, enviarás el trabajo de entrenamiento de GRPO de NeMo-RL a tu clúster de Ray.

Conéctate al panel de Ray

Para enviar trabajos y ver métricas, debes conectarte al panel de Ray. Como el panel está en GKE, usa el reenvío de puertos para acceder a él desde Cloud Shell:

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

Crea el script de ejecución

Crea un archivo llamado run_nemo_rl.sh. Esta secuencia de comandos se ejecutará en los trabajadores del clúster de Ray. Usamos cat << EOF para completar las variables de entorno que configuraste antes.

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

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

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

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

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

Crea un archivo .rayignore

Crea un archivo .rayignore para evitar que Ray suba directorios grandes o innecesarios:

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

Crea la configuración del entorno de ejecución

Crea un archivo JSON para pasar variables de entorno al trabajo de Ray:

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

Envía el trabajo

Usa la CLI de Ray para enviar el trabajo al extremo del panel. Si no se encuentra el comando ray en Cloud Shell, puedes instalarlo con pip install ray:

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

Verás registros transmitiéndose en tu terminal de Cloud Shell. El trabajo cargará el modelo, inicializará los trabajadores de Ray y comenzará el bucle de entrenamiento del GRPO.

7. Supervisa el rendimiento del entrenamiento

En este paso, observarás el rendimiento del sistema de archivos Lustre durante el entrenamiento y la creación de puntos de control.

Verifica los registros de entrenamiento

A medida que avance el entrenamiento, verás registros que indican que los puntos de control se guardan en /lustre/nemo_rl_qwen_72b_ds_cp/nemo-rl-grpo-test1. Ten en cuenta que la creación de puntos de control se realiza de forma asíncrona y no bloquea los trabajadores de Ray durante mucho tiempo.

Para ver la velocidad de los puntos de control, busca líneas de registro que indiquen que se guardaron puntos de control.

Cómo ver las métricas de Lustre en Cloud Console

Para ver las métricas de tu instancia de Lustre, sigue estos pasos:

1. En la consola de Google Cloud, busca Managed Service for Lustre.
2. Haz clic en el nombre de tu instancia (por ejemplo, ${CLUSTER_NAME}-lustre o rl-demo-gpu-lustre).
3. Haz clic en la pestaña Supervisión.

Aquí puedes observar lo siguiente:

• Capacidad de procesamiento (bytes/s): Observa los picos durante la creación de puntos de control.
• Capacidad: Supervisa cuánto espacio consumen los puntos de control.

Lustre puede escribir a una velocidad muy alta y escribir puntos de control en un tiempo mínimo

8. Limpia los recursos

Este paso es fundamental para evitar cargos continuos. Si se dejan en ejecución, los clústeres de GPU y Managed Lustre pueden generar costos significativos.

Ejecuta el siguiente comando en Cloud Shell para destruir toda la infraestructura aprovisionada (clúster de GKE, grupos de nodos de GPU, instancia de Managed Lustre y red de VPC) en un solo paso:

gcluster destroy "${CLUSTER_NAME}"

Este comando se ejecuta de forma síncrona en primer plano, lo que desmantela toda la infraestructura administrada por la implementación y genera registros de progreso en tu terminal. Espera a que el comando finalice por completo antes de cerrar tu sesión de Cloud Shell.

9. Felicitaciones

Completaste correctamente el codelab Scale Reinforcement Learning with GKE and Managed Lustre.

Qué aprendiste

• Cómo usar Cluster Toolkit para aprovisionar un clúster de GPU de GKE con instancias Spot y almacenamiento de Managed Lustre
• Cómo implementar un clúster de KubeRay y activar el almacenamiento de Lustre
• Cómo enviar una carga de trabajo de entrenamiento de GRPO de NeMo-RL
• Cómo observar el rendimiento del almacenamiento durante el entrenamiento

Próximos pasos

• Explora más funciones de NVIDIA NeMo-RL.
• Obtén más información sobre AI Hypercomputer de Google Cloud.
• Revisa la documentación del Servicio administrado para Lustre.