Google sử dụng công nghệ AI để dịch nội dung sang ngôn ngữ bạn ưu tiên. Bản dịch bằng AI có thể có lỗi.

Mở rộng quy mô hoạt động học tăng cường bằng GKE và Managed Lustre

1. Giới thiệu

Nếu muốn chạy trực tiếp các tập lệnh được đóng gói mà không cần hướng dẫn từng bước, bạn có thể tìm thấy các tập lệnh này trong kho lưu trữ GoogleCloudPlatform/devrel-demos.

Trong lớp học lập trình này, bạn sẽ tìm hiểu cách triển khai một quy trình huấn luyện hiệu suất cao cho Học tăng cường (RL) bằng cách sử dụng Google Kubernetes Engine (GKE) và Managed Lustre.

Các khối lượng công việc Học tăng cường, đặc biệt là những khối lượng công việc sử dụng các thuật toán như Tối ưu hoá chính sách tương đối theo nhóm (GRPO), tạo ra lượng dữ liệu khổng lồ trong quá trình "Tạo trải nghiệm" và yêu cầu kiểm tra thường xuyên. Bộ nhớ đối tượng tiêu chuẩn có thể gây ra tình trạng tắc nghẽn trong các đợt I/O này, khiến các trình tăng tốc đắt tiền không hoạt động.

Bạn sẽ sử dụng Managed Lustre, một hệ thống tệp song song, để loại bỏ những điểm tắc nghẽn này và đạt được thông lượng huấn luyện cao hơn.

Bạn sẽ thực hiện

Định cấu hình các biến môi trường cho một cụm Ray dựa trên GPU.
Cung cấp một cụm GPU tại chỗ trên GKE và một phiên bản Managed Lustre bằng Cluster Toolkit.
Triển khai một cụm KubeRay và gắn hệ thống tệp Lustre.
Gửi khối lượng công việc huấn luyện NeMo-RL.
Theo dõi thông lượng cao và độ trễ thấp của điểm kiểm tra bằng Cloud Monitoring.

Sơ đồ cấu trúc của GKE, KubeRay và Managed Lustre

Bạn cần có

Một trình duyệt web như Chrome.
Một dự án trên Google Cloud đã bật tính năng thanh toán.

Lớp học lập trình này dành cho người dùng kỹ thuật cao cấp, kỹ sư nền tảng và nhà nghiên cứu AI đã quen thuộc với GKE và các khái niệm về bộ nhớ.

Tổng thời lượng ước tính: 45 đến 60 phút cộng với 2 giờ thời gian đào tạo

2. Trước khi bắt đầu

Tạo một dự án trên Google Cloud

Trong Google Cloud Console, hãy chọn hoặc tạo một dự án trên đám mây của Google Cloud.
Đảm bảo bạn đã bật tính năng thanh toán cho dự án trên Cloud.

Khởi động Cloud Shell

Cloud Shell là một môi trường dòng lệnh chạy trong Google Cloud và được tải sẵn các công cụ cần thiết.

Nhấp vào Kích hoạt Cloud Shell ở đầu Cloud Console.
Sau khi kết nối với Cloud Shell, hãy xác minh thông tin xác thực của bạn:
```
gcloud auth list
```
Xác nhận rằng dự án của bạn đã được định cấu hình:
```
gcloud config get project
```
Nếu dự án của bạn không được thiết lập như mong đợi, hãy thiết lập dự án:
```
export PROJECT_ID=<YOUR_PROJECT_ID>
gcloud config set project $PROJECT_ID
```

Cài đặt Cluster Toolkit

Lớp học lập trình này sử dụng Cluster Toolkit (gcluster) để triển khai cụm GKE. Để biết hướng dẫn về cách thiết lập Cluster Toolkit, hãy xem Hướng dẫn thiết lập Cluster Toolkit.

Bật API

Chạy lệnh này trong Cloud Shell để bật tất cả các API bắt buộc:

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

3. Định cấu hình các biến môi trường

Để các lệnh trong lớp học lập trình này nhất quán, hãy thiết lập một số biến môi trường.

Tạo một tệp có tên là env.sh rồi điền cấu hình của bạn vào tệp đó. Bạn có thể sử dụng mẫu sau:

# 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

Thay thế <YOUR_PROJECT_ID> và <YOUR_HF_TOKEN> bằng các giá trị thực tế của bạn.

Nguồn tệp để tải các biến vào phiên hiện tại:

source env.sh

4. Triển khai Cụm GKE và Lustre được quản lý bằng Cluster Toolkit

Trong bước này, bạn sẽ sử dụng Cluster Toolkit (gcluster) để triển khai một cụm GKE có GPU Spot và tự động cung cấp bộ nhớ Lustre được quản lý bằng trình điều khiển Lustre CSI và PersistentVolumeClaim được định cấu hình sẵn (lustre-pvc).

Chuẩn bị Bản thiết kế

Trước khi triển khai, hãy xem lại bản thiết kế examples/gke-a4/gke-a4.yaml (để biết thông tin chi tiết, hãy xem phần Tạo một cụm A4):

Bật Managed Lustre: Huỷ nhận xét các phần mô-đun managed-lustre và lustre-pvc trong gke-a4.yaml.
Bật tiện ích bổ sung RayOperator: Đặt enable_ray_operator: true trong phần cài đặt mô-đun gke_cluster trong gke-a4.yaml.

Triển khai cơ sở hạ tầng

Đặt CIDR được uỷ quyền để truy cập Cloud Shell và triển khai bằng 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

Chờ quá trình triển khai hoàn tất. Cluster Toolkit tự động cung cấp mạng VPC, quan hệ ngang hàng Quyền truy cập vào dịch vụ riêng tư (PSA), hệ thống tệp Lustre được quản lý, trình điều khiển Lustre CSI, tiện ích bổ sung RayOperator và yêu cầu lưu trữ Kubernetes (lustre-pvc) trong một quy trình triển khai phối hợp.

5. Triển khai Cụm Ray trên GKE

Trong bước này, bạn sẽ triển khai một cụm KubeRay trên các nút GKE và gắn hệ thống tệp Lustre bằng PersistentVolumeClaim (lustre-pvc) do Cluster Toolkit tự động cung cấp.

Tạo cấu hình RayCluster

Tạo một tệp có tên là ray-cluster.yaml. Thao tác này chỉ định các nút đầu và nút worker của KubeRay, sử dụng loại trình tăng tốc nvidia-b200 và gắn ổ đĩa Lustre tại /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

Kết nối với Cụm

Đảm bảo phiên Cloud Shell của bạn được xác thực với cụm GKE:

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

Áp dụng cấu hình RayCluster

Áp dụng cấu hình cụm Ray:

kubectl apply -f ray-cluster.yaml

Xác minh trạng thái của Cụm

Theo dõi quá trình tạo nhóm:

kubectl get pods -w

Đợi cho đến khi các nhóm đầu và nhóm worker là Running.

6. Gửi khối lượng công việc học tăng cường

Trong bước này, bạn sẽ gửi quy trình huấn luyện GRPO NeMo-RL đến cụm Ray.

Kết nối với Trang tổng quan Ray

Để gửi các công việc và xem chỉ số, bạn cần kết nối với Trang tổng quan của Ray. Vì trang tổng quan nằm trong GKE, hãy sử dụng tính năng chuyển tiếp cổng để truy cập vào trang tổng quan này từ Cloud Shell:

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

Tạo tập lệnh thực thi

Tạo một tệp có tên là run_nemo_rl.sh. Tập lệnh này sẽ được thực thi trên các worker của cụm Ray. Chúng ta sẽ dùng cat << EOF để điền vào các biến môi trường mà bạn đã đặt trước đó.

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

Tạo tệp bỏ qua Ray

Tạo tệp .rayignore để ngăn Ray tải lên các thư mục lớn hoặc không cần thiết:

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

Tạo cấu hình môi trường thời gian chạy

Tạo một tệp JSON để truyền các biến môi trường đến công việc 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

Gửi yêu cầu

Sử dụng Ray CLI để gửi công việc đến điểm cuối trang tổng quan. Nếu không tìm thấy lệnh ray trong Cloud Shell, bạn có thể cài đặt lệnh này bằng pip install ray:

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

Bạn sẽ thấy nhật ký truyền trực tuyến trong thiết bị đầu cuối Cloud Shell. Công việc này sẽ tải mô hình, khởi động các worker Ray và bắt đầu vòng lặp huấn luyện GRPO.

7. Theo dõi hiệu suất huấn luyện

Trong bước này, bạn sẽ quan sát hiệu suất của hệ thống tệp Lustre trong quá trình huấn luyện và kiểm tra.

Kiểm tra nhật ký huấn luyện

Khi quá trình huấn luyện diễn ra, bạn sẽ thấy nhật ký cho biết các điểm kiểm tra đang được lưu vào /lustre/nemo_rl_qwen_72b_ds_cp/nemo-rl-grpo-test1. Lưu ý rằng quá trình tạo điểm kiểm tra diễn ra không đồng bộ và không chặn các worker Ray trong thời gian quá dài.

Để xem tốc độ tạo điểm kiểm tra, hãy tìm các dòng nhật ký cho biết các điểm kiểm tra đã lưu.

Xem các chỉ số Lustre trong Cloud Console

Cách xem các chỉ số cho phiên bản Lustre:

Trong Google Cloud Console, hãy tìm Managed Service for Lustre.
Nhấp vào tên phiên bản của bạn (ví dụ: ${CLUSTER_NAME}-lustre hoặc rl-demo-gpu-lustre).
Nhấp vào thẻ Giám sát.

Tại đây, bạn có thể quan sát:

Thông lượng (Bytes/giây): Xem các mức tăng đột biến trong quá trình tạo điểm kiểm tra.
Dung lượng: Theo dõi dung lượng mà các điểm kiểm tra đang sử dụng.

Biểu đồ hiệu suất của Lustre Lustre có khả năng ghi ở tốc độ rất cao, ghi các điểm kiểm tra trong thời gian tối thiểu

8. Dọn dẹp tài nguyên

Chạy lệnh sau trong Cloud Shell để huỷ tất cả cơ sở hạ tầng được cung cấp (cụm GKE, nhóm nút GPU, phiên bản Lustre được quản lý và mạng VPC) trong một bước:

gcluster destroy "${CLUSTER_NAME}"

Lệnh này chạy đồng bộ ở nền trước, phá huỷ mọi cơ sở hạ tầng do quá trình triển khai quản lý và xuất nhật ký tiến trình vào thiết bị đầu cuối. Vui lòng đợi lệnh hoàn tất hoàn toàn trước khi đóng phiên Cloud Shell.

9. Xin chúc mừng

Bạn đã hoàn tất thành công lớp học lập trình Mở rộng quy mô học tăng cường bằng GKE và Lustre được quản lý!

Kiến thức bạn học được

Cách sử dụng Cluster Toolkit để cung cấp một cụm GPU GKE bằng các phiên bản Spot và bộ nhớ Lustre được quản lý.
Cách triển khai một cụm KubeRay và gắn bộ nhớ Lustre.
Cách gửi khối lượng công việc huấn luyện GRPO NeMo-RL.
Cách quan sát hiệu suất bộ nhớ trong quá trình huấn luyện.

Các bước tiếp theo

Khám phá thêm các tính năng của NVIDIA NeMo-RL.
Tìm hiểu thêm về Siêu máy tính AI của Google Cloud.
Xem tài liệu về Dịch vụ được quản lý cho Lustre.