本教學課程說明如何在 Google Kubernetes Engine (GKE) 上,為強化學習 (RL) 協調分散式訓練環境。您可以使用 Ray 和 NVIDIA NeMo RL 架構設定分散式訓練環境,微調模型。
本教學課程著重於在 GKE 上,使用 Ray 和 NeMo RL 進行群組相對政策最佳化 (GRPO) 訓練管道。GRPO 是一種強化學習演算法,旨在提升模型的推理能力。這項演算法可節省記憶體,並透過淘汰 Critic 或價值模型,改用相對群組計算,簡化 RL 程序。
建議您先完成在 GKE 上使用 Verl 微調及擴充強化學習教學課程,再執行本教學課程。下列教學課程使用的叢集設定和配置,與使用 Verl 進行 RL 微調和調整大小的教學課程相同。
背景
以下各節簡要說明本教學課程使用的概念。
增強學習 (RL)
RL 會透過經驗、探索和意見回饋訓練模型,而非靜態模仿。雖然預先訓練會教導模型該說什麼,但人類回饋增強學習 (RLHF) 會教導模型如何提供實用、安全且合乎邏輯的內容。RL 是基礎模型與微調模型之間的橋樑,適用於特定用途。
詳情請參閱「什麼是強化學習?」一文。
群組相對政策最佳化 (GRPO)
GRPO 是 DeepSeek 普及的演算法,可移除 Critic 模型,為 LLM 對齊提供記憶體效率高的替代方案,取代近端策略最佳化 (PPO)。GRPO 不使用 Critic 網路,而是為相同提示產生一組回應,並將該組回應的平均獎勵做為基準。
詳情請參閱 GRPO。
NVIDIA NeMo RL
NeMo RL 是 NVIDIA 的開放原始碼後訓練程式庫,專為可擴充的 RL 設計。NeMo RL 是更廣泛的 NeMo 架構生態系統的一部分,可讓您在單一 GPU 上進行小規模實驗,也能在數千個 GPU 上部署多個節點。
詳情請參閱 NVIDIA NeMo RL。
GSM8k 資料集
在本教學課程中,您會使用 GSM8k 資料集,其中包含 8,500 個優質、語言多元的小學數學文字題。
模型會使用 GSM8k 和 GRPO,針對同一問題生成一組 n 個不同的回覆。GRPO 會比較這些回覆與群組平均值。相較於群組中的其他回覆,模型若能持續提供正確且合乎邏輯的回覆,就能獲得更多獎勵。隨著時間推移,模型會瞭解清楚說明步驟是獲得最多獎勵最可靠的方法,進而有效減少低成效答案的獎勵。
詳情請參閱「GSM8k」。
目標
本教學課程說明如何透過下列步驟,在 GKE 上使用 NeMo RL 設定 RL:
- 準備環境。
- 設定搭載 B200 或 H200 GPU 的 GKE 叢集。
- 設定 KubeRay 來管理分散式 Ray 叢集。
- 使用 Managed Lustre 儲存空間,享有高效能。
- 執行使用 NeMo RL 的 GRPO 訓練工作。
事前準備
- 登入 Google Cloud 帳戶。如果您是 Google Cloud新手,歡迎 建立帳戶,親自評估產品在實際工作環境中的成效。新客戶還能獲得價值 $300 美元的免費抵免額,可用於執行、測試及部署工作負載。
-
安裝 Google Cloud CLI。
-
若您採用的是外部識別資訊提供者 (IdP),請先使用聯合身分登入 gcloud CLI。
-
執行下列指令,初始化 gcloud CLI:
gcloud init -
選取或建立專案所需的角色
- 選取專案:選取專案時,不需要具備特定 IAM 角色,只要您在專案中獲派角色,即可選取該專案。
-
建立專案:如要建立專案,您需要「專案建立者」角色 (
roles/resourcemanager.projectCreator),其中包含resourcemanager.projects.create權限。瞭解如何授予角色。
-
建立 Google Cloud 專案:
gcloud projects create PROJECT_ID
將
PROJECT_ID替換為您要建立的 Google Cloud 專案名稱。 -
選取您建立的 Google Cloud 專案:
gcloud config set project PROJECT_ID
將
PROJECT_ID替換為 Google Cloud 專案名稱。
啟用必要的 API:
啟用 API 時所需的角色
如要啟用 API,您需要具備服務使用情形管理員 IAM 角色 (
roles/serviceusage.serviceUsageAdmin),其中包含serviceusage.services.enable權限。瞭解如何授予角色。gcloud services enable container.googleapis.com
storage.googleapis.com compute.googleapis.com -
安裝 Google Cloud CLI。
-
若您採用的是外部識別資訊提供者 (IdP),請先使用聯合身分登入 gcloud CLI。
-
執行下列指令,初始化 gcloud CLI:
gcloud init -
選取或建立專案所需的角色
- 選取專案:選取專案時,不需要具備特定 IAM 角色,只要您在專案中獲派角色,即可選取該專案。
-
建立專案:如要建立專案,您需要「專案建立者」角色 (
roles/resourcemanager.projectCreator),其中包含resourcemanager.projects.create權限。瞭解如何授予角色。
-
建立 Google Cloud 專案:
gcloud projects create PROJECT_ID
將
PROJECT_ID替換為您要建立的 Google Cloud 專案名稱。 -
選取您建立的 Google Cloud 專案:
gcloud config set project PROJECT_ID
將
PROJECT_ID替換為 Google Cloud 專案名稱。
啟用必要的 API:
啟用 API 時所需的角色
如要啟用 API,您需要具備服務使用情形管理員 IAM 角色 (
roles/serviceusage.serviceUsageAdmin),其中包含serviceusage.services.enable權限。瞭解如何授予角色。gcloud services enable container.googleapis.com
storage.googleapis.com compute.googleapis.com -
將角色授予使用者帳戶。針對下列每個 IAM 角色,執行一次下列指令:
roles/container.admin, roles/iam.serviceAccountAdmin, roles/storage.admingcloud projects add-iam-policy-binding PROJECT_ID --member="user:USER_IDENTIFIER" --role=ROLE
更改下列內容:
PROJECT_ID:專案 ID。USER_IDENTIFIER:使用者帳戶的 ID。 例如:myemail@example.com。ROLE:授予使用者帳戶的 IAM 角色。
- 如果您尚未建立 Hugging Face 帳戶,請先完成此步驟。
- 請確認您擁有 Hugging Face 權杖。
- 請確認專案有足夠的 B200 和 H200 GPU 配額。詳情請參閱「規劃 GPU 配額」和「GPU 配額」。
準備環境
在本教學課程中,您將使用 Cloud Shell。
在 Google Cloud 主控台視窗頂端,按一下「啟用 Cloud Shell」按鈕。
請設定下列環境變數:
export PROJECT_ID=$(gcloud config get project) export PROJECT_NUMBER=$(gcloud projects describe ${PROJECT_ID} --format="value(projectNumber)") export CONTROL_PLANE_REGION=CONTROL_PLANE_REGION export NODE_ZONE=NODE_ZONE export CLUSTER_NAME=CLUSTER_NAME export GPU_TYPE=GPU_TYPE export MACHINE_TYPE=MACHINE_TYPE export KSA_NAME=generic-ksa export NAMESPACE=default export RESERVATION=RESERVATION export LUSTRE_NAME=LUSTRE_NAME export HF_TOKEN=YOUR_HF_TOKEN export WANDB_API_KEY=YOUR_WANDB_API_KEY替換下列值:
CLUSTER_NAME:GKE 叢集的名稱。CONTROL_PLANE_REGION:GKE 叢集控制層的 Compute Engine 區域。NODE_ZONE:節點的區域。選取提供 NVIDIA B200 或 H200 GPU 的區域。GPU_TYPE:您在 Compute Engine 容量預留中預留的加速器。必須是下列其中一個值:nvidia-b200:NVIDIA B200 (180 GB)nvidia-h200-141gb:NVIDIA H200 (141 GB)
MACHINE_TYPE:要使用的機器類型:- 如要使用 NVIDIA B200 (180 GB) GPU,請使用
a4-highgpu-8g以上版本。 - 如要使用 NVIDIA H200 (141 GB) GPU,請使用
a3-ultragpu-8g以上版本。
- 如要使用 NVIDIA B200 (180 GB) GPU,請使用
RESERVATION:GPU 預留項目名稱。LUSTRE_NAME:Lustre 執行個體的名稱。YOUR_HF_TOKEN:您的 Hugging Face 權杖。YOUR_WANDB_API_KEY:您的 Wandb API 金鑰。
為網路建立下列環境變數:
export NETWORK="NETWORK-NAME" export GVNIC_NETWORK_PREFIX="GVNIC-NAME" export RDMA_NETWORK_PREFIX="RDMA-NAME"替換下列值:
NETWORK-NAME:GKE 的網路名稱。GVNIC-NAME:gVNIC 網路名稱的前置字元。您可以選擇任何前置字元。RDMA-NAME:遠端直接記憶體存取 (RDMA) 網路的前置字元。您可以選擇任何前置字元。
設定基礎架構
在本節中,您將建立虛擬私有雲網路和 GKE 叢集。
Create a VPC network
為 gVNIC 介面建立虛擬私有雲網路:
gcloud compute networks create ${NETWORK} --subnet-mode=auto gcloud compute networks create ${GVNIC_NETWORK_PREFIX}-net \ --subnet-mode=custom gcloud compute networks subnets create ${GVNIC_NETWORK_PREFIX}-sub \ --network=${GVNIC_NETWORK_PREFIX}-net \ --region=${CONTROL_PLANE_REGION} \ --range=192.168.0.0/24 gcloud compute firewall-rules create ${GVNIC_NETWORK_PREFIX}-internal \ --network=${GVNIC_NETWORK_PREFIX}-net \ --action=ALLOW \ --rules=tcp:0-65535,udp:0-65535,icmp \ --source-ranges=192.168.0.0/16建立 RDMA 專用的虛擬私有雲網路和子網路,包括八個 GPU 各自的子網路:
gcloud compute networks create ${RDMA_NETWORK_PREFIX}-net \ --network-profile=${NODE_ZONE}-vpc-roce \ --subnet-mode=custom for N in $(seq 0 7); do gcloud compute networks subnets create ${RDMA_NETWORK_PREFIX}-sub-$N \ --network=${RDMA_NETWORK_PREFIX}-net \ --region=${CONTROL_PLANE_REGION} \ --range=192.168.$((N+1)).0/24 & done wait
建立 GKE 叢集
您可以在 GKE Standard 叢集中設定 NeMo RL。
建立標準叢集:
gcloud container clusters create ${CLUSTER_NAME} \ --location=${CONTROL_PLANE_REGION} \ --workload-pool=${PROJECT_ID}.svc.id.goog \ --enable-dataplane-v2 \ --enable-ip-alias \ --enable-multi-networking \ --addons=RayOperator,LustreCsiDriver \ --enable-legacy-lustre-port \ --machine-type=n2-highmem-80 \ --num-nodes=1 \ --min-nodes=1 \ --max-nodes=5 \ --enable-autoscaling \ --network=${NETWORK}取得叢集憑證:
gcloud container clusters get-credentials ${CLUSTER_NAME} \ --location=${CONTROL_PLANE_REGION}建立 GPU 節點集區:
gcloud container node-pools create gpu-pool \ --cluster=${CLUSTER_NAME} \ --location=${CONTROL_PLANE_REGION} \ --node-locations=${NODE_ZONE} \ --machine-type=${MACHINE_TYPE} \ --accelerator=type=${GPU_TYPE},count=8,gpu-driver-version=DEFAULT \ --reservation-affinity=specific \ --reservation=${RESERVATION} \ --enable-autoscaling \ --num-nodes=0 \ --total-max-nodes=2 \ --additional-node-network=network=${GVNIC_NETWORK_PREFIX}-net,subnetwork=${GVNIC_NETWORK_PREFIX}-sub \ --additional-node-network=network=${RDMA_NETWORK_PREFIX}-net,subnetwork=${RDMA_NETWORK_PREFIX}-sub-0 \ --additional-node-network=network=${RDMA_NETWORK_PREFIX}-net,subnetwork=${RDMA_NETWORK_PREFIX}-sub-1 \ --additional-node-network=network=${RDMA_NETWORK_PREFIX}-net,subnetwork=${RDMA_NETWORK_PREFIX}-sub-2 \ --additional-node-network=network=${RDMA_NETWORK_PREFIX}-net,subnetwork=${RDMA_NETWORK_PREFIX}-sub-3 \ --additional-node-network=network=${RDMA_NETWORK_PREFIX}-net,subnetwork=${RDMA_NETWORK_PREFIX}-sub-4 \ --additional-node-network=network=${RDMA_NETWORK_PREFIX}-net,subnetwork=${RDMA_NETWORK_PREFIX}-sub-5 \ --additional-node-network=network=${RDMA_NETWORK_PREFIX}-net,subnetwork=${RDMA_NETWORK_PREFIX}-sub-6 \ --additional-node-network=network=${RDMA_NETWORK_PREFIX}-net,subnetwork=${RDMA_NETWORK_PREFIX}-sub-7安裝 NCCL RDMA 安裝程式:
kubectl apply -f https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/refs/heads/master/gpudirect-rdma/nccl-rdma-installer.yaml
設定網路對應
將下列資訊清單儲存為
network-mapping.yaml:套用資訊清單:
envsubst < network-mapping.yaml > network-mapping-updated.yaml kubectl apply -f network-mapping-updated.yaml
準備儲存空間
在本節中,您將建立 Managed Lustre 執行個體,為 RL 工作負載佈建所需的高效能儲存空間。
為私人服務連線分配 IP 位址範圍:
gcloud compute addresses create ${LUSTRE_NAME}-range \ --global --purpose=VPC_PEERING \ --prefix-length=20 --network=${NETWORK}連結對等互連:
gcloud services vpc-peerings connect \ --service=servicenetworking.googleapis.com \ --ranges=${LUSTRE_NAME}-range \ --network=${NETWORK}建立 Managed Lustre 執行個體:
gcloud lustre instances create ${LUSTRE_NAME} \ --per-unit-storage-throughput=500 \ --capacity-gib=18000 \ --filesystem=lustrefs \ --location=${NODE_ZONE} \ --network=projects/${PROJECT_ID}/global/networks/${NETWORK} \ --gke-support-enabled使用 Managed Lustre CSI 驅動程式存取現有的 Managed Lustre 執行個體
擷取 Managed Lustre 執行個體的 IP 位址。
export LUSTRE_IP=$(gcloud lustre instances describe ${LUSTRE_NAME} \ --location=$NODE_ZONE --format="value(mountPoint)" | awk -F'@' '{print $1}')檢查
lustre-pv.yaml資訊清單。套用資訊清單:
envsubst < lustre-pv.yaml > lustre-pv-updated.yaml kubectl apply -f lustre-pv-updated.yaml檢查
lustre-pvc.yaml資訊清單。套用資訊清單:
kubectl apply -f lustre-pvc.yaml
部署 RayCluster
在本節中,您將複製範例存放區、準備資訊清單,並部署 Ray 叢集:
複製範例存放區:
git clone https://github.com/GoogleCloudPlatform/kubernetes-engine-samples.git cd kubernetes-engine-samples前往工作目錄:
cd ai-ml/nemo-rl-on-gke/nemoRL檢查
values.yaml資訊清單:根據您在本教學課程中使用的加速器,將
NCCL_TUNER_CONFIG_PATH取代為下列任一值:- NVIDIA B200 (180 GB):
/usr/local/gib/configs/tuner_config_a4.txtpb - NVIDIA H200 (141 GB):
/usr/local/gib/configs/tuner_config_a3u.txtpb
在這個資訊清單中,主要節點會管理工作並代管 Ray 資訊主頁,工作站節點則會執行訓練工作。
- NVIDIA B200 (180 GB):
部署 Ray 叢集:
export REPLICA_COUNT=2 helm install ray-cluster . \ --set additionalWorkerGroups.worker-grp-0.replicas=$REPLICA_COUNT在本教學課程中,您會使用兩個工作節點。如要變更工作站節點數量,請變更
REPLICA_COUNT值。確認工作節點和頭部節點正在執行:
kubectl get pods輸出結果會與下列內容相似:
NAME READY STATUS RESTARTS AGE ray-cluster-kuberay-head-sw7dp 3/3 Running 0 33h ray-cluster-kuberay-worker-grp-0-worker-gkbxw 3/3 Running 0 33h ray-cluster-kuberay-worker-grp-0-worker-kdg62 3/3 Running 0 33h確認 Ray 叢集正在執行:
kubectl ray get cluster輸出結果會與下列內容相似:
NAME NAMESPACE DESIRED WORKERS AVAILABLE WORKERS CPUS GPUS TPUS MEMORY CONDITION STATUS AGE ray-cluster-kuberay default 2 2 618 17 0 1573741824k RayClusterProvisioned ready 33h
啟動 GRPO 工作
Ray 叢集準備就緒後,您可以將 Ray 工作提交至 GKE 上執行的 Ray 叢集。NeMo RL 會在執行 RL 訓練工作期間自動下載模型。
如要提交 Ray 工作,請啟動互動式工作階段來執行工作。
如要建立與 Ray 叢集的本機連線,請執行下列指令:
kubectl ray session ray-cluster-kuberay這項指令會在本機與 GKE 叢集中的 Ray 主節點之間啟動通訊埠轉送功能。請注意,這個工作階段啟用時,終端機將處於忙碌狀態;如要繼續,請開啟另一個終端機執行個體。
編輯
gemma3-27b-gsm8k.sh檔案:在
gemma3-27b-gsm8k.sh檔案中替換下列值:YOUR_WANDB_API_KEY:您的 WandB API 金鑰。YOUR_HF_TOKEN:您的 Hugging Face 權杖。
在這個檔案中,您可以看到在 GSM8k 資料集上使用 gemma3-27b-it 模型執行 Job 的設定。如要完成 GRPO 訓練管道,這個指令碼會定義下列參數:
num_prompts_per_step: 16和num_generations_per_prompt: 32:Gemma3-27b-it 模型會為每個提示產生大量回應。在這個設定中,模型會產生 512 個回覆 (16 × 32 = 512)。policy.generation.colocated.enabled=False:這項參數會停用共置生成功能,也就是說,模型不會在與訓練程序相同的節點中生成回應。在標準 RL 中,訓練和生成作業都由相同的 GPU 處理。在這個 NeMo RL 設定中,您會專門將特定節點 (使用policy.generation.colocated.resources參數管理) 用於 vLLM 推論,而叢集的其餘部分則專注於高負載的訓練數學。將這些工作負載分開,可避免記憶體密集型訓練緩衝區與運算密集型推論工作負載爭搶資源。
如要提交 Job,請執行下列指令:
bash gemma3-27b-it/gemma3-27b-gsm8k.sh作業執行時,輸出內容會顯示訓練結果、時間和效能指標。
監控 GRPO 工作狀態
Ray 完成工作後,NeMo RL 會將檢查點儲存在設定的路徑中。
如要檢查 GRPO 工作的輸出內容,請建立
ray-head容器的 SSH 工作階段:kubectl exec -it $(kubectl get pods -l ray.io/node-type=head -o name) -c ray-head -- bash在
ray-head容器的終端機中安裝 apt tree 公用程式:apt update && apt install -y tree輸出結果會與下列內容相似:
root@ray-cluster-kuberay-worker-grp-0-worker-gkbxw:/opt/nemo-rl# tree /data/nemo_rl_gemma3_27b_3_17/ /data/nemo_rl_gemma3_27b_3_17/ `-- step_10 |-- config.yaml |-- policy | |-- optimizer | | |-- __0_0.distcp | | |-- __10_0.distcp | | |-- __11_0.distcp | | |-- __12_0.distcp | | |-- __13_0.distcp | | |-- __14_0.distcp | | |-- __15_0.distcp | | |-- __1_0.distcp | | |-- __2_0.distcp | | |-- __3_0.distcp | | |-- __4_0.distcp | | |-- __5_0.distcp | | |-- __6_0.distcp | | |-- __7_0.distcp | | |-- __8_0.distcp | | `-- __9_0.distcp | |-- tokenizer | | |-- chat_template.jinja | | |-- special_tokens_map.json | | |-- tokenizer.json | | `-- tokenizer_config.json | `-- weights | |-- __0_0.distcp | |-- __10_0.distcp | |-- __11_0.distcp | |-- __12_0.distcp | |-- __13_0.distcp | |-- __14_0.distcp | |-- __15_0.distcp | |-- __1_0.distcp | |-- __2_0.distcp | |-- __3_0.distcp | |-- __4_0.distcp | |-- __5_0.distcp | |-- __6_0.distcp | |-- __7_0.distcp | |-- __8_0.distcp | `-- __9_0.distcp |-- train_dataloader.pt `-- training_info.json 6 directories, 39 files ```
清除所用資源
為避免產生費用,請刪除資源:
helm delete ray-cluster
gcloud container clusters delete ${CLUSTER_NAME} \
--location=${CONTROL_PLANE_REGION} \
--quiet
gcloud lustre instances delete ${LUSTRE_NAME} --location=${NODE_ZONE} --quiet
gcloud services vpc-peerings delete \
--service=servicenetworking.googleapis.com \
--network=${NETWORK}
gcloud compute addresses delete ${LUSTRE_NAME}-range --global --quiet
gcloud compute firewall-rules delete ${GVNIC_NETWORK_PREFIX}-internal --quiet
gcloud compute networks subnets delete ${GVNIC_NETWORK_PREFIX}-sub \
--region=${CONTROL_PLANE_REGION} --quiet
gcloud compute networks delete ${GVNIC_NETWORK_PREFIX}-net --quiet
for N in $(seq 0 7); do
gcloud compute networks subnets delete ${RDMA_NETWORK_PREFIX}-sub-$N \
--region=${CONTROL_PLANE_REGION} --quiet &
done
wait
gcloud compute networks delete ${RDMA_NETWORK_PREFIX}-net --quiet
gcloud compute networks delete ${NETWORK} --quiet