GKE और मैनेज किए गए Lustre की मदद से, रीइन्फ़ोर्समेंट लर्निंग को स्केल करना

1. परिचय

अगर आपको चरण-दर-चरण दिए गए ट्यूटोरियल के बिना, सीधे तौर पर पैकेज की गई स्क्रिप्ट को चलाना है, तो उन्हें GoogleCloudPlatform/devrel-demos रिपॉज़िटरी में देखा जा सकता है.

इस कोडलैब में, आपको Google Kubernetes Engine (GKE) और Managed Lustre का इस्तेमाल करके, रीइन्फ़ोर्समेंट लर्निंग (आरएल) के लिए ज़्यादा परफ़ॉर्मेंस वाली ट्रेनिंग पाइपलाइन को डिप्लॉय करने का तरीका बताया जाएगा.

रीइन्फ़ोर्समेंट लर्निंग वर्कलोड, खास तौर पर ग्रुप रिलेटिव पॉलिसी ऑप्टिमाइज़ेशन (जीआरपीओ) जैसे एल्गोरिदम का इस्तेमाल करने वाले वर्कलोड, "एक्सपीरियंस जनरेशन" के दौरान बहुत ज़्यादा डेटा जनरेट करते हैं. साथ ही, इन्हें बार-बार चेकपॉइंट करने की ज़रूरत होती है. इनपुट/आउटपुट के दौरान स्टैंडर्ड ऑब्जेक्ट स्टोरेज की वजह से, परफ़ॉर्मेंस में रुकावटें आ सकती हैं. इससे महंगे ऐक्सलरेटर का इस्तेमाल नहीं हो पाता.

इन समस्याओं को दूर करने और ट्रेनिंग थ्रूपुट को बेहतर बनाने के लिए, Managed Lustre का इस्तेमाल किया जाएगा. यह एक पैरलल फ़ाइल सिस्टम है.

आपको क्या करना होगा

• जीपीयू पर आधारित Ray क्लस्टर के लिए, एनवायरमेंट वैरिएबल कॉन्फ़िगर करें.
• क्लस्टर टूलकिट का इस्तेमाल करके, GKE पर स्पॉट जीपीयू क्लस्टर और मैनेज किया गया Lustre इंस्टेंस उपलब्ध कराएं.
• KubeRay क्लस्टर डिप्लॉय करें और Lustre फ़ाइल सिस्टम को माउंट करें.
• NeMo-RL ट्रेनिंग वर्कलोड सबमिट करें.
• Cloud Monitoring का इस्तेमाल करके, ज़्यादा थ्रूपुट और कम चेकपॉइंट लेटेंसी देखें.

आपको किन चीज़ों की ज़रूरत होगी

• कोई वेब ब्राउज़र, जैसे कि Chrome.
• बिलिंग की सुविधा वाला Google क्लाउड प्रोजेक्ट.

इस कोडलैब में, स्पॉट जीपीयू और मैनेज किए गए Lustre जैसे बेहतरीन संसाधनों का इस्तेमाल किया गया है. हम लागत कम करने के लिए स्पॉट कैपेसिटी का इस्तेमाल करते हैं. हालांकि, अगर इन संसाधनों को चालू रखा जाता है, तो इनका इस्तेमाल करना महंगा पड़ सकता है. अनचाहे शुल्क से बचने के लिए, लैब के आखिर में हमेशा क्लीन अप करने का तरीका अपनाएं.

यह कोडलैब, तकनीकी मामलों के जानकारों, प्लैटफ़ॉर्म इंजीनियरों, और एआई रिसर्चर के लिए है. इन्हें GKE और स्टोरेज के कॉन्सेप्ट के बारे में जानकारी है.

कुल अनुमानित अवधि: 45 से 60 मिनट और दो घंटे की ट्रेनिंग

2. शुरू करने से पहले

Google Cloud प्रोजेक्ट बनाना

1. Google Cloud Console में जाकर, Google Cloud प्रोजेक्ट चुनें या बनाएं.
2. पक्का करें कि आपके क्लाउड प्रोजेक्ट के लिए बिलिंग की सुविधा चालू हो.

Cloud Shell शुरू करना

Cloud Shell, Google Cloud में चलने वाला कमांड-लाइन एनवायरमेंट है. इसमें ज़रूरी टूल पहले से लोड होते हैं.

1. Google Cloud कंसोल में सबसे ऊपर मौजूद, Cloud Shell चालू करें पर क्लिक करें.
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 इंस्टॉल करना

इस कोडलैब में, GKE क्लस्टर को डिप्लॉय करने के लिए Cluster Toolkit (gcluster) का इस्तेमाल किया गया है. क्लस्टर टूलकिट को सेट अप करने के निर्देशों के लिए, क्लस्टर टूलकिट को सेट अप करने की गाइड देखें.

एपीआई चालू करें

सभी ज़रूरी एपीआई चालू करने के लिए, Cloud Shell में यह कमांड चलाएं:

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

3. एनवायरमेंट वैरिएबल कॉन्फ़िगर करना

इस कोडलैब में दिए गए निर्देशों को एक जैसा रखने के लिए, कुछ एनवायरमेंट वैरिएबल सेट अप करें.

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

HF_TOKEN आपका Hugging Face User Access Token है. मॉडल डाउनलोड करने के लिए, इसकी ज़रूरत होती है. इसे पाने के लिए, Hugging Face Security Tokens पर दिए गए निर्देशों का पालन करें.

<YOUR_PROJECT_ID> और <YOUR_HF_TOKEN> को अपनी असल वैल्यू से बदलें.

WANDB_API_KEY वैकल्पिक है. इसका इस्तेमाल Weights & Biases के लिए किया जाता है. यह मशीन लर्निंग के एक्सपेरिमेंट की निगरानी करने और उन्हें विज़ुअलाइज़ करने का टूल है. अगर आपके पास खाता नहीं है या आपको इसका इस्तेमाल नहीं करना है, तो इसे खाली छोड़ दें या जैसा है वैसा ही रहने दें.

अपने मौजूदा सेशन में वैरिएबल लोड करने के लिए, फ़ाइल को सोर्स करें:

source env.sh

4. Cluster Toolkit का इस्तेमाल करके, GKE क्लस्टर और मैनेज किए गए Lustre को डिप्लॉय करना

इस चरण में, स्पॉट जीपीयू के साथ GKE क्लस्टर को डिप्लॉय करने के लिए, Cluster Toolkit (gcluster) का इस्तेमाल किया जाता है. साथ ही, Lustre CSI ड्राइवर और पहले से कॉन्फ़िगर किए गए PersistentVolumeClaim (lustre-pvc) के साथ, मैनेज किए गए Lustre स्टोरेज को अपने-आप प्रोविज़न किया जाता है.

ब्लूप्रिंट तैयार करना

डिप्लॉय करने से पहले, examples/gke-a4/gke-a4.yaml ब्लूप्रिंट की समीक्षा करें. ज़्यादा जानकारी के लिए, A4 क्लस्टर बनाना लेख पढ़ें:

1. Managed Lustre चालू करें: gke-a4.yaml में managed-lustre और lustre-pvc मॉड्यूल सेक्शन से टिप्पणी हटाएं.
2. RayOperator ऐड-ऑन चालू करें: gke-a4.yaml में gke_cluster मॉड्यूल की सेटिंग में जाकर, enable_ray_operator: true को सेट करें.

इंफ़्रास्ट्रक्चर डिप्लॉय करना

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

--vars में spot=true को सेट करने पर, reservation_affinity में मैन्युअल तरीके से बदलाव करने की ज़रूरत नहीं होती. साथ ही, स्पॉट क्षमता का इस्तेमाल करके, GPU नोड पूल को डाइनैमिक तरीके से उपलब्ध कराया जाता है. ध्यान दें कि स्पॉट इंस्टेंस की क्षमता को रोका जा सकता है. प्रॉडक्शन एनवायरमेंट के लिए, मांग के हिसाब से क्षमता का इस्तेमाल करने या बुकिंग को कॉन्फ़िगर करने के लिए, spot=true को शामिल न करें.

डिप्लॉयमेंट पूरा होने तक इंतज़ार करें. क्लस्टर टूलकिट, वीपीसी नेटवर्क, Private Service Access (PSA) पियरिंग, मैनेज किया गया Lustre फ़ाइल सिस्टम, Lustre CSI ड्राइवर, RayOperator ऐड-ऑन, और Kubernetes स्टोरेज क्लेम (lustre-pvc) को एक साथ डिप्लॉयमेंट में अपने-आप उपलब्ध कराता है.

5. GKE पर Ray Cluster डिप्लॉय करना

इस चरण में, आपको अपने GKE नोड पर KubeRay क्लस्टर डिप्लॉय करना होगा. साथ ही, Cluster Toolkit की ओर से अपने-आप उपलब्ध कराए गए PersistentVolumeClaim (lustre-pvc) का इस्तेमाल करके, Lustre फ़ाइल सिस्टम को माउंट करना होगा.

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

क्लस्टर के स्टेटस की पुष्टि करना

पॉड बनने की प्रोसेस को मॉनिटर करें:

kubectl get pods -w

जब तक हेड और वर्कर पॉड 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 Ignore फ़ाइल बनाना

Ray को बड़ी या गैर-ज़रूरी डायरेक्ट्री अपलोड करने से रोकने के लिए, .rayignore फ़ाइल बनाएं:

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

रनटाइम एनवायरमेंट कॉन्फ़िगरेशन बनाना

Ray जॉब को एनवायरमेंट वैरिएबल पास करने के लिए, एक JSON फ़ाइल बनाएं:

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 का इस्तेमाल करें. अगर Cloud Shell में ray कमांड नहीं मिलती है, तो इसे 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 वर्कर को ज़्यादा समय तक ब्लॉक नहीं करती.

चेकपॉइंटिंग की स्पीड देखने के लिए, उन लॉग लाइनों को देखें जो सेव किए गए चेकपॉइंट के बारे में बताती हैं.

Cloud Console में Lustre मेट्रिक देखना

अपने Lustre इंस्टेंस की मेट्रिक देखने के लिए:

1. Google Cloud Console में, Managed Service for Lustre खोजें.
2. अपने इंस्टेंस के नाम पर क्लिक करें. उदाहरण के लिए, ${CLUSTER_NAME}-lustre या rl-demo-gpu-lustre.
3. निगरानी टैब पर क्लिक करें.

यहां आपको यह जानकारी दिखेगी:

• थ्रूपुट (बाइट/सेकंड): चेकपॉइंटिंग के दौरान स्पाइक देखें.
• क्षमता: मॉनिटर करें कि चेकपॉइंट कितना स्टोरेज इस्तेमाल कर रहे हैं.

Lustre बहुत तेज़ी से लिख सकता है और कम समय में चेकपॉइंट लिख सकता है

8. संसाधन मिटाना

यह चरण, मौजूदा शुल्क को रोकने के लिए ज़रूरी है. अगर इन्हें बंद नहीं किया जाता है, तो GPU क्लस्टर और मैनेज किए गए Lustre के लिए ज़्यादा शुल्क देना पड़ सकता है.

एक ही चरण में सभी प्रोविज़न किए गए इन्फ़्रास्ट्रक्चर (GKE (जीकेई) क्लस्टर, जीपीयू नोड पूल, मैनेज किया गया Lustre इंस्टेंस, और वीपीसी नेटवर्क) को बंद करने के लिए, Cloud Shell में यह निर्देश चलाएं:

gcluster destroy "${CLUSTER_NAME}"

यह कमांड, फ़ोरग्राउंड में सिंक्रोनस तरीके से चलती है. इससे, डिप्लॉयमेंट की मदद से मैनेज किए गए सभी इन्फ़्रास्ट्रक्चर को बंद कर दिया जाता है. साथ ही, प्रोग्रेस लॉग को आपके टर्मिनल पर आउटपुट कर दिया जाता है. Cloud Shell सेशन बंद करने से पहले, कमांड के पूरी तरह से पूरा होने का इंतज़ार करें.

9. बधाई हो

आपने GKE और मैनेज किए गए Lustre की मदद से, रीइन्फ़ोर्समेंट लर्निंग को स्केल करना कोडलैब पूरा कर लिया है!

आपने क्या सीखा

• स्पॉट इंस्टेंस और मैनेज किए गए Lustre स्टोरेज के साथ GKE GPU क्लस्टर बनाने के लिए, Cluster Toolkit का इस्तेमाल कैसे करें.
• KubeRay क्लस्टर को डिप्लॉय करने और Lustre स्टोरेज को माउंट करने का तरीका.
• NeMo-RL GRPO ट्रेनिंग वर्कलोड सबमिट करने का तरीका.
• ट्रेनिंग के दौरान स्टोरेज की परफ़ॉर्मेंस को कैसे मॉनिटर करें.

अगले चरण

• NVIDIA NeMo-RL की अन्य सुविधाओं के बारे में जानें.
• Google Cloud AI Hypercomputer के बारे में ज़्यादा जानें.
• Managed Service for Lustre के दस्तावेज़ पढ़ें.