Verstärkungslernen mit GKE und Managed Lustre skalieren

1. Einführung

Wenn Sie die verpackten Skripts lieber direkt ohne die Schritt-für-Schritt-Anleitung ausführen möchten, finden Sie sie im Repository GoogleCloudPlatform/devrel-demos.

In diesem Codelab erfahren Sie, wie Sie eine leistungsstarke Trainingspipeline für Reinforcement Learning (RL) mit Google Kubernetes Engine (GKE) und Managed Lustre bereitstellen.

Workloads für das verstärkende Lernen, insbesondere solche, bei denen Algorithmen wie Group Relative Policy Optimization (GRPO) verwendet werden, generieren während der „Experience Generation“ (Erlebnisgenerierung) riesige Datenmengen und erfordern häufige Prüfpunkte. Standardobjektspeicher kann bei diesen I/O-Spitzen zu Engpässen führen, sodass teure Beschleuniger nicht genutzt werden.

Sie verwenden Managed Lustre, ein paralleles Dateisystem, um diese Engpässe zu beseitigen und einen höheren Trainingsdurchsatz zu erzielen.

Aufgaben

• Konfigurieren Sie Umgebungsvariablen für einen GPU-basierten Ray-Cluster.
• Stellen Sie einen Spot-GPU-Cluster in GKE und eine Managed Lustre-Instanz mit dem Cluster Toolkit bereit.
• Stellen Sie einen KubeRay-Cluster bereit und stellen Sie das Lustre-Dateisystem bereit.
• Senden Sie eine NeMo-RL-Trainingsarbeitslast.
• Hohen Durchsatz und niedrige Checkpoint-Latenz beobachten mit Cloud Monitoring.

Voraussetzungen

• Ein Webbrowser wie Chrome.
• Google Cloud-Projekt mit aktivierter Abrechnungsfunktion.

In diesem Codelab werden High-End-Ressourcen wie Spot-GPUs und Managed Lustre verwendet. Wir verwenden Spot-Kapazität, um Kosten zu senken. Diese Ressourcen können jedoch teuer sein, wenn sie weiter ausgeführt werden. Führen Sie am Ende des Labs immer den Bereinigungsschritt aus, um unerwartete Kosten zu vermeiden.

Dieses Codelab richtet sich an fortgeschrittene technische Nutzer, Plattformentwickler und KI-Forscher, die mit GKE und Speicherkonzepten vertraut sind.

Geschätzte Gesamtdauer: 45 bis 60 Minuten plus 2 Stunden Trainingszeit

2. Hinweis

Google Cloud-Projekt erstellen

1. Wählen Sie in der Google Cloud Console ein Google Cloud-Projekt aus oder erstellen Sie eines.
2. Die Abrechnung für das Cloud-Projekt muss aktiviert sein.

Cloud Shell starten

Cloud Shell ist eine Befehlszeilenumgebung, die in Google Cloud ausgeführt wird und mit den erforderlichen Tools vorinstalliert ist.

1. Klicken Sie oben in der Google Cloud Console auf Cloud Shell aktivieren.
2. Prüfen Sie nach der Verbindung mit Cloud Shell Ihre Authentifizierung:

   gcloud auth list
   

3. Prüfen Sie, ob Ihr Projekt konfiguriert ist:

   gcloud config get project
   

4. Wenn Ihr Projekt nicht wie erwartet festgelegt ist, legen Sie es fest:

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

Cluster Toolkit installieren

In diesem Codelab wird das Cluster Toolkit (gcluster) verwendet, um den GKE-Cluster bereitzustellen. Eine Anleitung zum Einrichten von Cluster Toolkit finden Sie im Einrichtungsleitfaden für Cluster Toolkit.

APIs aktivieren

Führen Sie diesen Befehl in Cloud Shell aus, um alle erforderlichen APIs zu aktivieren:

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

3. Umgebungsvariablen konfigurieren

Damit die Befehle in diesem Codelab einheitlich sind, richten Sie einige Umgebungsvariablen ein.

Erstellen Sie eine Datei mit dem Namen env.sh und fügen Sie Ihre Konfiguration ein. Sie können die folgende Vorlage verwenden:

# 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 ist Ihr Hugging Face-Nutzerzugriffstoken, das zum Herunterladen von Modellen erforderlich ist. Eine Anleitung dazu finden Sie unter Hugging Face-Sicherheitstokens.

Ersetzen Sie <YOUR_PROJECT_ID> und <YOUR_HF_TOKEN> durch Ihre tatsächlichen Werte.

WANDB_API_KEY ist optional. Es wird für Weights & Biases verwendet, ein Tool zum Überwachen und Visualisieren von Machine-Learning-Tests. Wenn Sie kein Konto haben oder es nicht verwenden möchten, können Sie das Feld leer lassen oder so belassen.

Rufen Sie die Datei auf, um die Variablen in Ihre aktuelle Sitzung zu laden:

source env.sh

4. GKE-Cluster und Managed Lustre mit Cluster Toolkit bereitstellen

In diesem Schritt verwenden Sie das Cluster Toolkit (gcluster), um einen GKE-Cluster mit Spot-GPUs bereitzustellen und automatisch Managed Lustre-Speicher mit dem Lustre-CSI-Treiber und einem vorkonfigurierten PersistentVolumeClaim (lustre-pvc) bereitzustellen.

Blueprint vorbereiten

Sehen Sie sich vor der Bereitstellung den examples/gke-a4/gke-a4.yaml-Blueprint an (weitere Informationen finden Sie unter A4-Cluster erstellen):

1. Managed Lustre aktivieren: Entfernen Sie die Auskommentierung der Modulabschnitte managed-lustre und lustre-pvc in gke-a4.yaml.
2. RayOperator-Add-on aktivieren: Setzen Sie enable_ray_operator: true in den gke_cluster-Moduleinstellungen in gke-a4.yaml.

Infrastruktur bereitstellen

Legen Sie den autorisierten CIDR für den Cloud Shell-Zugriff fest und stellen Sie die Anwendung mit gcluster deploy bereit:

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

Wenn Sie spot=true in --vars festlegen, wird der GPU-Knotenpool dynamisch mit Spot-Kapazität bereitgestellt, ohne dass manuelle Änderungen an reservation_affinity erforderlich sind. Beachten Sie, dass Spot-Kapazität unterbrochen werden kann. Lassen Sie spot=true für Produktionsumgebungen weg, um On-Demand-Kapazität zu verwenden oder eine Reservierung zu konfigurieren.

Warten Sie, bis die Bereitstellung abgeschlossen ist. Mit dem Cluster Toolkit werden das VPC-Netzwerk, das PSA-Peering (Private Service Access), das Managed Lustre-Dateisystem, der Lustre-CSI-Treiber, das RayOperator-Add-on und der Kubernetes-Speicheranspruch (lustre-pvc) automatisch in einer koordinierten Bereitstellung bereitgestellt.

5. Ray-Cluster in GKE bereitstellen

In diesem Schritt stellen Sie einen KubeRay-Cluster auf Ihren GKE-Knoten bereit und hängen das Lustre-Dateisystem mit dem PersistentVolumeClaim (lustre-pvc) ein, der automatisch vom Cluster Toolkit bereitgestellt wird.

RayCluster-Konfiguration erstellen

Erstellen Sie eine Datei mit dem Namen ray-cluster.yaml. Damit werden die KubeRay-Head- und Worker-Knoten mit dem Beschleunigertyp nvidia-b200 angegeben und das Lustre-Volume wird unter /lustre bereitgestellt.

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

Verbindung zum Cluster herstellen

Prüfen Sie, ob Ihre Cloud Shell-Sitzung für Ihren GKE-Cluster authentifiziert ist:

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

RayCluster-Konfiguration anwenden

Wenden Sie die Ray-Clusterkonfiguration an:

kubectl apply -f ray-cluster.yaml

Clusterstatus prüfen

Überwachen Sie die Erstellung der Pods:

kubectl get pods -w

Warten Sie, bis die Head- und Worker-Pods den Status Running haben.

6. Reinforcement Learning-Arbeitslast senden

In diesem Schritt senden Sie den NeMo-RL GRPO-Trainingsjob an Ihren Ray-Cluster.

Verbindung zum Ray-Dashboard herstellen

Wenn Sie Jobs senden und Messwerte ansehen möchten, müssen Sie eine Verbindung zum Ray-Dashboard herstellen. Da sich das Dashboard in GKE befindet, verwenden Sie die Portweiterleitung, um darauf über Cloud Shell zuzugreifen:

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

Ausführungsskript erstellen

Erstellen Sie eine Datei mit dem Namen run_nemo_rl.sh. Dieses Skript wird auf den Ray-Cluster-Workern ausgeführt. Wir verwenden cat << EOF, um die Umgebungsvariablen einzufügen, die Sie zuvor festgelegt haben.

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

Ray-Ignore-Datei erstellen

Erstellen Sie eine .rayignore-Datei, um zu verhindern, dass Ray große oder unnötige Verzeichnisse hochlädt:

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

Laufzeitumgebung konfigurieren

Erstellen Sie eine JSON-Datei, um Umgebungsvariablen an den Ray-Job zu übergeben:

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

Job senden

Verwenden Sie die Ray CLI, um den Job an den Dashboard-Endpunkt zu senden. Wenn der Befehl ray in Cloud Shell nicht gefunden wird, können Sie ihn mit pip install ray installieren:

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

In Ihrem Cloud Shell-Terminal werden Logs gestreamt. Der Job lädt das Modell, initialisiert die Ray-Worker und beginnt mit der GRPO-Trainingsschleife.

7. Trainingsleistung überwachen

In diesem Schritt beobachten Sie die Leistung des Lustre-Dateisystems während des Trainings und der Checkpointing-Vorgänge.

Trainingslogs prüfen

Im Laufe des Trainings werden Logs angezeigt, die darauf hinweisen, dass Prüfpunkte in /lustre/nemo_rl_qwen_72b_ds_cp/nemo-rl-grpo-test1 gespeichert werden. Das Erstellen von Checkpoints erfolgt asynchron und blockiert die Ray-Worker nicht sehr lange.

Die Geschwindigkeit der Checkpointing-Vorgänge können Sie anhand von Logzeilen sehen, die auf gespeicherte Checkpoints hinweisen.

Lustre-Messwerte in der Cloud Console ansehen

So rufen Sie Messwerte für Ihre Lustre-Instanz auf:

1. Suchen Sie in der Google Cloud Console nach Managed Service for Lustre.
2. Klicken Sie auf den Namen Ihrer Instanz, z. B. ${CLUSTER_NAME}-lustre oder rl-demo-gpu-lustre.
3. Klicken Sie auf den Tab Monitoring.

Hier sehen Sie:

• Durchsatz (Byte/s): Hier sehen Sie die Spitzenwerte während der Erstellung von Prüfpunkten.
• Kapazität: Hier können Sie sehen, wie viel Speicherplatz von Checkpoints belegt wird.

Lustre kann mit sehr hoher Geschwindigkeit schreiben und Prüfpunkte in kürzester Zeit schreiben.

8. Ressourcen bereinigen

Dieser Schritt ist wichtig, um laufende Gebühren zu vermeiden. Wenn GPU-Cluster und Managed Lustre ausgeführt werden, können erhebliche Kosten entstehen.

Führen Sie den folgenden Befehl in Cloud Shell aus, um die gesamte bereitgestellte Infrastruktur (GKE-Cluster, GPU-Knotenpools, Managed Lustre-Instanz und VPC-Netzwerk) in einem einzigen Schritt zu löschen:

gcluster destroy "${CLUSTER_NAME}"

Dieser Befehl wird synchron im Vordergrund ausgeführt. Dabei wird die gesamte von der Bereitstellung verwaltete Infrastruktur abgebaut und es werden Fortschrittsprotokolle in Ihrem Terminal ausgegeben. Warten Sie, bis der Befehl vollständig ausgeführt wurde, bevor Sie Ihre Cloud Shell-Sitzung schließen.

9. Glückwunsch

Sie haben das Codelab Scale Reinforcement Learning with GKE and Managed Lustre (Verstärkungslernen mit GKE und Managed Lustre skalieren) erfolgreich abgeschlossen.

Das haben Sie gelernt

• So stellen Sie mit Cluster Toolkit einen GKE-GPU-Cluster mit Spot-Instanzen und Managed Lustre-Speicher bereit.
• So stellen Sie einen KubeRay-Cluster bereit und hängen Lustre-Speicher ein.
• So senden Sie eine NeMo-RL GRPO-Trainingsarbeitslast.
• Speicherleistung während des Trainings beobachten

Nächste Schritte

• Weitere Funktionen von NVIDIA NeMo-RL
• Weitere Informationen zum Google Cloud AI Hypercomputer
• Dokumentation zu Managed Service for Lustre