Scale Reinforcement Learning with GKE and Managed Lustre

1. מבוא

אם אתם מעדיפים להריץ את הסקריפטים הארוזים ישירות בלי להשתמש במדריך המפורט, תוכלו למצוא אותם במאגר GoogleCloudPlatform/devrel-demos.

ב-Codelab הזה תלמדו איך לפרוס צינור עיבוד נתונים לאימון ביצועים גבוהים ללמידת חיזוק (RL) באמצעות Google Kubernetes Engine ‏ (GKE) ו-Managed Lustre.

עומסי עבודה של למידת חיזוק, במיוחד כאלה שמשתמשים באלגוריתמים כמו Group Relative Policy Optimization ‏ (GRPO), יוצרים כמויות עצומות של נתונים במהלך 'יצירת חוויה' ודורשים יצירת נקודות ביקורת לעיתים קרובות. אחסון אובייקטים רגיל עלול לגרום לצווארי בקבוק במהלך פרצי קלט/פלט כאלה, ולהשאיר מאיצים יקרים ללא פעולה.

כדי למנוע את צווארי הבקבוק האלה ולהשיג תפוקת אימון גבוהה יותר, תשתמשו ב-Managed Lustre, מערכת קבצים מקבילה.

הפעולות שתבצעו:

  • הגדרת משתני סביבה עבור אשכול Ray מבוסס-GPU.
  • הקצאת אשכול GPU מסוג Spot ב-GKE ומופע Managed Lustre באמצעות Cluster Toolkit.
  • פורסים אשכול KubeRay ומטמיעים את מערכת הקבצים של Lustre.
  • שליחה של עומס עבודה לאימון של NeMo-RL.
  • התפוקה גבוהה והחביון של נקודות הבדיקה נמוך באמצעות Cloud Monitoring.

דיאגרמת ארכיטקטורה של GKE,‏ KubeRay ו-Managed Lustre

הדרישות

  • דפדפן אינטרנט כמו Chrome.
  • פרויקט ב-Google Cloud שהחיוב בו מופעל.

ה-Codelab הזה מיועד למשתמשים טכניים מתקדמים, למהנדסי פלטפורמות ולחוקרי AI שמכירים את GKE ואת מושגי האחסון.

משך זמן כולל משוער: 45 עד 60 דקות בתוספת שעתיים של זמן הדרכה

‫2. לפני שמתחילים

יצירת פרויקט ב-Google Cloud

  1. במסוף Google Cloud, בוחרים או יוצרים פרויקט בענן של Google.
  2. הקפידו לוודא שהחיוב מופעל בפרויקט שלכם ב-Cloud.

הפעלת Cloud Shell

‫Cloud Shell היא סביבת שורת פקודה שפועלת ב-Google Cloud וכוללת מראש את הכלים הנדרשים.

  1. לוחצים על Activate Cloud Shell בחלק העליון של מסוף Google Cloud.
  2. אחרי שמתחברים ל-Cloud Shell, מאמתים את האימות:
    gcloud auth list
  3. מוודאים שהפרויקט מוגדר:
    gcloud config get project
  4. אם הפרויקט לא מוגדר כמו שציפיתם, מגדירים אותו:
    export PROJECT_ID=<YOUR_PROJECT_ID>
gcloud config set project $PROJECT_ID

התקנת Cluster Toolkit

בשיעור הזה נשתמש ב-Cluster Toolkit ‏ (gcluster) כדי לפרוס את אשכול GKE. הוראות להגדרת Cluster Toolkit מופיעות במדריך ההגדרה של Cluster Toolkit.

הפעלת ממשקי ה-API

מריצים את הפקודה הזו ב-Cloud Shell כדי להפעיל את כל ממשקי ה-API הנדרשים:

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

3. הגדרת משתני סביבה

כדי לשמור על עקביות הפקודות ב-codelab הזה, צריך להגדיר כמה משתני סביבה.

יוצרים קובץ בשם env.sh וממלאים אותו בהגדרות. אפשר להשתמש בתבנית הבאה:

# 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

מחליפים את <YOUR_PROJECT_ID> ואת <YOUR_HF_TOKEN> בערכים האמיתיים.

מריצים את הקובץ כדי לטעון את המשתנים לסשן הנוכחי:

source env.sh

4. פריסת אשכול GKE ו-Managed Lustre באמצעות Cluster Toolkit

בשלב הזה, משתמשים ב-Cluster Toolkit ‏ (gcluster) כדי לפרוס אשכול GKE עם Spot GPU ולספק באופן אוטומטי אחסון Managed Lustre באמצעות מנהל התקנים של Lustre CSI ו-PersistentVolumeClaim (lustre-pvc) שהוגדר מראש.

הכנת ה-Blueprint

לפני הפריסה, כדאי לעיין בתוכנית האב examples/gke-a4/gke-a4.yaml (פרטים נוספים זמינים במאמר יצירת אשכול A4):

  1. הפעלת Lustre מנוהל: מבטלים את ההערה בקטעי המודולים managed-lustre ו-lustre-pvc ב-gke-a4.yaml.
  2. הפעלת התוסף RayOperator: מגדירים את enable_ray_operator: true בהגדרות המודול gke_cluster ב-gke-a4.yaml.

פריסת תשתית

מגדירים את ה-CIDR המורשה לגישה ל-Cloud Shell ומבצעים פריסה באמצעות 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

מחכים שהפריסה תסתיים. ‫Cluster Toolkit מקצה באופן אוטומטי את רשת ה-VPC, את הקישור בין רשתות שכנות (peering) של Private Service Access ‏(PSA), את Managed Lustre filesystem, את Lustre CSI driver, את התוסף RayOperator ואת בקשת האחסון של Kubernetes ‏(lustre-pvc) בפריסה מתואמת אחת.

5. פריסת אשכול Ray ב-GKE

בשלב הזה תפרסו אשכול KubeRay בצמתי GKE ותתקינו את מערכת הקבצים Lustre באמצעות PersistentVolumeClaim ‏ (lustre-pvc) שהוקצה אוטומטית על ידי Cluster Toolkit.

יצירת הגדרה של RayCluster

יוצרים קובץ בשם ray-cluster.yaml. הפקודה הזו מציינת את הצמתים הראשיים והצמתים של העובדים ב-KubeRay, באמצעות סוג המאיץ nvidia-b200 והרכבת נפח האחסון של 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

חיבור לאשכול

מוודאים שהסשן של Cloud Shell מאומת באשכול GKE:

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

החלת הגדרות RayCluster

החלת ההגדרה של אשכול Ray:

kubectl apply -f ray-cluster.yaml

אימות סטטוס האשכול

עוקבים אחרי יצירת ה-Pods:

kubectl get pods -w

מחכים עד שהפודים של ה-head וה-worker יהיו במצב Running.

6. שליחת עומס עבודה של למידה עם חיזוקים

בשלב הזה, תגישו את משימת האימון של NeMo-RL GRPO לאשכול Ray.

התחברות ללוח הבקרה של Ray

כדי לשלוח משימות ולראות מדדים, צריך להתחבר ללוח הבקרה של Ray. מכיוון שהלוח נמצא ב-GKE, צריך להשתמש בהעברת פורטים כדי לגשת אליו מ-Cloud Shell:

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

יצירת סקריפט ההפעלה

יוצרים קובץ בשם run_nemo_rl.sh. הסקריפט הזה יופעל על העובדים של אשכול Ray. אנחנו משתמשים ב-cat << EOF כדי למלא את משתני הסביבה שהגדרתם קודם.

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

כדי למנוע מ-Ray להעלות ספריות גדולות או לא נחוצות, יוצרים קובץ .rayignore:

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

יצירת הגדרה של סביבת זמן ריצה

יוצרים קובץ JSON כדי להעביר משתני סביבה למשימת 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

שליחת המשרה

משתמשים ב-Ray CLI כדי לשלוח את העבודה לנקודת הקצה של מרכז הבקרה. אם הפקודה ray לא נמצאת ב-Cloud Shell, אפשר להתקין אותה באמצעות pip install ray:

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

יומנים יזרמו לטרמינל של Cloud Shell. העבודה תטען את המודל, תפעיל את עובדי Ray ותתחיל את לולאת האימון של GRPO.

7. מעקב אחרי ביצועי האימון

בשלב הזה, תצפו בביצועים של מערכת הקבצים Lustre במהלך האימון והסימון של נקודות ביקורת.

בדיקת יומני האימון

במהלך האימון, יומנים יציינו שנקודות ביקורת נשמרות ב-/lustre/nemo_rl_qwen_72b_ds_cp/nemo-rl-grpo-test1. שימו לב ששמירת נקודות הביקורת מתבצעת באופן אסינכרוני ולא חוסמת את עובדי Ray למשך זמן רב מדי.

כדי לראות את המהירות של יצירת נקודות ביקורת, מחפשים שורות ביומן שמציינות נקודות ביקורת שנשמרו.

צפייה במדדים של Lustre ב-Cloud Console

כדי לראות את המדדים של מכונת Lustre:

  1. ב-מסוף Google Cloud, מחפשים את Managed Service for Lustre.
  2. לוחצים על שם המכונה (לדוגמה, ${CLUSTER_NAME}-lustre או rl-demo-gpu-lustre).
  3. לוחצים על הכרטיסייה מעקב.

כאן אפשר לראות:

  • קצב העברת נתונים (בייט לשנייה): אפשר לראות את העליות החדות במהלך יצירת נקודות ביקורת.
  • קיבולת: מעקב אחרי נפח האחסון שמשמש לנקודות בקרה.

תרשים ביצועים של LustreLustre יכול לכתוב במהירות גבוהה מאוד, ולכתוב נקודות ביקורת בזמן מינימלי

8. מחיקת משאבי הבדיקה

כדי להרוס את כל התשתית שהוקצתה (אשכול GKE, מאגרי צמתים של GPU, מכונת Lustre מנוהלת ורשת VPC) בפעולה אחת, מריצים את הפקודה הבאה ב-Cloud Shell:

gcluster destroy "${CLUSTER_NAME}"

הפקודה הזו פועלת באופן סינכרוני בחזית, ומפרקת את כל התשתית שמנוהלת על ידי הפריסה. יומני ההתקדמות מוצגים במסוף. צריך לחכות שהפקודה תסתיים לגמרי לפני שסוגרים את סשן Cloud Shell.

9. מזל טוב

סיימתם בהצלחה את ה-codelab Scale Reinforcement Learning with GKE and Managed Lustre.

מה למדתם

  • איך משתמשים ב-Cluster Toolkit כדי להקצות אשכול GPU ב-GKE עם מכונות Spot ואחסון Managed Lustre.
  • איך פורסים אשכול KubeRay וטוענים אחסון Lustre.
  • איך שולחים עומס עבודה לאימון של NeMo-RL GRPO.
  • איך בודקים את ביצועי האחסון במהלך האימון.

השלבים הבאים