Google uses AI technology to translate content into your preferred language. AI translations can contain errors.

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

Standard

Neste tutorial, mostramos como orquestrar um ambiente de treinamento distribuído para aprendizado por reforço (RL) 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 de grupo (GRPO) no GKE com Ray e NeMo RL. A GRPO é um algoritmo de aprendizado por reforço projetado para melhorar a capacidade de raciocínio de um modelo. Esse algoritmo eficiente em termos de memória simplifica o processo de RL eliminando o Critic, ou modelo de valor, e usando 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 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)

O RL ensina modelos por experiência, exploração e feedback, em vez de imitação estática. Embora o pré-treinamento ensine um modelo sobre o que dizer, 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 é aprendizado por reforço?

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

O GRPO, um algoritmo popularizado pela DeepSeek, oferece uma alternativa com eficiência de memória à otimização de política proximal (PPO, na sigla em inglês) para o alinhamento de LLMs, removendo o modelo de crítica. Em vez de uma rede de crítica, 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 pós-treinamento de código aberto da NVIDIA projetada para RL escalonável. Parte do ecossistema mais amplo do framework 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 seguindo estas 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 do 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, 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 CLI gcloud com sua identidade federada.

Para inicializar a CLI gcloud, 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, 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 CLI gcloud com sua identidade federada.

Para inicializar a CLI gcloud, 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 uma vez o seguinte comando 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_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
```
Substitua os seguintes valores:
- CLUSTER_NAME: o nome do cluster do GKE.
- CONTROL_PLANE_REGION: a região do Compute Engine para o plano de controle do cluster do GKE.
- NODE_ZONE: a zona dos seus 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.
- RESERVATION: o nome da sua reserva de GPU.
- LUSTRE_NAME: o nome da instância do Lustre.
- YOUR_HF_TOKEN: seu token do Hugging Face.
- YOUR_WANDB_API_KEY: sua chave de API do Wandb.
Crie as seguintes variáveis de ambiente para a rede:
```
export NETWORK="NETWORK-NAME"
export GVNIC_NETWORK_PREFIX="GVNIC-NAME"
export RDMA_NETWORK_PREFIX="RDMA-NAME"
```
Substitua os seguintes valores:
- NETWORK-NAME: o nome da rede para o GKE.
- 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 ${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

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=${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

Criar o cluster do GKE

É possível definir o NeMo RL em um cluster GKE Standard.

Criar um cluster padrão

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}

Receba as credenciais do cluster:

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

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

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:

envsubst < network-mapping.yaml > network-mapping-updated.yaml
kubectl apply -f network-mapping-updated.yaml

Preparar o armazenamento

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

Alocar um intervalo de endereços IP para acesso a serviços particulares:

gcloud compute addresses create ${LUSTRE_NAME}-range \
    --global --purpose=VPC_PEERING \
    --prefix-length=20 --network=${NETWORK}

Conecte o peering:

gcloud services vpc-peerings connect \
    --service=servicenetworking.googleapis.com \
    --ranges=${LUSTRE_NAME}-range \
    --network=${NETWORK}

Crie uma instância do 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

Acessar uma instância do Managed Lustre usando o driver CSI do Managed Lustre

Extraia o endereço IP da instância do Managed Lustre.

export LUSTRE_IP=$(gcloud lustre instances describe ${LUSTRE_NAME} \
    --location=$NODE_ZONE --format="value(mountPoint)" | awk -F'@' '{print $1}')

Inspecione o manifesto lustre-pv.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.

apiVersion: v1
kind: PersistentVolume
metadata:
  name: lustre-pv
spec:
  storageClassName: lustre-rwx-500mbps-per-tib
  capacity:
    storage: 18000Gi
  accessModes:
    - ReadWriteMany
  persistentVolumeReclaimPolicy: Retain
  volumeMode: Filesystem
  claimRef:
    namespace: default
    name: lustre-pvc
  csi:
    driver: lustre.csi.storage.gke.io
    volumeHandle: "${PROJECT_ID}/${NODE_ZONE}/${LUSTRE_NAME}"
    volumeAttributes:
      ip: ${LUSTRE_IP}
      filesystem: lustrefs

Aplique o manifesto:

envsubst < lustre-pv.yaml > lustre-pv-updated.yaml
kubectl apply -f lustre-pv-updated.yaml

Inspecione o manifesto lustre-pvc.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.

apiVersion: v1
kind: PersistentVolumeClaim
metadata:
  name: lustre-pvc
spec:
  accessModes:
    - ReadWriteMany
  storageClassName: lustre-rwx-500mbps-per-tib
  volumeName: lustre-pv
  resources:
    requests:
      storage: 18000Gi

Aplique o manifesto:
```
kubectl apply -f lustre-pvc.yaml
```

Implantar o RayCluster

Nesta seção, você vai clonar o repositório de exemplo, preparar os manifestos e implantar o cluster do Ray:

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"
  nodeSelector:
    cloud.google.com/gke-nodepool: default-pool
  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.

Implante o cluster do Ray:
```
export REPLICA_COUNT=2
helm install ray-cluster . \
  --set additionalWorkerGroups.worker-grp-0.replicas=$REPLICA_COUNT
```
Neste tutorial, você vai usar dois nós de trabalho. Se quiser mudar o número de nós de worker, altere o valor de REPLICA_COUNT.

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 cluster

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 2 nodes (16 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=10 \
  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=32 \
  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: 32: o modelo Gemma3-27b-it gera um grande grupo de respostas para cada comando. Nessa configuração, o modelo produz 512 respostas no total (16 × 32 = 512).
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. No RL padrão, as mesmas GPUs processam o treinamento e a geração. Nesta configuração do NeMo RL, você dedica nós específicos (gerenciados com o parâmetro policy.generation.colocated.resources) apenas à inferência do 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 com uso intensivo de memória e as cargas de trabalho de inferência com uso intensivo de 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.

Para verificar a saída do job do GRPO, crie uma sessão SSH para o contêiner ray-head:

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

Instale o utilitário de árvore apt no terminal do contêiner ray-head:

apt update && apt install -y tree

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_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

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 a grupos (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

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