GKE에서 NVIDIA NeMo RL을 사용하여 강화 학습 미세 조정 및 확장

이 튜토리얼에서는 Google Kubernetes Engine (GKE)에서 강화 학습 (RL)을 위한 분산 학습 환경을 오케스트레이션하는 방법을 보여줍니다. Ray와 NVIDIA NeMo RL 프레임워크를 사용하여 모델을 미세 조정하기 위한 분산 학습 환경을 설정합니다.

이 튜토리얼에서는 Ray 및 NeMo RL을 사용하여 GKE에서 그룹 상대 정책 최적화 (GRPO) 학습 파이프라인에 중점을 둡니다. GRPO는 모델의 추론 능력을 향상하도록 설계된 강화 학습 알고리즘입니다. 이 메모리 효율적인 알고리즘은 Critic 또는 값 모델을 제거하고 대신 상대적인 그룹 기반 계산을 사용하여 RL 프로세스를 간소화합니다.

이 튜토리얼을 실행하기 전에 GKE에서 verl을 사용하여 강화 학습 미세 조정 및 확장 튜토리얼을 완료하는 것이 좋습니다. 다음 튜토리얼에서는 미세 조정 및 확장 RL with verl 튜토리얼과 동일한 클러스터 설정 및 구성을 사용합니다.

배경

다음 섹션에서는 이 튜토리얼에서 사용되는 개념을 간략하게 설명합니다.

강화 학습 (RL)

RL은 정적인 모방이 아닌 경험, 탐색, 피드백을 통해 모델을 학습시킵니다. 사전 학습을 통해 모델에 말해야 할 내용을 가르치는 반면, 인간 피드백 기반 강화 학습 (RLHF)을 통해 유용하고 안전하며 논리적인 방법을 가르칩니다. RL은 기본 모델과 특수 사용 사례를 위해 미세 조정된 모델 간의 다리 역할을 합니다.

자세한 내용은 강화 학습이란 무엇인가요?를 참고하세요.

그룹 상대 정책 최적화 (GRPO)

DeepSeek에서 널리 사용되는 알고리즘인 GRPO는 비평가 모델을 삭제하여 LLM 정렬을 위한 근위 정책 최적화 (PPO)의 메모리 효율적인 대안을 제공합니다. 비평가 네트워크 대신 GRPO는 동일한 프롬프트에 대한 응답 그룹을 생성하고 해당 그룹의 평균 보상을 기준선으로 사용합니다.

자세한 내용은 GRPO를 참고하세요.

NVIDIA NeMo RL

NeMo RL은 확장 가능한 RL을 위해 설계된 NVIDIA의 오픈소스 사후 학습 라이브러리입니다. 더 광범위한 NeMo 프레임워크 생태계의 일부인 NeMo RL은 단일 GPU에서 소규모 실험과 수천 개의 GPU에 걸친 다중 노드 배포를 모두 지원합니다.

자세한 내용은 NVIDIA NeMo RL을 참고하세요.

GSM8k 데이터 세트

이 튜토리얼에서는 8,500개의 고품질의 언어적으로 다양한 초등학교 수학 단어 문제로 구성된 GSM8k 데이터 세트를 사용합니다.

GSM8k 및 GRPO를 사용하여 모델은 동일한 문제에 대해 n개의 서로 다른 대답을 생성합니다. GRPO는 이러한 응답을 그룹 평균과 비교합니다. 모델은 나머지 그룹에 비해 일관되게 정확하고 논리적으로 타당한 경로에 대해 더 많은 보상을 받습니다. 시간이 지남에 따라 모델은 단계를 명확하게 설명하는 것이 보상을 극대화하는 가장 신뢰할 수 있는 방법이라는 것을 학습하여 실적이 저조한 답변에 대한 보상을 효과적으로 줄입니다.

자세한 내용은 GSM8k를 참고하세요.

목표

이 튜토리얼에서는 다음 단계를 완료하여 GKE에서 NeMo RL로 RL을 설정하는 방법을 보여줍니다.

  1. 환경을 준비합니다.
  2. B200 또는 H200 GPU가 있는 GKE 클러스터를 설정합니다.
  3. 분산 Ray 클러스터를 관리하도록 KubeRay를 구성합니다.
  4. 고성능 스토리지를 위해 Managed Lustre를 사용합니다.
  5. NeMo RL을 사용하는 GRPO 학습 작업을 실행합니다.

시작하기 전에

  • Google Cloud 계정에 로그인합니다. Google Cloud를 처음 사용하는 경우 계정을 만들고 Google 제품의 실제 성능을 평가해 보세요. 신규 고객에게는 워크로드를 실행, 테스트, 배포하는 데 사용할 수 있는 $300의 무료 크레딧이 제공됩니다.

  • Google Cloud CLI를 설치합니다.

  • 외부 ID 공급업체(IdP)를 사용하는 경우 먼저 제휴 ID로 gcloud CLI에 로그인해야 합니다.

  • gcloud CLI를 초기화하려면, 다음 명령어를 실행합니다. 

    gcloud init

  • Google Cloud 프로젝트를 만들거나 선택합니다.

    프로젝트를 선택하거나 만드는 데 필요한 역할

    • 프로젝트 선택: 프로젝트를 선택하는 데는 특정 IAM 역할이 필요하지 않습니다. 역할이 부여된 프로젝트를 선택하면 됩니다.
    • 프로젝트 만들기: 프로젝트를 만들려면 resourcemanager.projects.create 권한이 포함된 프로젝트 생성자 역할(roles/resourcemanager.projectCreator)이 필요합니다. 역할 부여 방법 알아보기

    • Google Cloud 프로젝트를 만듭니다.

      gcloud projects create PROJECT_ID

      PROJECT_ID를 만들려는 Google Cloud 프로젝트의 이름으로 바꿉니다.

    • 생성한 Google Cloud 프로젝트를 선택합니다.

      gcloud config set project PROJECT_ID

      PROJECT_ID을 Google Cloud 프로젝트 이름으로 바꿉니다.

  • Google Cloud 프로젝트에 결제가 사용 설정되어 있는지 확인합니다.

  • 필요한 API를 사용 설정합니다.

    API 사용 설정에 필요한 역할

    API를 사용 설정하려면 serviceusage.services.enable 권한이 포함된 서비스 사용량 관리자 IAM 역할 (roles/serviceusage.serviceUsageAdmin)이 필요합니다. 역할 부여 방법 알아보기 

    gcloud services enable container.googleapis.com storage.googleapis.com compute.googleapis.com

  • Google Cloud CLI를 설치합니다.

  • 외부 ID 공급업체(IdP)를 사용하는 경우 먼저 제휴 ID로 gcloud CLI에 로그인해야 합니다.

  • gcloud CLI를 초기화하려면, 다음 명령어를 실행합니다. 

    gcloud init

  • Google Cloud 프로젝트를 만들거나 선택합니다.

    프로젝트를 선택하거나 만드는 데 필요한 역할

    • 프로젝트 선택: 프로젝트를 선택하는 데는 특정 IAM 역할이 필요하지 않습니다. 역할이 부여된 프로젝트를 선택하면 됩니다.
    • 프로젝트 만들기: 프로젝트를 만들려면 resourcemanager.projects.create 권한이 포함된 프로젝트 생성자 역할(roles/resourcemanager.projectCreator)이 필요합니다. 역할 부여 방법 알아보기

    • Google Cloud 프로젝트를 만듭니다.

      gcloud projects create PROJECT_ID

      PROJECT_ID를 만들려는 Google Cloud 프로젝트의 이름으로 바꿉니다.

    • 생성한 Google Cloud 프로젝트를 선택합니다.

      gcloud config set project PROJECT_ID

      PROJECT_ID을 Google Cloud 프로젝트 이름으로 바꿉니다.

  • Google Cloud 프로젝트에 결제가 사용 설정되어 있는지 확인합니다.

  • 필요한 API를 사용 설정합니다.

    API 사용 설정에 필요한 역할

    API를 사용 설정하려면 serviceusage.services.enable 권한이 포함된 서비스 사용량 관리자 IAM 역할 (roles/serviceusage.serviceUsageAdmin)이 필요합니다. 역할 부여 방법 알아보기 

    gcloud services enable container.googleapis.com storage.googleapis.com compute.googleapis.com
    •

  • 사용자 계정에 역할을 부여합니다. 다음 IAM 역할마다 다음 명령어를 1회 실행합니다. roles/container.admin, roles/iam.serviceAccountAdmin, roles/storage.admin 

    gcloud projects add-iam-policy-binding PROJECT_ID --member="user:USER_IDENTIFIER" --role=ROLE

    다음을 바꿉니다.

    • PROJECT_ID: 프로젝트 ID입니다.
    • USER_IDENTIFIER: 사용자 계정의 식별자입니다. 예를 들면 myemail@example.com입니다.
    • ROLE: 사용자 계정에 부여하는 IAM 역할입니다.

개발 환경 준비

이 튜토리얼에서는 Cloud Shell을 사용합니다.

  1. Google Cloud 콘솔로 이동합니다.

  2. Google Cloud 콘솔 창 상단에서 Cloud Shell 활성화 버튼을 클릭합니다.

  3. 다음 환경 변수를 설정합니다.

    export PROJECT_ID=$(gcloud config get project)
export PROJECT_NUMBER=$(gcloud projects describe ${PROJECT_ID} --format="value(projectNumber)")
export CONTROL_PLANE_LOCATION=CONTROL_PLANE_LOCATION
export NODE_LOCATION=NODE_LOCATION
export CLUSTER_NAME=CLUSTER_NAME
export GPU_TYPE=GPU_TYPE
export MACHINE_TYPE=MACHINE_TYPE
export GKE_VERSION=GKE_VERSION
export KSA_NAME=generic-ksa
export NAMESPACE=default
export GS_BUCKET=BUCKET_NAME-${PROJECT_ID}
export HF_TOKEN=YOUR_HUGGING_FACE_TOKEN

    다음 값을 바꿉니다.

    • CLUSTER_NAME: GKE 클러스터의 이름입니다.
    • CONTROL_PLANE_LOCATION: GKE 클러스터 컨트롤 플레인의 Compute Engine 리전입니다.
    • NODE_LOCATION: 노드의 위치입니다. NVIDIA B200 또는 H200 GPU를 사용할 수 있는 영역을 선택합니다.
    • GPU_TYPE: Compute Engine 용량 예약에서 예약한 가속기. 다음 값 중 하나여야 합니다.
      • nvidia-b200: NVIDIA B200 (180GB)
      • nvidia-h200-141gb: NVIDIA H200 (141GB)
    • MACHINE_TYPE: 사용할 머신 유형입니다.
      • NVIDIA B200 (180GB) GPU의 경우 a4-highgpu-8g 이상을 사용합니다.
      • NVIDIA H200 (141GB) GPU의 경우 a3-ultragpu-8g 이상을 사용합니다.

    • GKE_VERSION: 사용할 GKE 버전입니다.

      • NVIDIA B200 (180GB) GPU의 경우 1.32.2-gke.1422000 이상을 사용합니다.
      • NVIDIA H200 (141GB) GPU의 경우 1.31.4-gke.1183000 이상을 사용합니다.

    • BUCKET_NAME: Cloud Storage 버킷의 기본 이름입니다.

    • YOUR_HUGGING_FACE_TOKEN: Hugging Face 토큰입니다.

  4. 네트워크에 다음 환경 변수를 만듭니다.

    export GVNIC_NETWORK_PREFIX="GVNIC-NAME"
export RDMA_NETWORK_PREFIX="RDMA-NAME"

    다음 값을 바꿉니다.

    • GVNIC-NAME: gVNIC 네트워크 이름의 접두사입니다. 원하는 접두사를 사용할 수 있습니다.
    • RDMA-NAME: 원격 직접 메모리 액세스 (RDMA) 네트워크의 접두사입니다. 원하는 접두사를 사용할 수 있습니다.

인프라 설정

이 섹션에서는 VPC 네트워크와 GKE 클러스터를 만듭니다.

VPC 네트워크 만들기

  1. gVNIC 인터페이스용 VPC 네트워크를 만듭니다.

    gcloud compute networks create ${GVNIC_NETWORK_PREFIX}-net \
    --project=${PROJECT_ID} \
    --subnet-mode=custom
gcloud compute networks subnets create ${GVNIC_NETWORK_PREFIX}-sub \
    --network=${GVNIC_NETWORK_PREFIX}-net \
    --location=${CONTROL_PLANE_LOCATION} \
    --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

  2. 8개의 GPU용 8개의 서브넷이 포함된 RDMA용 VPC 네트워크 및 서브넷을 만듭니다.

    gcloud compute networks create ${RDMA_NETWORK_PREFIX}-net \
    --network-profile=${CONTROL_PLANE_LOCATION}-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 \
    --location=${CONTROL_PLANE_LOCATION} \
    --range=192.168.$((N+1)).0/24 &
done
wait

GKE 클러스터 만들기

GKE Autopilot 또는 Standard 클러스터에서 NeMo RL을 설정할 수 있습니다. 완전 관리형 Kubernetes 환경을 위해서는 Autopilot을 사용하는 것이 좋습니다. 워크로드에 가장 적합한 GKE 작업 모드를 선택하려면 GKE 작업 모드 정보를 참고하세요.

Autopilot

  1. Autopilot 클러스터를 만듭니다.

    gcloud container clusters create-auto ${CLUSTER_NAME} \
    --location=${CONTROL_PLANE_LOCATION} \
    --enable-multi-networking  \
    --enable-ray-operator

  2. 클러스터의 사용자 인증 정보를 가져옵니다.

    gcloud container clusters get-credentials ${CLUSTER_NAME} \
    --location=${CONTROL_PLANE_LOCATION}

  3. Autopilot용 NCCL RDMA 설치 프로그램을 설치합니다.

    kubectl apply -f https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/refs/heads/master/gpudirect-rdma/nccl-rdma-installer-autopilot.yaml

표준

  1. 표준 클러스터 만들기

    gcloud container clusters create ${CLUSTER_NAME} \
    --location=${CONTROL_PLANE_LOCATION} \
    --enable-dataplane-v2 \
    --enable-ip-alias \
    --enable-multi-networking \
    --addons=RayOperator \
    --num-nodes=1

  2. 클러스터의 사용자 인증 정보를 가져옵니다.

    gcloud container clusters get-credentials ${CLUSTER_NAME} \
    --location=${CONTROL_PLANE_LOCATION}

  3. GPU 노드 풀을 만듭니다.

    gcloud container node-pools create gpu-pool \
    --cluster=${CLUSTER_NAME} \
    --node-locations=${NODE_LOCATION} \
    --machine-type=${MACHINE_TYPE} \
    --accelerator=type=${GPU_TYPE},count=8 \
    --spot \
    --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

  4. NCCL RDMA 설치 프로그램을 설치합니다.

    kubectl apply -f https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/refs/heads/master/gpudirect-rdma/nccl-rdma-installer.yaml

네트워크 매핑 구성

  1. 다음 매니페스트를 network-mapping.yaml로 저장합니다.

    # Copyright 2026 Google LLC. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
apiVersion: networking.gke.io/v1
kind: GKENetworkParamSet
metadata:
  name: gvnic-1
spec:
  vpc: ${GVNIC_NETWORK_PREFIX}-net
  vpcSubnet: ${GVNIC_NETWORK_PREFIX}-sub
  deviceMode: NetDevice
---
apiVersion: networking.gke.io/v1
kind: Network
metadata:
  name: gvnic-1
spec:
  type: "Device"
  parametersRef:
    group: networking.gke.io
    kind: GKENetworkParamSet
    name: gvnic-1
---
apiVersion: networking.gke.io/v1
kind: GKENetworkParamSet
metadata:
  name: rdma-0
spec:
  vpc: ${RDMA_NETWORK_PREFIX}-net
  vpcSubnet: ${RDMA_NETWORK_PREFIX}-sub-0
  deviceMode: RDMA
---
apiVersion: networking.gke.io/v1
kind: Network
metadata:
  name: rdma-0
spec:
  type: "Device"
  parametersRef:
    group: networking.gke.io
    kind: GKENetworkParamSet
    name: rdma-0
---
apiVersion: networking.gke.io/v1
kind: GKENetworkParamSet
metadata:
  name: rdma-1
spec:
  vpc: ${RDMA_NETWORK_PREFIX}-net
  vpcSubnet: ${RDMA_NETWORK_PREFIX}-sub-1
  deviceMode: RDMA
---
apiVersion: networking.gke.io/v1
kind: Network
metadata:
  name: rdma-1
spec:
  type: "Device"
  parametersRef:
    group: networking.gke.io
    kind: GKENetworkParamSet
    name: rdma-1
---
apiVersion: networking.gke.io/v1
kind: GKENetworkParamSet
metadata:
  name: rdma-2
spec:
  vpc: ${RDMA_NETWORK_PREFIX}-net
  vpcSubnet: ${RDMA_NETWORK_PREFIX}-sub-2
  deviceMode: RDMA
---
apiVersion: networking.gke.io/v1
kind: Network
metadata:
  name: rdma-2
spec:
  type: "Device"
  parametersRef:
    group: networking.gke.io
    kind: GKENetworkParamSet
    name: rdma-2
---
apiVersion: networking.gke.io/v1
kind: GKENetworkParamSet
metadata:
  name: rdma-3
spec:
  vpc: ${RDMA_NETWORK_PREFIX}-net
  vpcSubnet: ${RDMA_NETWORK_PREFIX}-sub-3
  deviceMode: RDMA
---
apiVersion: networking.gke.io/v1
kind: Network
metadata:
  name: rdma-3
spec:
  type: "Device"
  parametersRef:
    group: networking.gke.io
    kind: GKENetworkParamSet
    name: rdma-3
---
apiVersion: networking.gke.io/v1
kind: GKENetworkParamSet
metadata:
  name: rdma-4
spec:
  vpc: ${RDMA_NETWORK_PREFIX}-net
  vpcSubnet: ${RDMA_NETWORK_PREFIX}-sub-4
  deviceMode: RDMA
---
apiVersion: networking.gke.io/v1
kind: Network
metadata:
  name: rdma-4
spec:
  type: "Device"
  parametersRef:
    group: networking.gke.io
    kind: GKENetworkParamSet
    name: rdma-4
---
apiVersion: networking.gke.io/v1
kind: GKENetworkParamSet
metadata:
  name: rdma-5
spec:
  vpc: ${RDMA_NETWORK_PREFIX}-net
  vpcSubnet: ${RDMA_NETWORK_PREFIX}-sub-5
  deviceMode: RDMA
---
apiVersion: networking.gke.io/v1
kind: Network
metadata:
  name: rdma-5
spec:
  type: "Device"
  parametersRef:
    group: networking.gke.io
    kind: GKENetworkParamSet
    name: rdma-5
---
apiVersion: networking.gke.io/v1
kind: GKENetworkParamSet
metadata:
  name: rdma-6
spec:
  vpc: ${RDMA_NETWORK_PREFIX}-net
  vpcSubnet: ${RDMA_NETWORK_PREFIX}-sub-6
  deviceMode: RDMA
---
apiVersion: networking.gke.io/v1
kind: Network
metadata:
  name: rdma-6
spec:
  type: "Device"
  parametersRef:
    group: networking.gke.io
    kind: GKENetworkParamSet
    name: rdma-6
---
apiVersion: networking.gke.io/v1
kind: GKENetworkParamSet
metadata:
  name: rdma-7
spec:
  vpc: ${RDMA_NETWORK_PREFIX}-net
  vpcSubnet: ${RDMA_NETWORK_PREFIX}-sub-7
  deviceMode: RDMA
---
apiVersion: networking.gke.io/v1
kind: Network
metadata:
  name: rdma-7
spec:
  type: "Device"
  parametersRef:
    group: networking.gke.io
    kind: GKENetworkParamSet
    name: rdma-7

  2. 매니페스트를 적용합니다.

    kubectl apply -f network-mapping.yaml

스토리지 준비

이 섹션에서는 RL 워크로드에 필요한 고성능 스토리지를 프로비저닝하는 Cloud Storage 버킷과 관리형 Lustre 인스턴스를 만듭니다.

  1. Cloud Storage 버킷을 만듭니다.

    gcloud storage buckets create gs://${GS_BUCKET} \
    --location=${CONTROL_PLANE_LOCATION} \
    --enable-hierarchical-namespace \
    --uniform-bucket-level-access

  2. Kubernetes 서비스 계정 (KSA)을 만들고 버킷에 바인딩합니다.

    kubectl create serviceaccount ${KSA_NAME} --namespace ${NAMESPACE}

gcloud storage buckets add-iam-policy-binding gs://${GS_BUCKET} \
    --member "principal://iam.googleapis.com/projects/${PROJECT_NUMBER}/locations/global/workloadIdentityPools/${PROJECT_ID}.svc.id.goog/subject/ns/${NAMESPACE}/sa/${KSA_NAME}" \
    --role "roles/storage.objectUser"

  3. 다음 단계를 완료하여 관리형 Lustre를 설정합니다.

    1. Managed Lustre 인스턴스 만들기의 단계에 따라 Managed Lustre 인스턴스를 만듭니다. 인스턴스가 GKE 클러스터와 동일한 네트워크를 사용하는지 확인합니다.
    2. 기존 Managed Lustre 인스턴스 액세스의 단계에 따라 Managed Lustre 인스턴스에 액세스합니다.

RayCluster 배포

이 섹션에서는 샘플 저장소를 클론하고, 매니페스트를 준비하고, launcher.sh 스크립트를 실행합니다.

  1. 샘플 저장소를 클론합니다.

    git clone https://github.com/GoogleCloudPlatform/kubernetes-engine-samples.git
cd kubernetes-engine-samples

  2. 작업 디렉터리로 이동합니다.

    cd ai-ml/nemo-rl-on-gke/nemoRL

  3. values.yaml 매니페스트를 검사합니다.

    # Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

image:
  repository: "nvcr.io/nvidia/nemo-rl"
  tag: "v0.5.0" 
  pullPolicy: Always

nameOverride: "kuberay"
fullnameOverride: ""

common:
  containerEnv: {}

configMap:
  fluentbit:
    data:
      fluent-bit.conf: |
        [INPUT]
            Name              tail
            Path              /tmp/ray/session_latest/logs/worker-*
            Tag               ray-worker
        [INPUT]
            Name              tail
            Path              /tmp/ray/session_latest/logs/raylet*
            Tag               raylet
        [INPUT]
            Name              tail
            Path              /tmp/ray/session_latest/logs/*
            Exclude_Path      /tmp/ray/session_latest/logs/debug_state.txt,/tmp/ray/session_latest/logs/raylet*,/tmp/ray/session_latest/logs/worker-*
            Tag               ray-misc
        [OUTPUT]
            Name              stackdriver
            Match             *
            resource          gce_instance
            labels_key        labels

# --- Head Node Configuration ---
head:
  enableInTreeAutoscaling: false
  serviceAccountName: ""
  rayStartParams:
    dashboard-host: '0.0.0.0'
  template:
    metadata:
      annotations:
        gke-gcsfuse/volumes: "true"
        networking.gke.io/default-interface: 'eth0'
  containerEnv:
  - name: RAY_GROUP
    value: "head"
  resources:
    limits:
      cpu: "64"
      memory: "500G"
      nvidia.com/gpu: 0
    requests:
      cpu: "64"
      memory: "500G"
      nvidia.com/gpu: 0
  tolerations:
    # - operator: "Exists"
    #   key: "components.gke.io/gke-managed-components"
    # - key: "nvidia.com/gpu"
    #   operator: "Exists"
    #   effect: "NoSchedule"
  volumeMounts:
    - mountPath: /data
      name: lustre-data

  volumes:
    - name: log-volume
      emptyDir: {}
    - name: fluentbit-config-volume
      configMap:
        name: "ray-cluster-kuberay-fluentbit-config"
    - name: lustre-data
      persistentVolumeClaim:
        claimName: lustre-pvc
  sidecarContainers:
    - name: fluent-bit
      image: fluent/fluent-bit:latest
      env:
      - name: RAY_GROUP
        value: "head"
      volumeMounts:
        - name: fluentbit-config-volume
          mountPath: /fluent-bit/etc/
        - mountPath: /tmp/ray
          name: log-volume

  # --- HEAD POD STARTUP SCRIPT ---
  command:
    - "bash"
    - "-c"
    - |
      set -ex
      echo "--- Head Pod Setup ---"
      apt-get update
      apt-get install -y sudo netcat-openbsd pciutils
      cd /opt/nemo-rl
      /usr/bin/python -m pip install uv
      /usr/bin/python -m uv venv
      echo "Head pod setup complete. Starting Ray..."

      exec ${KUBERAY_GEN_RAY_START_CMD}

  args: []
  headService: {}
  # nodeSelector:
  #   cloud.google.com/gke-accelerator: nvidia-b200 #cloud.google.com/gke-nodepool: cpu-node-pool-llama #cpu-node-pool

# --- Default Worker (Disabled) ---
worker:
  disabled: true

# --- A4 GPU Worker Groups ---
additionalWorkerGroups:
  worker-grp-0:
    disabled: false
    replicas: 4
    annotations:
      networking.gke.io/default-interface: 'eth0'
      networking.gke.io/interfaces: |
        [
          {"interfaceName":"eth0","network":"default"},
          {"interfaceName":"eth1","network":"gvnic-1"},
          {"interfaceName":"eth2","network":"rdma-0"},
          {"interfaceName":"eth3","network":"rdma-1"},
          {"interfaceName":"eth4","network":"rdma-2"},
          {"interfaceName":"eth5","network":"rdma-3"},
          {"interfaceName":"eth6","network":"rdma-4"},
          {"interfaceName":"eth7","network":"rdma-5"},
          {"interfaceName":"eth8","network":"rdma-6"},
          {"interfaceName":"eth9","network":"rdma-7"}
        ]
    containerEnv:
      - name: RAY_GROUP
        valueFrom:
          fieldRef:
            fieldPath: metadata.labels['ray.io/group']
      - name: NCCL_NET  
        value: "gIB"
      - name: NCCL_IB_GID_INDEX
        value: "3"   
      - name: GLOO_SOCKET_IFNAME
        value: "eth0"
      - name: NCCL_CROSS_NIC
        value: "0"
      - name: NCCL_SOCKET_IFNAME
        value: "eth0"
      - name: TP_SOCKET_IFNAME # Specific to DTensor/PyTorch Distributed
        value: "eth0"
      - name: NCCL_TUNER_CONFIG_PATH
        value: "/usr/local/gib/configs/tuner_config_a4.txtpb"
      - name: NCCL_NET_GDR_LEVEL
        value: "PIX"
      - name: LD_LIBRARY_PATH
        value: /usr/local/nvidia/lib64
    resources:
      limits:
        nvidia.com/gpu: 8
        cpu: "206"
        memory: "2400Gi"
      requests:
        nvidia.com/gpu: 8
        cpu: "206"
        memory: "2400Gi"

    nodeSelector:
      cloud.google.com/gke-accelerator: nvidia-b200
    tolerations:
      - operator: "Exists"
        key: "nvidia.com/gpu"
      - operator: "Exists"
        key: "cloud.google.com/impending-node-termination"
      - operator: "Exists"
        key: "user-workload"
    securityContext:
      privileged: true
    volumes:
      - name: log-volume
        emptyDir: {}
      - name: shared-memory
        emptyDir:
          medium: "Memory"
          sizeLimit: 240Gi
      - name: ray-tmp
        emptyDir:
          medium: "Memory"
      - name: fluentbit-config-volume
        configMap:
          name: "ray-cluster-kuberay-fluentbit-config"
      - name: nvidia-install-dir-host
        hostPath:
          path: /home/kubernetes/bin/nvidia
      - name: gib-nccl-plugin-volume
        hostPath: 
          path: /home/kubernetes/bin/gib
      - name: lustre-data
        persistentVolumeClaim:
          claimName: lustre-pvc
    volumeMounts:
      - mountPath: /tmp/ray
        name: log-volume
      - name: shared-memory
        mountPath: /dev/shm
      - name: nvidia-install-dir-host
        mountPath: /usr/local/nvidia
      - name: gib-nccl-plugin-volume
        mountPath: /usr/local/gib
      - mountPath: /data
        name: lustre-data   
    # --- WORKER POD STARTUP SCRIPT ---
    command:
      - "bash"
      - "-c"
      - |
        set -ex

        echo "--- Worker Pod Setup ---"
        apt-get update
        apt-get install -y sudo netcat-openbsd pciutils
        cd /opt/nemo-rl
        /usr/bin/python -m pip install uv
        /usr/bin/python -m uv venv

        ldconfig /usr/local/nvidia/lib64/
        ldconfig -p | grep libcuda | sed 's/^/  /'
        export LD_LIBRARY_PATH="/usr/local/gib/lib64:$LD_LIBRARY_PATH"
        source /usr/local/gib/scripts/set_nccl_env.sh

        echo "Worker pod setup complete. Starting Ray..."

        exec ${KUBERAY_GEN_RAY_START_CMD}


    sidecarContainers:
      - name: fluent-bit
        env:
          - name: RAY_GROUP
            valueFrom:
              fieldRef:
                fieldPath: metadata.labels['ray.io/group']
        image: fluent/fluent-bit:latest
        volumeMounts:
          - name: fluentbit-config-volume
            mountPath: /fluent-bit/etc/
          - mountPath: /tmp/ray
            name: log-volume

# --- Service Config ---
service:
  type: ClusterIP

    이 튜토리얼에서 사용하는 액셀러레이터에 따라 NCCL_TUNER_CONFIG_PATH를 다음 값 중 하나로 바꿉니다.

    • NVIDIA B200 (180GB): /usr/local/gib/configs/tuner_config_a4.txtpb
    • NVIDIA H200 (141GB): /usr/local/gib/configs/tuner_config_a3u.txtpb

    이 매니페스트에서 헤드 노드는 작업을 관리하고 Ray 대시보드를 호스팅합니다. 워커 노드는 학습 작업을 실행합니다.

  4. Ray 클러스터를 설치합니다.

    export REPLICA_COUNT=2
helm install ray-cluster . \
  --set values.additionalWorkerGroups.worker-grp-0.replicas=$REPLICA_COUNT

    이 튜토리얼에서는 작업자 노드 두 개를 사용합니다. 작업자 노드 수를 변경하려면 REPLICA_COUNT 값을 변경합니다.

  5. Ray 클러스터를 배포하려면 launcher.sh 스크립트를 실행합니다.

    bash launcher.sh

  6. 작업자 및 헤드 노드가 실행 중인지 확인합니다.

    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

  7. Ray 클러스터가 실행 중인지 확인합니다.

    kubectl ray get clusters

    출력은 다음과 비슷합니다.

    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 클러스터가 준비되면 GKE에서 실행 중인 Ray 클러스터에 Ray 작업을 제출할 수 있습니다. NeMo RL은 RL 학습 작업을 실행하는 동안 모델을 자동으로 다운로드합니다.

Ray 작업을 제출하려면 대화형 세션을 시작하여 작업을 실행합니다.

  1. Ray 클러스터에 로컬 연결을 설정하려면 다음 명령어를 실행합니다.

      kubectl ray session ray-cluster-kuberay

    이 명령어는 로컬 머신과 GKE 클러스터의 Ray 헤드 노드 간에 포트 전달을 시작합니다. 이 세션이 활성 상태인 동안에는 터미널이 사용 중입니다. 계속하려면 별도의 터미널 인스턴스를 여세요.

  2. gemma3-27b-gsm8k.sh 파일을 수정합니다.

    # Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

#!/bin/bash
WANDB_API_KEY='YOUR_WANDB_API_KEY' # Update this with your WANDB API key
HF_TOKEN='YOUR_HF_TOKEN' # Update this with your HF token
WORLD_SIZE=16

# --- Step 1: Find the Ray Head Pod ---
echo "Finding Ray head pod..."
export HEAD_POD_NAME=$(kubectl get pods --selector=ray.io/node-type=head -o jsonpath='{.items[0].metadata.name}')
if [ -z "$HEAD_POD_NAME" ]; then
    echo "Error: No running Ray head pod found. Please check your cluster."
    exit 1
fi
echo "Found head pod: $HEAD_POD_NAME"
echo ""

# --- Step 2: Define the Job Script to Run ---
# This is the script that will be executed *inside* the head pod.
# It assumes the 'uv venv' setup from the values.yaml is already done.
JOB_SCRIPT=$(cat <<EOF
set -ex

echo "--- Running on Ray Head Pod ($HOSTNAME) ---"
cd /opt/nemo-rl

git pull && git checkout main

sed -i 's/subset: Optional\[str\] = None/subset: Optional[str] = "main"/' /opt/nemo-rl/nemo_rl/data/datasets/response_datasets/response_dataset.py
sed -i 's/raw_dataset = load_dataset(data_path)/raw_dataset = load_dataset(data_path, "main")/' /opt/nemo-rl/nemo_rl/data/datasets/utils.py

echo "Setting environment variables..."
export WANDB_API_KEY=$WANDB_API_KEY
export HF_TOKEN=$HF_TOKEN
export HF_HOME=/opt/nemo-rl/

###-----Example to launch Gemma3-27B on 3 nodes (24 GPUs)----------
uv run python examples/run_grpo_math.py \
  --config examples/configs/recipes/llm/grpo-gemma3-27b-it-8n4g-fsdp2tp4-actckpt-long.yaml \
  cluster.num_nodes=2 \
  cluster.gpus_per_node=8 \
  grpo.max_num_steps=300 \
  checkpointing.checkpoint_dir=/data/nemo_rl_gemma3_27b_3_17 \
  data.dataset_name=ResponseDataset \
  +data.train_data_path=openai/gsm8k \
  +data.val_data_path=openai/gsm8k \
  +data.val_split=test \
  +data.train_split=train \
  +data.subset="main" \
  +data.input_key="question" \
  +data.output_key="answer" \
  logger.tensorboard_enabled=False \
  logger.wandb_enabled=True \
  logger.wandb.name='nemo_rl_gemma3_27b_3_17' \
  grpo.num_prompts_per_step=16 \
  grpo.num_generations_per_prompt=64 \
  policy.generation.colocated.enabled=False \
  policy.generation.colocated.resources.num_nodes=1 \
  policy.generation.colocated.resources.gpus_per_node=8 \
  policy.generation.vllm_cfg.tensor_parallel_size=8 \
  policy.generation.vllm_cfg.gpu_memory_utilization=0.9 \
  policy.dtensor_cfg.tensor_parallel_size=8

echo "--- Job Finished ---"
EOF
)

# --- Step 3: Execute the Job ---
echo "Submitting job to $HEAD_POD_NAME..."
echo "$JOB_SCRIPT" | tr -d '\r' | kubectl exec -i $HEAD_POD_NAME -c ray-head -- /bin/bash

echo ""
echo "Job submission complete."

    gemma3-27b-gsm8k.sh 파일에서 다음 값을 바꿉니다.

    • YOUR_WANDB_API_KEY: WandB API 키
    • YOUR_HF_TOKEN: Hugging Face 토큰입니다.

    이 파일에서는 GSM8k 데이터 세트에서 gemma3-27b-it 모델로 작업을 실행하는 구성을 확인할 수 있습니다. GRPO 학습 파이프라인을 완료하기 위해 이 스크립트는 다음 매개변수를 정의합니다.

    • num_prompts_per_step: 16 및 num_generations_per_prompt: 64: Gemma3-27b-it 모델은 프롬프트마다 많은 응답을 생성합니다. 이 구성에서 모델은 총 1,024개의 응답을 생성합니다 (16 × 64 = 1,024).
    • policy.generation.colocated.enabled=False: 이 파라미터는 공동 배치 생성 기능을 사용 중지합니다. 즉, 모델이 학습 프로세스와 동일한 노드에서 대답을 생성하지 않습니다. 표준 RL에서는 동일한 GPU가 학습과 생성을 모두 처리합니다. 이 NeMo RL 설정에서는 policy.generation.colocated.resources 매개변수로 관리되는 특정 노드를 vLLM 추론 전용으로 사용하고 나머지 클러스터는 대규모 학습 수학에 집중합니다. 이러한 워크로드를 분리하면 메모리 집약적인 학습 버퍼와 컴퓨팅 집약적인 추론 워크로드 간의 리소스 경합을 방지할 수 있습니다.

  3. 작업을 제출하려면 다음 명령어를 실행합니다.

    bash gemma3-27b-it/gemma3-27b-gsm8k.sh

    작업이 실행되면 출력에 학습 결과, 타이밍, 성능 측정항목이 표시됩니다.

GRPO 작업 상태 모니터링

Ray가 작업을 완료하면 NeMo RL이 구성된 경로에 체크포인트를 저장합니다.

  1. apt 트리 유틸리티를 설치합니다.

    apt install tree

  2. GRPO 작업의 상태를 모니터링하려면 Ray 헤드 노드의 로그를 확인하세요.

    kubectl exec -it $(kubectl get pods -l ray.io/node-type=head -o name) -c ray-head -- bash

    출력은 다음과 비슷합니다.

    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_LOCATION}
gcloud storage rm -r gs://${GS_BUCKET}

다음 단계