Ajustar e escalonar a aprendizagem por reforço com o NVIDIA NeMo RL no GKE

Autopilot Standard

Neste tutorial, mostramos como orquestrar um ambiente de treinamento distribuído para aprendizado por reforço (RL, na sigla em inglês) no Google Kubernetes Engine (GKE). Você usa o Ray e o framework NVIDIA NeMo RL para configurar um ambiente de treinamento distribuído e ajustar um modelo.

Este tutorial se concentra no pipeline de treinamento da otimização de política relativa a grupos (GRPO) no GKE com Ray e NeMo RL. O GRPO é um algoritmo de aprendizado por reforço projetado para melhorar a capacidade de raciocínio de um modelo. Esse algoritmo com eficiência de memória simplifica o processo de RL ao eliminar o Critic, ou modelo de valor, e usar um cálculo relativo baseado em grupo.

Antes de executar este tutorial, recomendamos que você conclua o tutorial Ajustar e escalonar o aprendizado por reforço com o verl no GKE. O tutorial a seguir usa a mesma configuração e configuração de cluster que o tutorial de ajuste e escalonamento de RL com verl.

Contexto

As seções a seguir oferecem uma breve visão geral dos conceitos usados neste tutorial.

Aprendizado por reforço (RL)

A RL ensina modelos por experiência, exploração e feedback, em vez de imitação estática. Embora o pré-treinamento ensine um modelo a falar, o aprendizado por reforço com feedback humano (RLHF) ensina como ser útil, seguro e lógico. O RL serve como ponte entre um modelo de base e um modelo refinado para um caso de uso especializado.

Para mais informações, consulte O que é aprendizagem por reforço?

Otimização de política relativa de grupo (GRPO)

O GRPO, um algoritmo popularizado pela DeepSeek, oferece uma alternativa com uso eficiente de memória à otimização de política proximal (PPO, na sigla em inglês) para o alinhamento de LLMs ao remover o modelo de crítica. Em vez de uma rede de críticos, o GRPO gera um grupo de respostas para o mesmo comando e usa a recompensa média desse grupo como base.

Para mais informações, consulte GRPO.

NVIDIA NeMo RL

O NeMo RL é uma biblioteca de código aberto da NVIDIA para pós-treinamento criada para RL escalonável. Parte do ecossistema mais amplo da estrutura NeMo, o NeMo RL permite experimentos em pequena escala em uma única GPU e implantações de vários nós em milhares de GPUs.

Para mais informações, consulte NVIDIA NeMo RL.

Conjunto de dados GSM8k

Neste tutorial, você vai usar o conjunto de dados GSM8k, que contém 8.500 problemas de matemática de alta qualidade e linguisticamente diversos para o ensino fundamental.

Ao usar o GSM8k e o GRPO, o modelo gera um grupo de n respostas diferentes para o mesmo problema. O GRPO compara essas respostas com a média do grupo. O modelo é mais recompensado por caminhos que são consistentemente corretos e logicamente válidos em comparação com o restante do grupo. Com o tempo, o modelo aprende que articular as etapas com clareza é a maneira mais confiável de maximizar a recompensa, reduzindo efetivamente a recompensa por respostas de baixa performance.

Para mais informações, consulte GSM8k.

Objetivos

Neste tutorial, mostramos como configurar o RL no GKE com o NeMo RL concluindo as seguintes etapas:

Prepare seu ambiente.
Configure um cluster do GKE com GPUs B200 ou H200.
Configure o KubeRay para gerenciar um cluster distribuído do Ray.
Use o Managed Lustre para armazenamento de alta performance.
Execute um job de treinamento de GRPO que usa o NeMo RL.

Antes de começar

Faça login na sua conta do Google Cloud . Se você começou a usar o Google Cloud, crie uma conta para avaliar o desempenho de nossos produtos em situações reais. Clientes novos também recebem US$ 300 em créditos para executar, testar e implantar cargas de trabalho.

Instale a CLI do Google Cloud.

Observação:se você já instalou a CLI gcloud anteriormente, execute o comando gcloud components update para verificar se tem a versão mais recente.

Ao usar um provedor de identidade (IdP) externo, primeiro faça login na gcloud CLI com sua identidade federada.

Para inicializar a gcloud CLI, execute o seguinte comando:

gcloud init

Crie ou selecione um Google Cloud projeto.

Funções necessárias para selecionar ou criar um projeto

Selecionar um projeto: não é necessário um papel específico do IAM para selecionar um projeto. Você pode escolher qualquer projeto em que tenha recebido um papel.
Criar um projeto: para criar um projeto, é necessário ter o papel de Criador de projetos (roles/resourcemanager.projectCreator), que contém a permissão resourcemanager.projects.create. Saiba como conceder papéis.

Crie um projeto do Google Cloud :
```
gcloud projects create PROJECT_ID
```
Substitua PROJECT_ID por um nome para o projeto Google Cloud que você está criando.
Selecione o projeto Google Cloud que você criou:
```
gcloud config set project PROJECT_ID
```
Substitua PROJECT_ID pelo nome do projeto do Google Cloud .

Verifique se o faturamento está ativado para o projeto do Google Cloud .

Ative as APIs necessárias:

Funções necessárias para ativar APIs

Para ativar as APIs, é necessário ter o papel do IAM de administrador de uso do serviço (roles/serviceusage.serviceUsageAdmin), que contém a permissão serviceusage.services.enable. Saiba como conceder papéis.

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

Instale a CLI do Google Cloud.

Observação:se você já instalou a CLI gcloud anteriormente, execute o comando gcloud components update para verificar se tem a versão mais recente.

Ao usar um provedor de identidade (IdP) externo, primeiro faça login na gcloud CLI com sua identidade federada.

Para inicializar a gcloud CLI, execute o seguinte comando:

gcloud init

Crie ou selecione um Google Cloud projeto.

Funções necessárias para selecionar ou criar um projeto

Selecionar um projeto: não é necessário um papel específico do IAM para selecionar um projeto. Você pode escolher qualquer projeto em que tenha recebido um papel.
Criar um projeto: para criar um projeto, é necessário ter o papel de Criador de projetos (roles/resourcemanager.projectCreator), que contém a permissão resourcemanager.projects.create. Saiba como conceder papéis.

Crie um projeto do Google Cloud :
```
gcloud projects create PROJECT_ID
```
Substitua PROJECT_ID por um nome para o projeto Google Cloud que você está criando.
Selecione o projeto Google Cloud que você criou:
```
gcloud config set project PROJECT_ID
```
Substitua PROJECT_ID pelo nome do projeto do Google Cloud .

Verifique se o faturamento está ativado para o projeto do Google Cloud .

Ative as APIs necessárias:

Funções necessárias para ativar APIs

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

Atribua papéis à sua conta de usuário. Execute o seguinte comando uma vez para cada um dos seguintes papéis do IAM: roles/container.admin, roles/iam.serviceAccountAdmin, roles/storage.admin
```
gcloud projects add-iam-policy-binding PROJECT_ID --member="user:USER_IDENTIFIER" --role=ROLE
```
Substitua:
- PROJECT_ID: o ID do projeto.
- USER_IDENTIFIER: o identificador da sua conta de usuário . Por exemplo, myemail@example.com.
- ROLE: o papel do IAM concedido à sua conta de usuário.

Crie uma conta do Hugging Face caso ainda não tenha uma.
Verifique se você tem um token do Hugging Face.
Verifique se o projeto tem cota suficiente para GPUs B200 e H200. Para saber mais, consulte Planejar a cota de GPU e Cota de GPU.

Preparar o ambiente

Neste tutorial, você vai usar o Cloud Shell.

Acesse o console doGoogle Cloud .
Na parte de cima da janela do console do Google Cloud , clique no botão Ativar Cloud Shell.
Configure as variáveis de ambiente a seguir:
```
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
```
Substitua os seguintes valores:
- CLUSTER_NAME: o nome do cluster do GKE.
- CONTROL_PLANE_LOCATION: a região do Compute Engine para o plano de controle do cluster do GKE.
- NODE_LOCATION: o local dos nós. Selecione uma zona em que as GPUs NVIDIA B200 ou H200 estejam disponíveis.
- GPU_TYPE: o acelerador que você reservou na reserva de capacidade do Compute Engine. Precisa ser um dos valores abaixo:
  - nvidia-b200: NVIDIA B200 (180 GB)
  - nvidia-h200-141gb: NVIDIA H200 (141 GB)
- MACHINE_TYPE: o tipo de máquina a ser usado:
  - Para GPUs NVIDIA B200 (180 GB), use a4-highgpu-8g ou mais recente.
  - Para GPUs NVIDIA H200 (141 GB), use a3-ultragpu-8g ou mais recente.
- GKE_VERSION: a versão do GKE a ser usada:
  - Para GPUs NVIDIA B200 (180 GB), use 1.32.2-gke.1422000 ou mais recente.
  - Para GPUs NVIDIA H200 (141 GB), use 1.31.4-gke.1183000 ou mais recente.
- BUCKET_NAME: o nome base do bucket do Cloud Storage.
- YOUR_HUGGING_FACE_TOKEN: seu token do Hugging Face.
Crie as seguintes variáveis de ambiente para a rede:
```
export GVNIC_NETWORK_PREFIX="GVNIC-NAME"
export RDMA_NETWORK_PREFIX="RDMA-NAME"
```
Substitua os seguintes valores:
- GVNIC-NAME: o prefixo do nome da rede gVNIC. Você pode usar qualquer prefixo.
- RDMA-NAME: o prefixo da rede de acesso direto à memória (RDMA) remota. Você pode usar qualquer prefixo.

Configurar a infraestrutura

Nesta seção, você cria redes VPC e um cluster do GKE.

Crie uma rede VPC

Crie uma rede VPC para a interface gVNIC:

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

Crie uma rede e sub-redes VPC para RDMA que incluam oito sub-redes para oito GPUs:

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

Criar o cluster do GKE

É possível definir o NeMo RL em um cluster do GKE Autopilot ou Standard. Recomendamos que você use um cluster do Autopilot para ter uma experiência totalmente gerenciada do Kubernetes. Para escolher o modo de operação do GKE mais adequado para suas cargas de trabalho, consulte Sobre os modos de operação do GKE.

Piloto automático

Criar um cluster do Autopilot:

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

Receba as credenciais do cluster:

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

Instale o instalador do NCCL RDMA para o Autopilot:

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

Padrão

Criar um cluster padrão

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

Receba as credenciais do cluster:

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

Crie o pool de nós de 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

Instale o instalador do NCCL RDMA:

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

Configurar mapeamentos de rede

Salve o seguinte manifesto como 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

Aplique o manifesto:
```
kubectl apply -f network-mapping.yaml
```

Preparar o armazenamento

Nesta seção, você vai criar buckets do Cloud Storage e uma instância do Managed Lustre, que provisiona o armazenamento de alta performance necessário para sua carga de trabalho de RL.

Crie um bucket do Cloud Storage:

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

Crie uma conta de serviço do Kubernetes (KSA) e vincule-a ao bucket:

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"

Para configurar o Managed Lustre, siga estas etapas:
1. Crie uma instância do Managed Lustre seguindo as etapas em Criar uma instância do Managed Lustre. Verifique se a instância usa a mesma rede do cluster do GKE.
2. Siga as etapas em Acessar uma instância do Managed Lustre.

Implantar o RayCluster

Nesta seção, você vai clonar o repositório de amostra, preparar os manifestos e executar um script launcher.sh:

Clone o repositório de amostra:

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

Navegue até o diretório de trabalho:
```
cd ai-ml/nemo-rl-on-gke/nemoRL
```

Inspecione o manifesto 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

Substitua NCCL_TUNER_CONFIG_PATH por qualquer um dos valores a seguir, com base no acelerador usado neste tutorial:

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

Neste manifesto, o nó principal gerencia o job e hospeda o painel do Ray. Os nós de trabalho executam os jobs de treinamento.

Instale o cluster do Ray:
```
export REPLICA_COUNT=2
helm install ray-cluster . \
  --set values.additionalWorkerGroups.worker-grp-0.replicas=$REPLICA_COUNT
```
Neste tutorial, você usa dois nós de trabalho. Se quiser mudar o número de nós de worker, altere o valor de REPLICA_COUNT.
Para implantar o cluster do Ray, execute o script launcher.sh:
```
bash launcher.sh
```

Verifique se os nós de trabalho e de cabeçalho estão em execução:

kubectl get pods

O resultado será o seguinte:

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

Verifique se o cluster do Ray está em execução:

kubectl ray get clusters

O resultado será o seguinte:

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

Iniciar o job do GRPO

Depois que o cluster do Ray estiver pronto, envie um job do Ray para o cluster do Ray em execução no GKE. O NeMo RL baixa automaticamente o modelo durante a execução do job de treinamento de RL.

Para enviar um job do Ray, inicie uma sessão interativa para executar o job.

Para estabelecer uma conexão local com o cluster do Ray, execute este comando:
```
  kubectl ray session ray-cluster-kuberay
```
Esse comando inicia o encaminhamento de porta entre a máquina local e o nó principal do Ray no cluster do GKE. Observe que o terminal vai ficar ocupado enquanto essa sessão estiver ativa. Para continuar, abra uma instância de terminal separada.

Edite o arquivo 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."

Substitua os seguintes valores no arquivo gemma3-27b-gsm8k.sh:

YOUR_WANDB_API_KEY: sua chave de API do WandB.
YOUR_HF_TOKEN: seu token do Hugging Face.

Nesse arquivo, você pode conferir a configuração para executar um job com o modelo gemma3-27b-it no conjunto de dados GSM8k. Para concluir o pipeline de treinamento do GRPO, esse script define os seguintes parâmetros:

num_prompts_per_step: 16 e num_generations_per_prompt: 64: o modelo Gemma3-27b-it gera um grande grupo de respostas para cada solicitação. Nessa configuração, o modelo produz 1.024 respostas no total (16 × 64 = 1.024).
policy.generation.colocated.enabled=False: esse parâmetro desativa o recurso de geração colocada, o que significa que o modelo não gera respostas no mesmo nó que o processo de treinamento. Na RL padrão, as mesmas GPUs processam o treinamento e a geração. Nessa configuração do NeMo RL, você dedica nós específicos (gerenciados com o parâmetro policy.generation.colocated.resources) somente à inferência vLLM, enquanto o restante do cluster se concentra na matemática de treinamento pesado. Ao separar essas cargas de trabalho, você evita a disputa de recursos entre os buffers de treinamento que exigem muita memória e as cargas de trabalho de inferência que exigem muita computação.

Para enviar o job, execute o seguinte comando:
```
bash gemma3-27b-it/gemma3-27b-gsm8k.sh
```
Quando o job está em execução, a saída mostra os resultados do treinamento, o tempo e as métricas de performance.

Monitorar a integridade do job de GRPO

Depois que o Ray termina o job, o NeMo RL armazena os pontos de verificação no caminho configurado.

Instale o utilitário de árvore apt:
```
apt install tree
```

Para monitorar a integridade do job GRPO, verifique os registros do nó principal do Ray:

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

O resultado será o seguinte:

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

Limpar

Para evitar cobranças, exclua os recursos:

helm delete ray-cluster
gcloud container clusters delete ${CLUSTER_NAME} --location=${CONTROL_PLANE_LOCATION}
gcloud storage rm -r gs://${GS_BUCKET}

Ajustar e escalonar a aprendizagem por reforço com o NVIDIA NeMo RL no GKE Mantenha tudo organizado com as coleções Salve e categorize o conteúdo com base nas suas preferências.

Contexto

Aprendizado por reforço (RL)

Otimização de política relativa de grupo (GRPO)

NVIDIA NeMo RL

Conjunto de dados GSM8k

Objetivos

Antes de começar

Preparar o ambiente

Configurar a infraestrutura

Crie uma rede VPC

Criar o cluster do GKE

Piloto automático

Padrão

Configurar mapeamentos de rede

Preparar o armazenamento

Implantar o RayCluster

Iniciar o job do GRPO

Monitorar a integridade do job de GRPO

Limpar

A seguir

Ajustar e escalonar a aprendizagem por reforço com o NVIDIA NeMo RL no GKE