Google utilise la technologie IA pour traduire le contenu dans votre langue préférée. Les traductions générées par IA peuvent contenir des erreurs.

Mettre à l'échelle l'apprentissage par renforcement avec GKE et Managed Lustre

1. Introduction

Si vous préférez exécuter les scripts packagés directement sans le tutoriel pas à pas, vous les trouverez dans le dépôt GoogleCloudPlatform/devrel-demos.

Dans cet atelier de programmation, vous allez apprendre à déployer un pipeline d'entraînement hautes performances pour le Reinforcement Learning (RL) à l'aide de Google Kubernetes Engine (GKE) et de Managed Lustre.

Les charges de travail d'apprentissage par renforcement, en particulier celles qui utilisent des algorithmes tels que l'optimisation des stratégies relatives aux groupes (GRPO, Group Relative Policy Optimization), génèrent d'énormes quantités de données lors de la "génération d'expérience" et nécessitent des points de contrôle fréquents. Le stockage d'objets standard peut entraîner des goulots d'étranglement lors de ces pics d'E/S, laissant les coûteux accélérateurs inactifs.

Vous utiliserez Managed Lustre, un système de fichiers parallèle, pour éliminer ces goulots d'étranglement et obtenir un débit d'entraînement plus élevé.

Objectifs de l'atelier

Configurez des variables d'environnement pour un cluster Ray basé sur GPU.
Provisionnez un cluster GPU Spot sur GKE et une instance Managed Lustre à l'aide de Cluster Toolkit.
Déployez un cluster KubeRay et installez le système de fichiers Lustre.
Envoyez une charge de travail d'entraînement NeMo-RL.
Observez le haut débit et la faible latence des points de contrôle à l'aide de Cloud Monitoring.

Diagramme de l'architecture de GKE, KubeRay et Managed Lustre

Prérequis

Un navigateur Web tel que Chrome.
Un projet Google Cloud avec facturation activée.

Cet atelier de programmation s'adresse aux utilisateurs techniques avancés, aux ingénieurs de plate-forme et aux chercheurs en IA qui connaissent les concepts de GKE et de stockage.

Durée totale estimée : 45 à 60 minutes plus 2 heures de formation

2. Avant de commencer

Créer un projet Google Cloud

Dans la console Google Cloud, sélectionnez ou créez un projet Google Cloud.
Assurez-vous que la facturation est activée pour votre projet Cloud.

Démarrer Cloud Shell

Cloud Shell est un environnement de ligne de commande exécuté dans Google Cloud et fourni avec les outils nécessaires.

Cliquez sur Activer Cloud Shell en haut de la console Google Cloud.
Une fois connecté à Cloud Shell, vérifiez votre authentification :
```
gcloud auth list
```
Vérifiez que votre projet est configuré :
```
gcloud config get project
```

Si votre projet n'est pas défini comme prévu, définissez-le :

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

Installer Cluster Toolkit

Cet atelier de programmation utilise Cluster Toolkit (gcluster) pour déployer le cluster GKE. Pour savoir comment configurer Cluster Toolkit, consultez le guide de configuration de Cluster Toolkit.

Activer les API

Exécutez cette commande dans Cloud Shell pour activer toutes les API requises :

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

3. Configurer les variables d'environnement

Pour que les commandes de cet atelier de programmation soient cohérentes, configurez quelques variables d'environnement.

Créez un fichier nommé env.sh et renseignez-le avec votre configuration. Vous pouvez utiliser le modèle suivant :

# 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

Remplacez <YOUR_PROJECT_ID> et <YOUR_HF_TOKEN> par vos valeurs réelles.

Importez le fichier pour charger les variables dans votre session actuelle :

source env.sh

4. Déployer un cluster GKE et Managed Lustre à l'aide de Cluster Toolkit

Dans cette étape, vous allez utiliser Cluster Toolkit (gcluster) pour déployer un cluster GKE avec des GPU Spot et provisionner automatiquement le stockage Lustre géré avec le pilote CSI Lustre et PersistentVolumeClaim (lustre-pvc) préconfiguré.

Préparer le plan

Avant de déployer le plan examples/gke-a4/gke-a4.yaml, consultez-le (pour en savoir plus, consultez Créer un cluster A4) :

Activer Managed Lustre : annulez la mise en commentaire des sections de module managed-lustre et lustre-pvc dans gke-a4.yaml.
Activer le module complémentaire RayOperator : définissez enable_ray_operator: true sous les paramètres du module gke_cluster dans gke-a4.yaml.

Déployer l'infrastructure

Définissez votre CIDR autorisé pour l'accès à Cloud Shell et déployez-le à l'aide de 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

Attendez la fin du déploiement. Cluster Toolkit provisionne automatiquement le réseau VPC, l'appairage Private Service Access (PSA), le système de fichiers Lustre géré, le pilote CSI Lustre, le module complémentaire RayOperator et la revendication de stockage Kubernetes (lustre-pvc) dans un déploiement coordonné.

5. Déployer un cluster Ray sur GKE

Dans cette étape, vous allez déployer un cluster KubeRay sur vos nœuds GKE et monter le système de fichiers Lustre à l'aide du PersistentVolumeClaim (lustre-pvc) provisionné automatiquement par Cluster Toolkit.

Créer une configuration RayCluster

Créez un fichier nommé ray-cluster.yaml. Cela spécifie les nœuds principaux et de calcul KubeRay, en utilisant le type d'accélérateur nvidia-b200 et en montant le volume Lustre à /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

Se connecter au cluster

Assurez-vous que votre session Cloud Shell est authentifiée auprès de votre cluster GKE :

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

Appliquer la configuration RayCluster

Appliquez la configuration du cluster Ray :

kubectl apply -f ray-cluster.yaml

Vérifier l'état du cluster

Surveillez la création des pods :

kubectl get pods -w

Attendez que les pods principaux et de nœuds de calcul soient Running.

6. Envoyer une charge de travail d'apprentissage par renforcement

Dans cette étape, vous allez envoyer le job d'entraînement GRPO NeMo-RL à votre cluster Ray.

Se connecter au tableau de bord Ray

Pour envoyer des jobs et afficher des métriques, vous devez vous connecter au tableau de bord Ray. Comme le tableau de bord se trouve dans GKE, utilisez le transfert de port pour y accéder depuis Cloud Shell :

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

Créer le script d'exécution

Créez un fichier nommé run_nemo_rl.sh. Ce script sera exécuté sur les nœuds de calcul du cluster Ray. Nous utilisons cat << EOF pour renseigner les variables d'environnement que vous avez définies précédemment.

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

Créer un fichier Ray Ignore

Créez un fichier .rayignore pour empêcher Ray d'importer des répertoires volumineux ou inutiles :

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

Créer une configuration d'environnement d'exécution

Créez un fichier JSON pour transmettre des variables d'environnement au 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

Envoyer le job

Utilisez la CLI Ray pour envoyer le job au point de terminaison du tableau de bord. Si la commande ray est introuvable dans Cloud Shell, vous pouvez l'installer avec pip install ray :

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

Vous verrez les journaux s'afficher dans votre terminal Cloud Shell. Le job chargera le modèle, initialisera les nœuds de calcul Ray et lancera la boucle d'entraînement GRPO.

7. Surveiller les performances d'entraînement

Dans cette étape, vous allez observer les performances du système de fichiers Lustre pendant l'entraînement et la création de points de contrôle.

Consulter les journaux d'entraînement

Au fur et à mesure de l'entraînement, des journaux indiquent que des points de contrôle sont enregistrés dans /lustre/nemo_rl_qwen_72b_ds_cp/nemo-rl-grpo-test1. Notez que la création de points de contrôle se fait de manière asynchrone et ne bloque pas les workers Ray très longtemps.

Pour afficher la vitesse de checkpointing, recherchez les lignes de journal indiquant les points de contrôle enregistrés.

Afficher les métriques Lustre dans la console Cloud

Pour afficher les métriques de votre instance Lustre :

Dans la console Google Cloud, recherchez Managed Service for Lustre.
Cliquez sur le nom de votre instance (par exemple, ${CLUSTER_NAME}-lustre ou rl-demo-gpu-lustre).
Cliquez sur l'onglet Surveillance.

Vous pouvez y observer les éléments suivants :

Débit (octets/s) : observez les pics pendant la création de points de contrôle.
Capacité : surveillez l'espace consommé par les points de contrôle.

Graphique des performances de Lustre Lustre est capable d'écrire à très grande vitesse et d'écrire des points de contrôle en un minimum de temps.

8. Nettoyer les ressources

Exécutez la commande suivante dans Cloud Shell pour détruire toute l'infrastructure provisionnée (cluster GKE, pools de nœuds GPU, instance Managed Lustre et réseau VPC) en une seule étape :

gcluster destroy "${CLUSTER_NAME}"

Cette commande s'exécute de manière synchrone au premier plan, en supprimant toute l'infrastructure gérée par le déploiement et en affichant les journaux de progression dans votre terminal. Attendez que la commande soit entièrement exécutée avant de fermer votre session Cloud Shell.

9. Félicitations

Vous avez terminé l'atelier de programmation Faire évoluer l'apprentissage par renforcement avec GKE et Managed Lustre.

Connaissances acquises

Découvrez comment utiliser Cluster Toolkit pour provisionner un cluster GPU GKE avec des instances Spot et un stockage Lustre géré.
Découvrez comment déployer un cluster KubeRay et monter le stockage Lustre.
Comment envoyer une charge de travail d'entraînement NeMo-RL GRPO.
Comment observer les performances de stockage pendant l'entraînement.

Étapes suivantes

Découvrez d'autres fonctionnalités de NVIDIA NeMo-RL.
En savoir plus sur Google Cloud AI Hypercomputer
Consultez la documentation du service géré pour Lustre.