Ajuster et mettre à l'échelle l'apprentissage par renforcement avec NVIDIA NeMo RL sur GKE

Autopilot Standard

Ce tutoriel explique comment orchestrer un environnement d'entraînement distribué pour l'apprentissage par renforcement (RL) sur Google Kubernetes Engine (GKE). Vous utilisez Ray et le framework NVIDIA NeMo RL pour configurer un environnement d'entraînement distribué afin d'affiner un modèle.

Ce tutoriel se concentre sur le pipeline d'entraînement Group Relative Policy Optimization (GRPO) sur GKE avec Ray et NeMo RL. GRPO est un algorithme d'apprentissage par renforcement conçu pour améliorer la capacité de raisonnement d'un modèle. Cet algorithme économe en mémoire simplifie le processus de RL en éliminant le Critic, ou modèle de valeur, et en utilisant à la place un calcul relatif basé sur les groupes.

Avant de suivre ce tutoriel, nous vous recommandons de suivre le tutoriel Ajuster et mettre à l'échelle l'apprentissage par renforcement avec Verl sur GKE. Le tutoriel suivant utilise la même configuration de cluster que le tutoriel sur le scaling et le fine-tuning du RL avec Verl.

Arrière-plan

Les sections suivantes présentent brièvement les concepts utilisés dans ce tutoriel.

Apprentissage par renforcement

L'apprentissage par renforcement enseigne aux modèles par l'expérience, l'exploration et le retour d'informations plutôt que par l'imitation statique. Bien que le pré-entraînement apprenne à un modèle quoi dire, l'apprentissage par renforcement qui utilise le feedback humain (RLHF) lui apprend à être utile, sûr et logique. L'apprentissage par renforcement sert de pont entre un modèle de base et un modèle affiné pour un cas d'utilisation spécialisé.

Pour en savoir plus, consultez Qu'est-ce que l'apprentissage par renforcement ?

Optimisation des stratégies relatives aux groupes (GRPO)

GRPO, un algorithme popularisé par DeepSeek, offre une alternative économe en mémoire à l'optimisation de la politique proximale (PPO) pour l'alignement des LLM en supprimant le modèle Critic. Au lieu d'un réseau Critic, GRPO génère un groupe de réponses pour la même requête et utilise la récompense moyenne de ce groupe comme référence.

Pour en savoir plus, consultez GRPO.

NVIDIA NeMo RL

NeMo RL est la bibliothèque Open Source de post-entraînement de NVIDIA conçue pour le RL évolutif. Faisant partie de l'écosystème plus vaste du framework NeMo, NeMo RL permet à la fois des expériences à petite échelle sur un seul GPU et des déploiements multinœuds sur des milliers de GPU.

Pour en savoir plus, consultez NVIDIA NeMo RL.

Ensemble de données GSM8k

Dans ce tutoriel, vous allez utiliser l'ensemble de données GSM8k, qui contient 8 500 problèmes de mathématiques de haute qualité, linguistiquement diversifiés et adaptés au niveau de l'école primaire.

En utilisant GSM8k et GRPO, le modèle génère un groupe de n réponses différentes pour le même problème. GRPO compare ces réponses à la moyenne du groupe. Le modèle est davantage récompensé pour les chemins qui sont constamment corrects et logiques par rapport au reste du groupe. Au fil du temps, le modèle apprend que la meilleure façon de maximiser la récompense est d'articuler clairement ses étapes, ce qui réduit efficacement la récompense pour les réponses peu performantes.

Pour en savoir plus, consultez GSM8k.

Objectifs

Ce tutoriel vous explique comment configurer le RL sur GKE avec NeMo RL en suivant les étapes suivantes :

Préparez votre environnement.
Configurez un cluster GKE avec des GPU B200 ou H200.
Configurez KubeRay pour gérer un cluster Ray distribué.
Utilisez Managed Lustre pour un stockage hautes performances.
Exécutez une tâche d'entraînement GRPO qui utilise NeMo RL.

Avant de commencer

Connectez-vous à votre compte Google Cloud . Si vous débutez sur Google Cloud, créez un compte pour évaluer les performances de nos produits en conditions réelles. Les nouveaux clients bénéficient également de 300 $de crédits sans frais pour exécuter, tester et déployer des charges de travail.

Installez la Google Cloud CLI.

Remarque : Si vous avez déjà installé la gcloud CLI, assurez-vous que vous disposez de la dernière version en exécutant gcloud components update.

Si vous utilisez un fournisseur d'identité (IdP) externe, vous devez d'abord vous connecter à la gcloud CLI avec votre identité fédérée.

Pour initialiser la gcloud CLI, exécutez la commande suivante :

gcloud init

Créez ou sélectionnez un projet Google Cloud .

Rôles requis pour sélectionner ou créer un projet

Sélectionnez un projet : la sélection d'un projet ne nécessite pas de rôle IAM spécifique. Vous pouvez sélectionner n'importe quel projet pour lequel un rôle vous a été attribué.
Créer un projet : pour créer un projet, vous devez disposer du rôle Créateur de projet (roles/resourcemanager.projectCreator), qui contient l'autorisation resourcemanager.projects.create. Découvrez comment attribuer des rôles.

Créez un projet Google Cloud :
```
gcloud projects create PROJECT_ID
```
Remplacez PROJECT_ID par le nom du projet Google Cloud que vous créez.
Sélectionnez le projet Google Cloud que vous avez créé :
```
gcloud config set project PROJECT_ID
```
Remplacez PROJECT_ID par le nom de votre projet Google Cloud .

Vérifiez que la facturation est activée pour votre projet Google Cloud .

Activez les API requises :

Rôles requis pour activer les API

Pour activer les API, vous avez besoin du rôle IAM Administrateur Service Usage (roles/serviceusage.serviceUsageAdmin), qui contient l'autorisation serviceusage.services.enable. Découvrez comment attribuer des rôles.

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

Installez la Google Cloud CLI.

Remarque : Si vous avez déjà installé la gcloud CLI, assurez-vous que vous disposez de la dernière version en exécutant gcloud components update.

Si vous utilisez un fournisseur d'identité (IdP) externe, vous devez d'abord vous connecter à la gcloud CLI avec votre identité fédérée.

Pour initialiser la gcloud CLI, exécutez la commande suivante :

gcloud init

Créez ou sélectionnez un projet Google Cloud .

Rôles requis pour sélectionner ou créer un projet

Sélectionnez un projet : la sélection d'un projet ne nécessite pas de rôle IAM spécifique. Vous pouvez sélectionner n'importe quel projet pour lequel un rôle vous a été attribué.
Créer un projet : pour créer un projet, vous devez disposer du rôle Créateur de projet (roles/resourcemanager.projectCreator), qui contient l'autorisation resourcemanager.projects.create. Découvrez comment attribuer des rôles.

Créez un projet Google Cloud :
```
gcloud projects create PROJECT_ID
```
Remplacez PROJECT_ID par le nom du projet Google Cloud que vous créez.
Sélectionnez le projet Google Cloud que vous avez créé :
```
gcloud config set project PROJECT_ID
```
Remplacez PROJECT_ID par le nom de votre projet Google Cloud .

Vérifiez que la facturation est activée pour votre projet Google Cloud .

Activez les API requises :

Rôles requis pour activer les API

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

Attribuez des rôles à votre compte utilisateur. Exécutez la commande suivante une fois pour chacun des rôles IAM suivants : roles/container.admin, roles/iam.serviceAccountAdmin, roles/storage.admin
```
gcloud projects add-iam-policy-binding PROJECT_ID --member="user:USER_IDENTIFIER" --role=ROLE
```
Remplacez les éléments suivants :
- PROJECT_ID : ID de votre projet
- USER_IDENTIFIER : identifiant de votre compte d'utilisateur. Par exemple, myemail@example.com.
- ROLE : rôle IAM que vous accordez à votre compte utilisateur.

Créez un compte Hugging Face si vous n'en possédez pas.
Assurez-vous de disposer d'un jeton Hugging Face.
Assurez-vous que votre projet dispose d'un quota suffisant pour les GPU B200 et H200. Pour en savoir plus, consultez Planifier le quota de GPU et Quota de GPU.

Préparer votre environnement

Dans ce tutoriel, vous utilisez Cloud Shell.

Accédez à la consoleGoogle Cloud .
En haut de la fenêtre de la console Google Cloud , cliquez sur le bouton Activer Cloud Shell.
Définissez les variables d'environnement suivantes :
```
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
```
Remplacez les valeurs suivantes :
- CLUSTER_NAME : nom de votre cluster GKE.
- CONTROL_PLANE_LOCATION : région Compute Engine du plan de contrôle du cluster GKE.
- NODE_LOCATION : emplacement de vos nœuds. Sélectionnez une zone où les GPU NVIDIA B200 ou H200 sont disponibles.
- GPU_TYPE : accélérateur que vous avez réservé dans la réservation de capacité Compute Engine. Il doit s'agir de l'une des options suivantes :
  - nvidia-b200 : NVIDIA B200 (180 Go)
  - nvidia-h200-141gb : NVIDIA H200 (141 Go)
- MACHINE_TYPE : type de machine à utiliser :
  - Pour les GPU NVIDIA B200 (180 Go), utilisez la version a4-highgpu-8g ou ultérieure.
  - Pour les GPU NVIDIA H200 (141 Go), utilisez a3-ultragpu-8g ou version ultérieure.
- GKE_VERSION : version de GKE à utiliser :
  - Pour les GPU NVIDIA B200 (180 Go), utilisez la version 1.32.2-gke.1422000 ou ultérieure.
  - Pour les GPU NVIDIA H200 (141 Go), utilisez 1.31.4-gke.1183000 ou version ultérieure.
- BUCKET_NAME : nom de base de votre bucket Cloud Storage.
- YOUR_HUGGING_FACE_TOKEN : votre jeton Hugging Face.
Créez les variables d'environnement suivantes pour le réseau :
```
export GVNIC_NETWORK_PREFIX="GVNIC-NAME"
export RDMA_NETWORK_PREFIX="RDMA-NAME"
```
Remplacez les valeurs suivantes :
- GVNIC-NAME : préfixe du nom du réseau gVNIC. Vous pouvez utiliser le préfixe de votre choix.
- RDMA-NAME : préfixe du réseau d'accès direct à la mémoire à distance (RDMA). Vous pouvez utiliser le préfixe de votre choix.

Configurer l'infrastructure

Dans cette section, vous allez créer des réseaux VPC et un cluster GKE.

Créer un réseau VPC

Créez un réseau VPC pour l'interface 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

Créez un réseau et des sous-réseaux VPC pour RDMA, y compris huit sous-réseaux pour huit GPU :

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

Créer le cluster GKE

Vous pouvez définir NeMo RL dans un cluster GKE Autopilot ou GKE Standard. Nous vous recommandons d'utiliser un cluster GKE Autopilot pour une expérience Kubernetes entièrement gérée. Pour choisir le mode de fonctionnement GKE le mieux adapté à vos charges de travail, consultez À propos des modes de fonctionnement GKE.

Autopilot

Créez un cluster Autopilot :

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

Obtenez les identifiants de votre cluster :

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

Installez le programme d'installation NCCL RDMA pour Autopilot :

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

Standard

Créez un cluster standard :

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

Obtenez les identifiants de votre cluster :

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

Créez le pool de nœuds 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

Installez le programme d'installation NCCL RDMA :

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

Configurer les mappages réseau

Enregistrez le manifeste suivant sous le nom 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

Appliquez le fichier manifeste :
```
kubectl apply -f network-mapping.yaml
```

Préparer l'espace de stockage

Dans cette section, vous allez créer des buckets Cloud Storage et une instance Lustre gérée, qui provisionne le stockage hautes performances requis pour votre charge de travail RL.

Créez un bucket Cloud Storage :

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

Créez un compte de service Kubernetes (KSA) et associez-le au 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"

Pour configurer Managed Lustre, procédez comme suit :
1. Créez une instance Managed Lustre en suivant la procédure décrite dans Créer une instance Managed Lustre. Assurez-vous que l'instance utilise le même réseau que votre cluster GKE.
2. Accédez à l'instance Managed Lustre en suivant la procédure décrite dans Accéder à une instance Managed Lustre existante.

Déployer un RayCluster

Dans cette section, vous allez cloner l'exemple de dépôt, préparer les fichiers manifestes et exécuter un script launcher.sh :

Clonez l'exemple de dépôt :

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

Accédez au répertoire de travail :
```
cd ai-ml/nemo-rl-on-gke/nemoRL
```

Inspectez le fichier manifeste 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

Remplacez NCCL_TUNER_CONFIG_PATH par l'une des valeurs suivantes, en fonction de l'accélérateur que vous utilisez dans ce tutoriel :

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

Dans ce fichier manifeste, le nœud principal gère le job et héberge le tableau de bord Ray. Les nœuds de calcul exécutent les tâches d'entraînement.

Installez le cluster Ray :
```
export REPLICA_COUNT=2
helm install ray-cluster . \
  --set values.additionalWorkerGroups.worker-grp-0.replicas=$REPLICA_COUNT
```
Pour ce tutoriel, vous utilisez deux nœuds de calcul. Si vous souhaitez modifier le nombre de nœuds de calcul, modifiez la valeur REPLICA_COUNT.
Pour déployer le cluster Ray, exécutez le script launcher.sh :
```
bash launcher.sh
```

Vérifiez que les nœuds de calcul et principaux sont en cours d'exécution :

kubectl get pods

Le résultat ressemble à ce qui suit :

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

Vérifiez que le cluster Ray est en cours d'exécution :

kubectl ray get clusters

Le résultat ressemble à ce qui suit :

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

Lancer le job GRPO

Une fois votre cluster Ray prêt, vous pouvez envoyer un job Ray à votre cluster Ray en cours d'exécution sur GKE. NeMo RL télécharge automatiquement le modèle lors de l'exécution du job d'entraînement RL.

Pour envoyer un job Ray, démarrez une session interactive pour l'exécuter.

Pour établir une connexion locale à votre cluster Ray, exécutez cette commande :
```
  kubectl ray session ray-cluster-kuberay
```
Cette commande lance le transfert de port entre votre machine locale et le nœud principal Ray de votre cluster GKE. Notez que votre terminal sera occupé pendant que cette session est active. Pour continuer, ouvrez une autre instance de terminal.

Modifiez le fichier gemma3-27b-gsm8k.sh comme suit :

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

Remplacez les valeurs suivantes dans le fichier gemma3-27b-gsm8k.sh :

YOUR_WANDB_API_KEY : votre clé API WandB.
YOUR_HF_TOKEN : votre jeton Hugging Face.

Dans ce fichier, vous pouvez voir la configuration permettant d'exécuter un job avec le modèle gemma3-27b-it sur l'ensemble de données GSM8k. Pour exécuter le pipeline d'entraînement GRPO, ce script définit les paramètres suivants :

num_prompts_per_step: 16 et num_generations_per_prompt: 64 : le modèle Gemma3-27b-it génère un grand nombre de réponses pour chaque requête. Dans cette configuration, le modèle produit 1 024 réponses au total (16 x 64 = 1 024).
policy.generation.colocated.enabled=False : ce paramètre désactive la fonctionnalité de génération colocalisée, ce qui signifie que le modèle ne génère pas de réponses dans le même nœud que le processus d'entraînement. Dans l'apprentissage par renforcement standard, les mêmes GPU gèrent l'entraînement et la génération. Dans cette configuration NeMo RL, vous dédiez des nœuds spécifiques (gérés avec le paramètre policy.generation.colocated.resources) uniquement à l'inférence vLLM, tandis que le reste du cluster se concentre sur les calculs d'entraînement intensifs. En séparant ces charges de travail, vous évitez les conflits de ressources entre les tampons d'entraînement gourmands en mémoire et les charges de travail d'inférence gourmandes en ressources de calcul.

Pour envoyer le job, exécutez la commande suivante :
```
bash gemma3-27b-it/gemma3-27b-gsm8k.sh
```
Lorsque le job est en cours d'exécution, la sortie affiche les résultats de l'entraînement, le timing et les métriques de performances.

Surveiller l'état du job GRPO

Une fois le job terminé, NeMo RL stocke les points de contrôle dans le chemin configuré.

Installez l'utilitaire apt-tree :
```
apt install tree
```

Pour surveiller l'état du job GRPO, vérifiez les journaux du nœud principal Ray :

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

Le résultat ressemble à ce qui suit :

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

Effectuer un nettoyage

Pour éviter que des frais ne vous soient facturés, supprimez les ressources :

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

Ajuster et mettre à l'échelle l'apprentissage par renforcement avec NVIDIA NeMo RL sur GKE Restez organisé à l'aide des collections Enregistrez et classez les contenus selon vos préférences.

Arrière-plan

Apprentissage par renforcement

Optimisation des stratégies relatives aux groupes (GRPO)

NVIDIA NeMo RL

Ensemble de données GSM8k

Objectifs

Avant de commencer

Préparer votre environnement

Configurer l'infrastructure

Créer un réseau VPC

Créer le cluster GKE

Autopilot

Standard

Configurer les mappages réseau

Préparer l'espace de stockage

Déployer un RayCluster

Lancer le job GRPO

Surveiller l'état du job GRPO

Effectuer un nettoyage

Étapes suivantes

Ajuster et mettre à l'échelle l'apprentissage par renforcement avec NVIDIA NeMo RL sur GKE