Scalare l'apprendimento per rinforzo con GKE e Managed Lustre

1. Introduzione

Se preferisci eseguire gli script in pacchetto direttamente senza il tutorial passo passo, puoi trovarli nel repository GoogleCloudPlatform/devrel-demos.

In questo codelab imparerai a eseguire il deployment di una pipeline di addestramento ad alte prestazioni per il Reinforcement Learning (RL) utilizzando Google Kubernetes Engine (GKE) e Managed Lustre.

I carichi di lavoro di apprendimento per rinforzo, in particolare quelli che utilizzano algoritmi come l'ottimizzazione relativa delle norme di gruppo (GRPO), generano enormi quantità di dati durante la "generazione di esperienze" e richiedono un checkpointing frequente. L'archiviazione di oggetti standard può causare colli di bottiglia durante questi burst di I/O, lasciando inattivi gli acceleratori costosi.

Utilizzerai Managed Lustre, un file system parallelo, per eliminare questi colli di bottiglia e ottenere una velocità effettiva di addestramento superiore.

In questo lab proverai a:

• Configura le variabili di ambiente per un cluster Ray basato su GPU.
• Esegui il provisioning di un cluster GPU spot su GKE e di un'istanza Managed Lustre utilizzando Cluster Toolkit.
• Esegui il deployment di un cluster KubeRay e monta il file system Lustre.
• Invia un carico di lavoro di addestramento NeMo-RL.
• Osserva un throughput elevato e una bassa latenza del checkpoint utilizzando Cloud Monitoring.

Che cosa ti serve

• Un browser web come Chrome.
• Un progetto Google Cloud con la fatturazione abilitata.

Questo codelab utilizza risorse di fascia alta, tra cui GPU spot e Managed Lustre. Sebbene utilizziamo la capacità spot per ridurre i costi, queste risorse possono essere costose se lasciate in esecuzione. Segui sempre il passaggio di pulizia alla fine del lab per evitare addebiti imprevisti.

Questo codelab è destinato a utenti tecnici avanzati, ingegneri di piattaforma e ricercatori di AI che hanno familiarità con GKE e i concetti di archiviazione.

Durata totale stimata: 45-60 minuti più 2 ore di tempo di addestramento

2. Prima di iniziare

Crea un progetto Google Cloud

1. Nella console Google Cloud, seleziona o crea un progetto Google Cloud.
2. Verifica che la fatturazione sia attivata per il tuo progetto Cloud.

Avvia Cloud Shell

Cloud Shell è un ambiente a riga di comando in esecuzione in Google Cloud che viene precaricato con gli strumenti necessari.

1. Fai clic su Attiva Cloud Shell nella parte superiore della console Google Cloud.
2. Una volta connesso a Cloud Shell, verifica l'autenticazione:

   gcloud auth list
   

3. Verifica che il progetto sia configurato:

   gcloud config get project
   

4. Se il progetto non è impostato come previsto, impostalo:

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

Installa Cluster Toolkit

Questo codelab utilizza Cluster Toolkit (gcluster) per eseguire il deployment del cluster GKE. Per istruzioni su come configurare Cluster Toolkit, consulta la guida alla configurazione di Cluster Toolkit.

Abilita API

Esegui questo comando in Cloud Shell per abilitare tutte le API richieste:

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

3. Configura le variabili di ambiente

Per mantenere la coerenza dei comandi in questo codelab, configura alcune variabili di ambiente.

Crea un file denominato env.sh e inserisci la configurazione. Puoi utilizzare il seguente modello:

# 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 è il token di accesso utente Hugging Face, necessario per scaricare i modelli. Puoi ottenerne uno seguendo le istruzioni riportate in Token di sicurezza di Hugging Face.

Sostituisci <YOUR_PROJECT_ID> e <YOUR_HF_TOKEN> con i tuoi valori effettivi.

WANDB_API_KEY è facoltativo. Viene utilizzato per Weights & Biases, uno strumento per monitorare e visualizzare gli esperimenti di machine learning. Se non hai un account o non vuoi utilizzarlo, puoi lasciarlo vuoto o così com'è.

Recupera il file per caricare le variabili nella sessione corrente:

source env.sh

4. Esegui il deployment del cluster GKE e di Managed Lustre utilizzando Cluster Toolkit

In questo passaggio, utilizzi Cluster Toolkit (gcluster) per eseguire il deployment di un cluster GKE con GPU spot e il provisioning automatico dello spazio di archiviazione Managed Lustre con il driver CSI Lustre e PersistentVolumeClaim (lustre-pvc) preconfigurato.

Preparare il progetto

Prima del deployment, esamina il examples/gke-a4/gke-a4.yamlblueprint (per i dettagli, vedi Creare un cluster A4):

1. Abilita Managed Lustre: rimuovi il commento dalle sezioni dei moduli managed-lustre e lustre-pvc in gke-a4.yaml.
2. Attiva il componente aggiuntivo RayOperator: imposta enable_ray_operator: true nelle impostazioni del modulo gke_cluster in gke-a4.yaml.

Esegui il deployment dell'infrastruttura

Imposta il CIDR autorizzato per l'accesso a Cloud Shell ed esegui il deployment utilizzando 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

L'impostazione spot=true in --vars esegue il provisioning dinamico del pool di nodi GPU utilizzando la capacità Spot senza richiedere modifiche manuali a reservation_affinity. Tieni presente che la capacità spot può essere prerilasciata. Per gli ambienti di produzione, ometti spot=true per utilizzare la capacità on demand o configura una prenotazione.

Attendi il completamento del deployment. Cluster Toolkit esegue automaticamente il provisioning della rete VPC, del peering dell'accesso privato ai servizi (PSA), del file system Managed Lustre, del driver CSI Lustre, del componente aggiuntivo RayOperator e della richiesta di archiviazione Kubernetes (lustre-pvc) in un unico deployment coordinato.

5. Esegui il deployment di un cluster Ray su GKE

In questo passaggio, eseguirai il deployment di un cluster KubeRay sui nodi GKE e monterai il file system Lustre utilizzando PersistentVolumeClaim (lustre-pvc) di cui viene eseguito automaticamente il provisioning da Cluster Toolkit.

Crea configurazione RayCluster

Crea un file denominato ray-cluster.yaml. Specifica i nodi head e worker di KubeRay, utilizzando il tipo di acceleratore nvidia-b200 e montando il volume Lustre in /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

Connettiti al cluster

Assicurati che la sessione di Cloud Shell sia autenticata nel cluster GKE:

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

Applica la configurazione di RayCluster

Applica la configurazione del cluster Ray:

kubectl apply -f ray-cluster.yaml

Verifica lo stato del cluster

Monitora la creazione dei pod:

kubectl get pods -w

Attendi che i pod head e worker siano Running.

6. Inviare il workload di apprendimento per rinforzo

In questo passaggio, invierai il job di addestramento NeMo-RL GRPO al tuo cluster Ray.

Connettersi alla dashboard Ray

Per inviare job e visualizzare le metriche, devi connetterti alla dashboard di Ray. Poiché la dashboard si trova in GKE, utilizza il port forwarding per accedervi da 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 lo script di esecuzione

Crea un file denominato run_nemo_rl.sh. Questo script verrà eseguito sui worker del cluster Ray. Utilizziamo cat << EOF per compilare le variabili di ambiente che hai impostato in precedenza.

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 file di ignoranza di Ray

Crea un file .rayignore per impedire a Ray di caricare directory di grandi dimensioni o non necessarie:

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

Crea configurazione dell'ambiente di runtime

Crea un file JSON per passare le variabili di ambiente al job 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

Invia il job

Utilizza Ray CLI per inviare il job all'endpoint della dashboard. Se il comando ray non viene trovato in Cloud Shell, puoi installarlo 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

Vedrai lo streaming dei log nel terminale Cloud Shell. Il job caricherà il modello, inizializzerà i worker Ray e inizierà il ciclo di addestramento GRPO.

7. Monitorare il rendimento dell'addestramento

In questo passaggio, osserverai le prestazioni del file system Lustre durante l'addestramento e il checkpoint.

Controllare i log di addestramento

Man mano che l'addestramento procede, vedrai i log che indicano che i checkpoint vengono salvati in /lustre/nemo_rl_qwen_72b_ds_cp/nemo-rl-grpo-test1. Tieni presente che il checkpointing avviene in modo asincrono e non blocca i worker Ray per molto tempo.

Per visualizzare la velocità di checkpointing, cerca le righe del log che indicano i checkpoint salvati.

Visualizzare le metriche di Lustre nella console Cloud

Per visualizzare le metriche della tua istanza Lustre:

1. Nella console Google Cloud, cerca Managed Service for Lustre.
2. Fai clic sul nome dell'istanza (ad esempio, ${CLUSTER_NAME}-lustre o rl-demo-gpu-lustre).
3. Fai clic sulla scheda Monitoraggio.

Qui puoi osservare:

• Throughput (byte/sec): visualizza i picchi durante la creazione dei checkpoint.
• Capacità: monitora la quantità di spazio utilizzata dai checkpoint.

Lustre è in grado di scrivere a velocità molto elevata, scrivendo i checkpoint in un tempo minimo

8. Pulizia delle risorse

Questo passaggio è fondamentale per evitare addebiti continui. Se lasciati in esecuzione, i cluster GPU e Managed Lustre possono comportare costi significativi.

Esegui questo comando in Cloud Shell per eliminare tutta l'infrastruttura di cui è stato eseguito il provisioning (cluster GKE, pool di nodi GPU, istanza Managed Lustre e rete VPC) in un unico passaggio:

gcluster destroy "${CLUSTER_NAME}"

Questo comando viene eseguito in modo sincrono in primo piano, smantellando tutta l'infrastruttura gestita dal deployment e restituendo i log di avanzamento al terminale. Attendi il completamento del comando prima di chiudere la sessione di Cloud Shell.

9. Complimenti

Hai completato correttamente il codelab Scalare il reinforcement learning con GKE e Managed Lustre.

Cosa hai imparato

• Come utilizzare Cluster Toolkit per eseguire il provisioning di un cluster GPU GKE con istanze Spot e spazio di archiviazione Managed Lustre.
• Come eseguire il deployment di un cluster KubeRay e montare lo spazio di archiviazione Lustre.
• Come inviare un workload di addestramento GRPO NeMo-RL.
• Come osservare il rendimento dello spazio di archiviazione durante l'addestramento.

Passaggi successivi

• Esplora altre funzionalità di NVIDIA NeMo-RL.
• Scopri di più su Google Cloud AI Hypercomputer.
• Consulta la documentazione di Managed Service for Lustre.