Skalowanie uczenia ze wzmocnieniem za pomocą GKE i Managed Lustre

1. Wprowadzenie

Jeśli wolisz uruchamiać spakowane skrypty bezpośrednio bez samouczka krok po kroku, znajdziesz je w repozytorium GoogleCloudPlatform/devrel-demos.

Z tego ćwiczenia dowiesz się, jak wdrożyć wydajny potok trenowania do uczenia ze wzmocnieniem (RL) przy użyciu Google Kubernetes Engine (GKE) i Managed Lustre.

Obciążenia związane z uczeniem ze wzmocnieniem, zwłaszcza te, które korzystają z algorytmów takich jak Group Relative Policy Optimization (GRPO), generują ogromne ilości danych podczas „generowania doświadczeń” i wymagają częstego tworzenia punktów kontrolnych. Standardowa pamięć obiektowa może powodować wąskie gardła podczas tych skoków operacji wejścia/wyjścia, pozostawiając drogie akceleratory bezczynne.

Aby wyeliminować te wąskie gardła i uzyskać większą przepustowość trenowania, użyjesz Managed Lustre, czyli równoległego systemu plików.

Jakie zadania wykonasz

• Skonfiguruj zmienne środowiskowe dla klastra Ray opartego na GPU.
• Zarezerwuj klaster GPU typu Spot w GKE i instancję Managed Lustre za pomocą Cluster Toolkit.
• Wdróż klaster KubeRay i zamontuj system plików Lustre.
• Prześlij zadanie trenowania NeMo-RL.
• Obserwuj wysoką przepustowość i krótkie czasy oczekiwania na punkty kontrolne za pomocą Cloud Monitoring.

Czego potrzebujesz

• przeglądarka, np. Chrome;
• projekt Google Cloud z włączonymi płatnościami;

W tym ćwiczeniu wykorzystywane są zasoby wysokiej klasy, w tym GPU typu spot i usługa Managed Lustre. Chociaż korzystamy z zasobów typu spot, aby obniżyć koszty, mogą one być drogie, jeśli pozostaną włączone. Aby uniknąć nieoczekiwanych opłat, zawsze wykonuj krok „Wyczyść” na końcu modułu.

To ćwiczenie jest przeznaczone dla zaawansowanych użytkowników technicznych, inżynierów platform i badaczy AI, którzy znają GKE i koncepcje związane z pamięcią masową.

Szacowany łączny czas trwania: 45–60 minut plus 2 godziny szkolenia

2. Zanim zaczniesz

Tworzenie projektu Google Cloud

1. W konsoli Google Cloud wybierz lub utwórz projekt w chmurze Google Cloud.
2. Sprawdź, czy w projekcie Cloud włączone są płatności.

Uruchamianie Cloud Shell

Cloud Shell to środowisko wiersza poleceń działające w Google Cloud, które zawiera niezbędne narzędzia.

1. U góry konsoli Google Cloud kliknij Aktywuj Cloud Shell.
2. Po połączeniu z Cloud Shell sprawdź uwierzytelnianie:

   gcloud auth list
   

3. Sprawdź, czy projekt jest skonfigurowany:

   gcloud config get project
   

4. Jeśli projekt nie jest ustawiony zgodnie z oczekiwaniami, ustaw go:

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

Instalowanie narzędzia Cluster Toolkit

W tym ćwiczeniu do wdrożenia klastra GKE używamy narzędzia Cluster Toolkit (gcluster). Instrukcje konfigurowania narzędzia Cluster Toolkit znajdziesz w przewodniku konfiguracji narzędzia Cluster Toolkit.

Włącz interfejsy API

Aby włączyć wszystkie wymagane interfejsy API, uruchom to polecenie w Cloud Shell:

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

3. Konfigurowanie zmiennych środowiskowych

Aby polecenia w tym samouczku były spójne, skonfiguruj kilka zmiennych środowiskowych.

Utwórz plik o nazwie env.sh i wypełnij go konfiguracją. Możesz użyć tego szablonu:

# 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 to token dostępu użytkownika Hugging Face, który jest wymagany do pobierania modeli. Możesz go uzyskać, postępując zgodnie z instrukcjami na stronie Hugging Face Security Tokens (Tokeny zabezpieczające Hugging Face).

Zastąp <YOUR_PROJECT_ID> i <YOUR_HF_TOKEN> rzeczywistymi wartościami.

Element WANDB_API_KEY jest opcjonalny. Jest on używany w Weights & Biases, narzędziu do monitorowania i wizualizacji eksperymentów związanych z uczeniem maszynowym. Jeśli nie masz konta lub nie chcesz go używać, możesz pozostawić to pole puste lub bez zmian.

Uruchom plik, aby wczytać zmienne do bieżącej sesji:

source env.sh

4. Wdrażanie klastra GKE i zarządzanego systemu Lustre za pomocą narzędzia Cluster Toolkit

W tym kroku użyjesz narzędzi Cluster Toolkit (gcluster), aby wdrożyć klaster GKE z GPU Spot i automatycznie udostępnić pamięć masową Managed Lustre za pomocą sterownika CSI Lustre i wstępnie skonfigurowanego elementu PersistentVolumeClaim (lustre-pvc).

Przygotowywanie planu

Przed wdrożeniem zapoznaj się z examples/gke-a4/gke-a4.yamlplanem (szczegóły znajdziesz w artykule Tworzenie klastra A4):

1. Włącz zarządzany system Lustre: usuń komentarze z sekcji modułów managed-lustre i lustre-pvc w pliku gke-a4.yaml.
2. Włącz dodatek RayOperator: ustaw enable_ray_operator: true w ustawieniach modułu gke_cluster w gke-a4.yaml.

Wdrażanie infrastruktury

Ustaw autoryzowany zakres CIDR dla dostępu do Cloud Shell i wdróż go za pomocą polecenia 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

Ustawienie spot=true w --vars dynamicznie udostępnia pulę węzłów GPU przy użyciu mocy obliczeniowej Spot bez konieczności ręcznego edytowania reservation_affinity. Pamiętaj, że zasoby typu spot mogą zostać wywłaszczone. W przypadku środowisk produkcyjnych pomiń spot=true, aby używać pojemności na żądanie, lub skonfiguruj rezerwację.

Poczekaj na zakończenie wdrażania. Cluster Toolkit automatycznie udostępnia sieć VPC, połączenie równorzędne Private Service Access (PSA), Managed Lustre filesystem, sterownik CSI Lustre, dodatek RayOperator i roszczenie dotyczące pamięci Kubernetes (lustre-pvc) w ramach jednego skoordynowanego wdrożenia.

5. Wdrażanie klastra Ray w GKE

W tym kroku wdrożysz klaster KubeRay na węzłach GKE i zamontujesz system plików Lustre za pomocą elementu PersistentVolumeClaim (lustre-pvc) automatycznie udostępnionego przez Cluster Toolkit.

Tworzenie konfiguracji RayCluster

Utwórz plik o nazwie ray-cluster.yaml. Określa węzły główne i robocze KubeRay, używając nvidia-b200 typu akceleratora i montując wolumin Lustre w /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

Łączenie z klastrem

Sprawdź, czy sesja Cloud Shell jest uwierzytelniona w klastrze GKE:

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

Stosowanie konfiguracji RayCluster

Zastosuj konfigurację klastra Ray:

kubectl apply -f ray-cluster.yaml

Sprawdzanie stanu klastra

Monitoruj tworzenie podów:

kubectl get pods -w

Poczekaj, aż pody węzła głównego i procesu roboczego będą w stanie Running.

6. Przesyłanie zadania uczenia się przez wzmacnianie

W tym kroku prześlesz zadanie trenowania NeMo-RL GRPO do klastra Ray.

Łączenie z panelem Ray

Aby przesyłać zadania i wyświetlać dane, musisz połączyć się z panelem Ray. Ponieważ panel znajduje się w GKE, użyj przekierowania portów, aby uzyskać do niego dostęp z Cloud Shell:

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

Tworzenie skryptu wykonania

Utwórz plik o nazwie run_nemo_rl.sh. Ten skrypt zostanie wykonany na instancjach roboczych klastra Ray. Używamy cat << EOF, aby wypełnić zmienne środowiskowe ustawione wcześniej.

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

Tworzenie pliku ignorowania Ray

Utwórz plik .rayignore, aby zapobiec przesyłaniu przez Ray dużych lub niepotrzebnych katalogów:

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

Tworzenie konfiguracji środowiska wykonawczego

Utwórz plik JSON, aby przekazać zmienne środowiskowe do zadania 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

Przesyłanie zadania

Użyj interfejsu wiersza poleceń Ray, aby przesłać zadanie do punktu końcowego panelu. Jeśli polecenie ray nie zostanie znalezione w Cloud Shell, możesz je zainstalować za pomocą polecenia pip install ray:

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

W terminalu Cloud Shell zobaczysz przesyłane strumieniowo logi. Zadanie wczyta model, zainicjuje instancje robocze Ray i rozpocznie pętlę trenowania GRPO.

7. Monitorowanie wydajności trenowania

Na tym etapie obserwujesz wydajność systemu plików Lustre podczas trenowania i tworzenia punktów kontrolnych.

Sprawdzanie dzienników trenowania

W trakcie trenowania zobaczysz logi wskazujące, że punkty kontrolne są zapisywane w /lustre/nemo_rl_qwen_72b_ds_cp/nemo-rl-grpo-test1. Zwróć uwagę, że tworzenie punktów kontrolnych odbywa się asynchronicznie i nie blokuje procesów roboczych Ray przez zbyt długi czas.

Aby sprawdzić szybkość tworzenia punktów kontrolnych, poszukaj wierszy dziennika wskazujących zapisane punkty kontrolne.

Wyświetlanie wskaźników Lustre w konsoli Cloud

Aby wyświetlić dane dotyczące instancji Lustre:

1. W konsoli Google Cloud wyszukaj Managed Service for Lustre.
2. Kliknij nazwę instancji (np. ${CLUSTER_NAME}-lustre lub rl-demo-gpu-lustre).
3. Kliknij kartę Monitorowanie.

Możesz tu zobaczyć:

• Przepustowość (bajty/s): obserwuj skoki podczas tworzenia punktów kontrolnych.
• Pojemność: monitoruj ilość miejsca zajmowanego przez punkty kontrolne.

Lustre może zapisywać dane z bardzo dużą szybkością, a punkty kontrolne są zapisywane w minimalnym czasie.

8. Usuwanie zasobów

Ten krok jest niezbędny, aby zapobiec dalszym obciążeniom. Pozostawione włączone klastry GPU i Managed Lustre mogą generować znaczne koszty.

Aby zniszczyć całą udostępnioną infrastrukturę (klaster GKE, pule węzłów GPU, instancję Managed Lustre i sieć VPC) w jednym kroku, uruchom w Cloud Shell to polecenie:

gcluster destroy "${CLUSTER_NAME}"

To polecenie jest wykonywane synchronicznie na pierwszym planie, usuwając całą infrastrukturę zarządzaną przez wdrożenie i wyświetlając w terminalu dzienniki postępu. Zanim zamkniesz sesję Cloud Shell, poczekaj na całkowite zakończenie działania polecenia.

9. Gratulacje

Udało Ci się ukończyć ćwiczenie w Codelabs Skalowanie uczenia ze wzmocnieniem za pomocą GKE i Managed Lustre.

Czego się nauczysz

• Jak używać narzędzia Cluster Toolkit do udostępniania klastra GKE z procesorami GPU z instancjami typu spot i zarządzanym magazynem Lustre.
• Jak wdrożyć klaster KubeRay i zamontować pamięć Lustre.
• Jak przesłać zadanie trenowania GRPO NeMo-RL.
• Jak obserwować wydajność pamięci masowej podczas trenowania.

Dalsze kroki

• Poznaj więcej funkcji NVIDIA NeMo-RL.
• Dowiedz się więcej o superkomputerze AI od Google Cloud.
• Zapoznaj się z dokumentacją usługi Managed Service for Lustre.