ปรับขนาดการเรียนรู้แบบเสริมกำลังด้วย GKE และ Lustre ที่มีการจัดการ

1. บทนำ

หากต้องการเรียกใช้สคริปต์ที่แพ็กไว้โดยตรงโดยไม่ต้องทำตามบทแนะนำทีละขั้นตอน คุณสามารถค้นหาสคริปต์ได้ในที่เก็บ GoogleCloudPlatform/devrel-demos

ใน Codelab นี้ คุณจะได้เรียนรู้วิธีทําให้ไปป์ไลน์การฝึกที่มีประสิทธิภาพสูงสําหรับการเรียนรู้แบบเสริมกําลัง (RL) ใช้งานได้โดยใช้ Google Kubernetes Engine (GKE) และ Managed Lustre

ภาระงานการเรียนรู้แบบเสริมกำลัง โดยเฉพาะอย่างยิ่งภาระงานที่ใช้อัลกอริทึม เช่น การเพิ่มประสิทธิภาพนโยบายแบบสัมพัทธ์ของกลุ่ม (GRPO) จะสร้างข้อมูลจำนวนมหาศาลในระหว่าง "การสร้างประสบการณ์" และต้องมีการตรวจสอบเป็นระยะๆ ที่เก็บข้อมูลออบเจ็กต์มาตรฐานอาจทำให้เกิดคอขวดระหว่างการเพิ่ม I/O เหล่านี้ ซึ่งทำให้ตัวเร่งความเร็วที่มีราคาแพงไม่ได้ใช้งาน

คุณจะใช้ Managed Lustre ซึ่งเป็นระบบไฟล์แบบขนานเพื่อขจัดจุดคอขวดเหล่านี้และเพิ่มอัตราการส่งข้อมูลการฝึก

สิ่งที่คุณต้องทำ

  • กำหนดค่าตัวแปรสภาพแวดล้อมสำหรับคลัสเตอร์ Ray ที่ใช้ GPU
  • จัดสรรคลัสเตอร์ GPU แบบสปอตใน GKE และอินสแตนซ์ Lustre ที่มีการจัดการโดยใช้ชุดเครื่องมือคลัสเตอร์
  • ทําให้คลัสเตอร์ KubeRay ใช้งานได้และติดตั้งระบบไฟล์ Lustre
  • ส่งภาระงานการฝึก NeMo-RL
  • สังเกตการณ์การส่งข้อความปริมาณมากและเวลาในการตอบสนองของจุดตรวจสอบต่ำโดยใช้ Cloud Monitoring

แผนภาพสถาปัตยกรรมของ GKE, KubeRay และ Lustre ที่มีการจัดการ

สิ่งที่คุณต้องมี

  • เว็บเบราว์เซอร์ เช่น Chrome
  • โปรเจ็กต์ Google Cloud ที่เปิดใช้การเรียกเก็บเงิน

Codelab นี้เหมาะสำหรับผู้ใช้ด้านเทคนิคขั้นสูง วิศวกรแพลตฟอร์ม และนักวิจัย AI ที่คุ้นเคยกับแนวคิดของ GKE และพื้นที่เก็บข้อมูล

ระยะเวลารวมโดยประมาณ: 45-60 นาที บวกเวลาฝึกอบรม 2 ชั่วโมง

2. ก่อนเริ่มต้น

สร้างโปรเจ็กต์ Google Cloud

  1. ในคอนโซล Google Cloud ให้เลือกหรือสร้างโปรเจ็กต์ Google Cloud
  2. ตรวจสอบว่าได้เปิดใช้การเรียกเก็บเงินสำหรับโปรเจ็กต์ที่อยู่ในระบบคลาวด์แล้ว

เริ่มต้น Cloud Shell

Cloud Shell คือสภาพแวดล้อมบรรทัดคำสั่งที่ทำงานใน Google Cloud ซึ่งโหลดเครื่องมือที่จำเป็นไว้ล่วงหน้า

  1. คลิกเปิดใช้งาน 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

ติดตั้งชุดเครื่องมือของคลัสเตอร์

Codelab นี้ใช้ 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. กำหนดค่าตัวแปรสภาพแวดล้อม

ตั้งค่าตัวแปรสภาพแวดล้อม 2-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 และ Lustre ที่มีการจัดการโดยใช้ Cluster Toolkit

ในขั้นตอนนี้ คุณจะใช้ Cluster Toolkit (gcluster) เพื่อติดตั้งใช้งานคลัสเตอร์ GKE ด้วย GPU แบบสปอต และจัดสรรพื้นที่เก็บข้อมูล 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 การเข้าถึงบริการส่วนตัว (PSA), ระบบไฟล์ Lustre ที่มีการจัดการ, ไดรเวอร์ Lustre CSI, ส่วนเสริม RayOperator และการอ้างสิทธิ์พื้นที่เก็บข้อมูล Kubernetes (lustre-pvc) โดยอัตโนมัติในการติดตั้งใช้งานที่ประสานงานเดียว

5. ติดตั้งใช้งานคลัสเตอร์ Ray บน GKE

ในขั้นตอนนี้ คุณจะทำให้คลัสเตอร์ KubeRay ใช้งานได้ในโหนด GKE และติดตั้งระบบไฟล์ Lustre โดยใช้ PersistentVolumeClaim (lustre-pvc) ที่ Cluster Toolkit จัดสรรให้โดยอัตโนมัติ

สร้างการกำหนดค่า RayCluster

สร้างไฟล์ชื่อ ray-cluster.yaml ซึ่งจะระบุโหนดหัวหน้าและโหนด Worker ของ 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

รอจนกว่าพ็อดส่วนหัวและพ็อดของ Worker จะเป็น Running

6. ส่งภาระงานการเรียนรู้แบบเสริมกำลัง

ในขั้นตอนนี้ คุณจะส่งงานการฝึก GRPO ของ NeMo-RL ไปยังคลัสเตอร์ 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 สคริปต์นี้จะดำเนินการใน Worker ของคลัสเตอร์ 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

สร้างไฟล์ .rayignore เพื่อป้องกันไม่ให้ Ray อัปโหลดไดเรกทอรีขนาดใหญ่หรือไดเรกทอรีที่ไม่จำเป็น

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 งานจะโหลดโมเดล เริ่มต้นใช้งาน Worker ของ Ray และเริ่มลูปการฝึก GRPO

7. ตรวจสอบประสิทธิภาพการฝึก

ในขั้นตอนนี้ คุณจะสังเกตประสิทธิภาพของระบบไฟล์ Lustre ระหว่างการฝึกและการตรวจสอบ

ตรวจสอบบันทึกการฝึก

เมื่อการฝึกดำเนินไป คุณจะเห็นบันทึกที่ระบุว่าระบบกำลังบันทึกจุดตรวจไปยัง /lustre/nemo_rl_qwen_72b_ds_cp/nemo-rl-grpo-test1 โปรดทราบว่าการสร้างจุดตรวจสอบจะเกิดขึ้นแบบไม่พร้อมกันและจะไม่บล็อก Worker ของ Ray เป็นเวลานาน

หากต้องการดูความเร็วของการตรวจสอบ ให้มองหาบรรทัดบันทึกที่ระบุจุดตรวจสอบที่บันทึกไว้

ดูเมตริก Lustre ใน Cloud Console

วิธีดูเมตริกสำหรับอินสแตนซ์ Lustre

  1. ใน คอนโซล Google Cloud ให้ค้นหา Managed Service for Lustre
  2. คลิกชื่ออินสแตนซ์ (เช่น ${CLUSTER_NAME}-lustre หรือ rl-demo-gpu-lustre)
  3. คลิกแท็บการตรวจสอบ

โดยคุณจะดูข้อมูลต่อไปนี้ได้

  • อัตราการส่งข้อมูล (ไบต์/วินาที): ดูจุดสูงสุดระหว่างการตรวจสอบ
  • ความจุ: ตรวจสอบว่าจุดตรวจสอบใช้พื้นที่เท่าใด

แผนภูมิประสิทธิภาพ LustreLustre สามารถเขียนด้วยความเร็วสูงมากและเขียนจุดตรวจสอบได้ในเวลาอันสั้น

8. ล้างข้อมูลทรัพยากร

เรียกใช้คำสั่งต่อไปนี้ใน Cloud Shell เพื่อทำลายโครงสร้างพื้นฐานที่จัดสรรทั้งหมด (คลัสเตอร์ GKE, Node Pool ของ GPU, อินสแตนซ์ Lustre ที่มีการจัดการ และเครือข่าย VPC) ในขั้นตอนเดียว

gcluster destroy "${CLUSTER_NAME}"

คำสั่งนี้จะทำงานแบบซิงโครนัสในเบื้องหน้า โดยจะทำลายโครงสร้างพื้นฐานทั้งหมดที่การติดตั้งใช้งานจัดการ และส่งบันทึกความคืบหน้าไปยังเทอร์มินัล โปรดรอให้คำสั่งเสร็จสมบูรณ์ก่อนปิดเซสชัน Cloud Shell

9. ขอแสดงความยินดี

คุณทำ Codelab Scale Reinforcement Learning with GKE and Managed Lustre เสร็จสมบูรณ์แล้ว

สิ่งที่คุณได้เรียนรู้

  • วิธีใช้ชุดเครื่องมือคลัสเตอร์เพื่อจัดสรรคลัสเตอร์ GPU ของ GKE ด้วยอินสแตนซ์สปอตและพื้นที่เก็บข้อมูล Lustre ที่มีการจัดการ
  • วิธีติดตั้งใช้งานคลัสเตอร์ KubeRay และเมานต์พื้นที่เก็บข้อมูล Lustre
  • วิธีส่งภาระงานการฝึก GRPO ของ NeMo-RL
  • วิธีสังเกตประสิทธิภาพของพื้นที่เก็บข้อมูลระหว่างการฝึก

ขั้นตอนถัดไป