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, 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, conclua o tutorial Ajustar e escalonar o aprendizado por reforço com o verl no GKE. Este tutorial usa a mesma configuração e configuração de cluster do 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 é 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) 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 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:

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

Observação: este tutorial usa a mesma configuração e configuração de cluster do tutorial Ajustar e escalonar o aprendizado por reforço com verl no GKE. Se você já tiver concluído esse tutorial, pule para as seções Preparar seu ambiente e Configurar a infraestrutura. Na infraestrutura já configurada, acesse a seção Implantar RayCluster para implantar um cluster do Ray.

Antes de começar

1. 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.

2. Configure a CLI gcloud para usar sua identidade federada.

   Para mais informações, consulte Fazer login na CLI gcloud com sua identidade federada.

3. Para inicializar a CLI gcloud, execute o seguinte comando:

   gcloud init

4. 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.
   Observação: se você não pretende manter os recursos criados neste procedimento, crie um projeto novo em vez de selecionar um que já existe. Depois de concluir essas etapas, é possível excluir o projeto. Para fazer isso, basta remover todos os recursos associados a ele.

   • 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 .

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

6. Ative as APIs necessárias:

   Funções necessárias para ativar APIs

   Para ativar APIs, você precisa da permissão serviceusage.services.enable. Se você criou o projeto, provavelmente já tem essa permissão com o papel de Proprietário (roles/owner). Caso contrário, é possível receber essa permissão com o papel de Administrador do Service Usage (roles/serviceusage.serviceUsageAdmin). Saiba como conceder papéis.

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

7. 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 conta. Para exemplos, consulte Representar usuários do pool de força de trabalho nas políticas do IAM.
   • ROLE: o papel do IAM concedido à sua conta de usuário.
8. Crie uma conta do Hugging Face caso ainda não tenha uma.
9. Verifique se você tem um token do Hugging Face com read access.
10. Crie uma conta do Weights & Biases (Wandb), caso ainda não tenha uma.
11. Crie uma chave de API do Wandb.
12. Verifique se o projeto Google Cloud 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.

1. Acesse o console doGoogle Cloud .

2. Na parte de cima da janela do console do Google Cloud , clique no botão Ativar Cloud Shell.

3. Configure as variáveis de ambiente a seguir:

   export CONTROL_PLANE_REGION="YOUR_REGION"
   export NODE_ZONE="YOUR_ZONE"
   export CLUSTER_NAME="YOUR_CLUSTER_NAME"
   export GPU_TYPE="YOUR_GPU_TYPE"
   export MACHINE_TYPE="YOUR_MACHINE_TYPE"
   export KSA_NAME="generic-ksa"
   export NAMESPACE="default"
   export RESERVATION="RESERVATION_NAME"
   export LUSTRE_NAME="CHOSEN_LUSTRE_NAME"
   export HF_TOKEN="YOUR_HF_TOKEN"
   export WANDB_API_KEY="YOUR_WANDB_API_KEY"
   
   export PROJECT_ID=$(gcloud config get project)
   export PROJECT_NUMBER=$(gcloud projects describe "${PROJECT_ID}" --format="value(projectNumber)")

   Substitua os seguintes valores:

   • YOUR_REGION: a região do Compute Engine para o plano de controle do cluster do GKE.
   • YOUR_NODE_ZONE: a zona dos nós. Selecione uma zona em que as GPUs NVIDIA B200 ou H200 estejam disponíveis.
   • YOUR_CLUSTER_NAME: o nome do cluster do GKE.
   • YOUR_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)
   • YOUR_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.
   • YOUR_RESERVATION_NAME: o nome da sua reserva de GPU.
   • CHOSEN_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.

4. Crie as seguintes variáveis de ambiente para a rede:

   export NETWORK="YOUR_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

1. 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

2. 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.

1. 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}

2. Receba as credenciais do cluster:

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

3. 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

4. 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

1. 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

2. Aplique o manifesto:

   envsubst < network-mapping.yaml | kubectl apply -f -

Preparar o armazenamento

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

1. 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}

2. Conecte o peering:

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

3. 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

4. Acesse uma instância do Managed Lustre usando o driver CSI do Managed Lustre.

   1. 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}')

   2. 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

   3. Aplique o manifesto:

      envsubst < lustre-pv.yaml | kubectl apply -f -

   4. 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

   5. 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:

1. Clone o repositório de amostra:

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

2. Navegue até o diretório de trabalho:

   cd kubernetes-engine-samples/ai-ml/nemo-rl-on-gke/nemoRL

3. 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.

4. 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.

5. 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                2/2   Running 0      33h
   ray-cluster-kuberay-worker-grp-0-worker-gkbxw 2/2   Running 0      33h
   ray-cluster-kuberay-worker-grp-0-worker-kdg62 2/2   Running 0      33h
   

6. 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.

1. 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.

2. Em um terminal separado, navegue até kubernetes-engine-samples/ai-ml/nemo-rl-on-gke/nemoRL/gemma3-27b-it e 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. 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) apenas à 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.

3. 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.

Sucesso: o NeMo RL está orquestrando um ciclo de vida distribuído nas VMs de GPU no cluster.

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.

1. 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

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

   apt update && apt install -y tree

3. Liste a estrutura de diretórios do contêiner ray-head:

   tree /data/nemo_rl_gemma3_27b_3_17/

   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
   

Observação: os pontos de verificação gerados pelo NeMo RL estão no formato de ponto de verificação distribuído do PyTorch. Para avaliar ou disponibilizar o modelo, é necessário converter o checkpoint para o formato de peso do Hugging Face. Para instruções sobre como converter o checkpoint, consulte Converter o formato PyTorch DCP para o formato HF.

Limpar

Para evitar cobranças na sua conta do Google Cloud pelos recursos usados neste tutorial, exclua os recursos individuais ou o projeto que os contém.

Excluir os recursos

1. Exclua o cluster do Slurm:

   helm delete ray-cluster

2. Exclua seu Cluster do GKE:

   gcloud container clusters delete ${CLUSTER_NAME} \
       --location=${CONTROL_PLANE_REGION} \
       --quiet

3. Exclua o sistema de arquivos Lustre:

   gcloud lustre instances delete ${LUSTRE_NAME} --location=${NODE_ZONE} --quiet

4. Exclua o peering de VPC:

   gcloud services vpc-peerings delete \
       --service=servicenetworking.googleapis.com \
       --network=${NETWORK}

5. Exclua o intervalo de endereços IP privados do Lustre:

   gcloud compute addresses delete ${LUSTRE_NAME}-range --global --quiet

6. Exclua as sub-redes RDMA e gVNIC:

   gcloud compute networks subnets delete ${GVNIC_NETWORK_PREFIX}-sub \
       --region=${CONTROL_PLANE_REGION} --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

7. Exclua as regras de firewall e as redes:

   echo "[$(date)] ========== Deleting firewall rules and networks =========="
   
   NETWORKS=(
       "${RDMA_NETWORK_PREFIX}-net"
       "${GVNIC_NETWORK_PREFIX}-net"
       "${NETWORK}"
   )
   
   for NW in "${NETWORKS[@]}"; do
   
     echo "========== Deleting firewall rules for ${NW} =========="
     while true; do
         rules=$(gcloud compute firewall-rules list \
           --filter="network:${NW}" \
           --format="value(name)" \
           --project="${PROJECT_ID}")
   
           if [[ -z "${rules}" ]]; then
             echo "No firewall rules remain for ${NW}"
             break
           fi
   
           for rule in ${rules}; do
             echo "Deleting firewall rule ${rule}..."
             gcloud compute firewall-rules delete "${rule}" --project="${PROJECT_ID}" --quiet || true
           done
   
           sleep 3
         done
   
         echo "[$(date)] ========== Deleting network ${NW} =========="
         gcloud compute networks delete ${NW} --quiet || true
   done

Excluir o projeto

Cuidado: a exclusão de um projeto tem os seguintes efeitos:

• Tudo no projeto é excluído. Se você tiver usado um projeto atual para as tarefas neste documento, a exclusão dele incluirá a exclusão de quaisquer outros trabalhos feitos nele.
• Os IDs do projeto personalizados são perdidos. Ao criar o projeto, você pode ter criado um ID do projeto personalizado para ser usado no futuro. Para preservar os URLs que usam o ID do projeto, como um URL appspot.com, exclua recursos específicos do projeto, em vez de excluir o projeto inteiro.

Se você planeja explorar várias arquiteturas, tutoriais e guias de início rápido, a reutilização de projetos pode evitar que você exceda limites da cota.

Excluir um projeto do Google Cloud :

gcloud projects delete PROJECT_ID

A seguir

• Saiba mais sobre o Ray no GKE
• Confira a documentação do NVIDIA NeMo RL