Escalonar o aprendizado por reforço com o GKE e o Managed Lustre

1. Introdução

Se você preferir executar os scripts empacotados diretamente sem o tutorial detalhado, eles estão disponíveis no repositório GoogleCloudPlatform/devrel-demos.

Neste codelab, você vai aprender a implantar um pipeline de treinamento de alta performance para aprendizado por reforço (RL, na sigla em inglês) usando o Google Kubernetes Engine (GKE) e o Managed Lustre.

As cargas de trabalho de aprendizagem por reforço, principalmente aquelas que usam algoritmos como a otimização de política relativa de grupo (GRPO, na sigla em inglês), geram grandes quantidades de dados durante a "geração de experiência" e exigem checkpoints frequentes. O armazenamento de objetos padrão pode causar gargalos durante esses picos de E/S, deixando aceleradores caros ociosos.

Você vai usar o Managed Lustre, um sistema de arquivos paralelo, para eliminar esses gargalos e alcançar uma capacidade de processamento de treinamento maior.

Atividades deste laboratório

• Configure variáveis de ambiente para um cluster do Ray baseado em GPU.
• Provisione um cluster de GPU do Spot no GKE e uma instância Managed Lustre usando o Cluster Toolkit.
• Implante um cluster KubeRay e monte o sistema de arquivos Lustre.
• Envie uma carga de trabalho de treinamento do NeMo-RL.
• Observe alta capacidade de processamento e baixa latência de checkpoint usando o Cloud Monitoring.

O que é necessário

• Um navegador da Web, como o Chrome.
• Ter um projeto do Google Cloud com o faturamento ativado.

Este codelab usa recursos sofisticados, incluindo GPUs Spot e Managed Lustre. Embora usemos a capacidade spot para reduzir custos, esses recursos podem ser caros se ficarem em execução. Sempre siga a etapa de limpeza no final do laboratório para evitar cobranças inesperadas.

Este codelab é destinado a usuários técnicos avançados, engenheiros de plataforma e pesquisadores de IA que conhecem o GKE e os conceitos de armazenamento.

Duração total estimada: 45 a 60 minutos mais 2 horas de treinamento

2. Antes de começar

Criar um projeto do Google Cloud

1. No Console do Google Cloud, selecione ou crie um projeto na nuvem do Google Cloud.
2. Verifique se o faturamento está ativado para seu projeto do Cloud.

Iniciar o Cloud Shell

O Cloud Shell é um ambiente de linha de comando executado no Google Cloud que vem pré-carregado com as ferramentas necessárias.

1. Clique em Ativar o Cloud Shell na parte de cima do console do Google Cloud.
2. Depois de se conectar ao Cloud Shell, verifique sua autenticação:

   gcloud auth list
   

3. Confirme se o projeto está configurado:

   gcloud config get project
   

4. Se o projeto não estiver definido como esperado, faça o seguinte:

   export PROJECT_ID=<YOUR_PROJECT_ID>
   gcloud config set project $PROJECT_ID
   

Instalar o Cluster Toolkit

Este codelab usa o Cluster Toolkit (gcluster) para implantar o cluster do GKE. Para instruções sobre como configurar o Cluster Toolkit, consulte o guia de configuração do Cluster Toolkit.

Ativar APIs

Execute este comando no Cloud Shell para ativar todas as APIs necessárias:

gcloud services enable \
  container.googleapis.com \
  lustre.googleapis.com \
  compute.googleapis.com \
  servicenetworking.googleapis.com

3. Configurar variáveis de ambiente

Para manter a consistência dos comandos neste codelab, configure algumas variáveis de ambiente.

Crie um arquivo chamado env.sh e o preencha com sua configuração. Você pode usar o seguinte modelo:

# Environment Variables for the RL Demo execution
export PROJECT_ID="{{'<var>'}}PROJECT_ID{{'</var>'}}"
export ZONE="us-east1-b"
export REGION="us-east1"
export CLUSTER_NAME="ray-a4-gpu-spot"
export HF_TOKEN="{{'<var>'}}YOUR_HF_TOKEN{{'</var>'}}" # Required for downloading models
export WANDB_API_KEY="{{'<var>'}}YOUR_WANDB_API_KEY{{'</var>'}}" # Optional

# Topology defaults
export NUM_NODES="8"
export GPUS_PER_NODE="8" # Fixed for A4/B200 architecture

HF_TOKEN é seu token de acesso de usuário do Hugging Face, que é necessário para baixar modelos. Para conseguir um, siga as instruções em Tokens de segurança do Hugging Face.

Substitua <YOUR_PROJECT_ID> e <YOUR_HF_TOKEN> pelos seus valores reais.

WANDB_API_KEY é opcional. Ele é usado para Weights & Biases (em inglês), uma ferramenta para monitorar e visualizar experimentos de aprendizado de máquina. Se você não tiver uma conta ou não quiser usá-la, deixe o campo em branco ou no estado em que se encontra.

Crie o arquivo para carregar as variáveis na sessão atual:

source env.sh

4. Implantar o cluster do GKE e o Managed Lustre usando o Cluster Toolkit

Nesta etapa, você vai usar o Cluster Toolkit (gcluster) para implantar um cluster do GKE com GPUs Spot e provisionar automaticamente o armazenamento Managed Lustre com o driver CSI do Lustre e o PersistentVolumeClaim pré-configurado (lustre-pvc).

Preparar o blueprint

Antes de fazer a implantação, revise o blueprint do examples/gke-a4/gke-a4.yaml. Para mais detalhes, consulte Criar um cluster A4:

1. Ativar o Managed Lustre: remova o comentário das seções de módulo managed-lustre e lustre-pvc em gke-a4.yaml.
2. Ative o complemento RayOperator: defina enable_ray_operator: true nas configurações do módulo gke_cluster em gke-a4.yaml.

Implantar infraestrutura

Defina o CIDR autorizado para acesso ao Cloud Shell e implante usando gcluster deploy:

export AUTHORIZED_CIDR="$(curl -s ifconfig.me)/32"

gcluster deploy examples/gke-a4/gke-a4.yaml \
  --vars project_id=${PROJECT_ID},deployment_name=${CLUSTER_NAME},region=${REGION},zone=${ZONE},static_node_count=${NUM_NODES},authorized_cidr=${AUTHORIZED_CIDR},spot=true

Definir spot=true em --vars provisiona dinamicamente o pool de nós da GPU usando a capacidade do Spot sem exigir edições manuais em reservation_affinity. A capacidade do Spot pode ser interrompida. Em ambientes de produção, omita spot=true para usar a capacidade sob demanda ou configure uma reserva.

Aguarde até que a implantação seja concluída. O Cluster Toolkit provisiona automaticamente a rede VPC, o peering de acesso a serviços particulares (PSA, na sigla em inglês), o sistema de arquivos Managed Lustre, o driver CSI do Lustre, o complemento RayOperator e a reivindicação de armazenamento do Kubernetes (lustre-pvc) em uma implantação coordenada.

5. Implantar um cluster do Ray no GKE

Nesta etapa, você vai implantar um cluster do KubeRay nos nós do GKE e montar o sistema de arquivos Lustre usando o PersistentVolumeClaim (lustre-pvc) provisionado automaticamente pelo Cluster Toolkit.

Criar configuração do RayCluster

Crie um arquivo chamado ray-cluster.yaml. Isso especifica os nós principais e de trabalho do KubeRay, usando o tipo de acelerador nvidia-b200 e montando o volume do Lustre em /lustre.

cat << EOF > ray-cluster.yaml
apiVersion: ray.io/v1
kind: RayCluster
metadata:
  name: ${CLUSTER_NAME}
  namespace: default
spec:
  rayVersion: '2.54.0'
  headGroupSpec:
    rayStartParams:
      dashboard-host: '0.0.0.0'
    template:
      spec:
        nodeSelector:
          cloud.google.com/gke-accelerator: nvidia-b200
        tolerations:
        - key: "nvidia.com/gpu"
          operator: "Exists"
          effect: "NoSchedule"
        containers:
        - name: ray-head
          image: nvcr.io/nvidia/nemo-rl:v0.4.0
          ports:
          - containerPort: 6379
            name: gcs-server
          - containerPort: 8265
            name: dashboard
          - containerPort: 10001
            name: client
          resources:
            limits:
              cpu: "32"
              memory: "1000Gi"
            requests:
              cpu: "8"
              memory: "64Gi"
          volumeMounts:
          - mountPath: /lustre
            name: lustre-storage
        volumes:
        - name: lustre-storage
          persistentVolumeClaim:
            claimName: lustre-pvc
  workerGroupSpecs:
  - groupName: gpu-worker-group
    replicas: ${NUM_NODES}
    minReplicas: ${NUM_NODES}
    maxReplicas: ${NUM_NODES}
    rayStartParams: {}
    template:
      spec:
        nodeSelector:
          cloud.google.com/gke-accelerator: nvidia-b200
        tolerations:
        - key: "nvidia.com/gpu"
          operator: "Exists"
          effect: "NoSchedule"
        containers:
        - name: ray-worker
          image: nvcr.io/nvidia/nemo-rl:v0.4.0
          resources:
            limits:
              nvidia.com/gpu: "8"
              cpu: "100"
              memory: "1000Gi"
            requests:
              nvidia.com/gpu: "8"
              cpu: "100"
              memory: "1000Gi"
          volumeMounts:
          - mountPath: /lustre
            name: lustre-storage
          - mountPath: /dev/shm
            name: dshm
        volumes:
        - name: lustre-storage
          persistentVolumeClaim:
            claimName: lustre-pvc
        - name: dshm
          emptyDir:
            medium: Memory
EOF

Conectar ao cluster

Verifique se a sessão do Cloud Shell está autenticada no cluster do GKE:

gcloud container clusters get-credentials ${CLUSTER_NAME} \
  --region ${REGION} \
  --project ${PROJECT_ID}

Aplicar configuração do RayCluster

Aplique a configuração do cluster do Ray:

kubectl apply -f ray-cluster.yaml

Verificar o status do cluster

Monitore a criação dos pods:

kubectl get pods -w

Aguarde até que os pods principal e de worker estejam Running.

6. Enviar carga de trabalho de aprendizado por reforço

Nesta etapa, você vai enviar o job de treinamento do GRPO NeMo-RL para o cluster do Ray.

Conectar-se ao painel do Ray

Para enviar jobs e ver métricas, é necessário se conectar ao painel do Ray. Como o painel está no GKE, use o encaminhamento de portas para acessá-lo no Cloud Shell:

# Run this in a separate Cloud Shell tab or in the background
kubectl port-forward service/${CLUSTER_NAME}-head-svc 8265:8265 &

Criar o script de execução

Crie um arquivo chamado run_nemo_rl.sh. Esse script será executado nos workers do cluster do Ray. Usamos cat << EOF para preencher as variáveis de ambiente definidas anteriormente.

cat << EOF > run_nemo_rl.sh
#!/bin/bash
set -ex

# Override job runtime conflicts (NeMo-RL passes os.environ to ray.init)
export RAY_OVERRIDE_JOB_RUNTIME_ENV=1

echo "--- Running on Ray Cluster ---"
cd /opt/nemo-rl

# Ensure directories exist on the high-speed Lustre drive
mkdir -p /lustre/huggingface_cache
mkdir -p /lustre/nemo_rl_qwen_72b_ds_cp

echo "Launching NeMo-RL GRPO training..."
uv run python examples/run_grpo_math.py \
  --config examples/configs/grpo_math_70B_megatron.yaml \
  policy.model_name='Qwen/Qwen2.5-72B-Instruct' \
  policy.megatron_cfg.converter_type='Qwen2ForCausalLM' \
  logger.wandb_enabled=False \
  cluster.num_nodes=${NUM_NODES} \
  cluster.gpus_per_node=${GPUS_PER_NODE} \
  logger.wandb.name='nemo-rl-grpo-test1' \
  grpo.max_num_steps=20 \
  grpo.num_generations_per_prompt=8 \
  grpo.num_prompts_per_step=32 \
  policy.train_global_batch_size=256 \
  checkpointing.enabled=True \
  checkpointing.save_period=2 \
  checkpointing.keep_top_k=2 \
  checkpointing.metric_name=null \
  checkpointing.checkpoint_dir=/lustre/nemo_rl_qwen_72b_ds_cp/nemo-rl-grpo-test1 \
  data.dataset_name='DeepScaler'
EOF
chmod +x run_nemo_rl.sh

Criar arquivo de ignorar do Ray

Crie um arquivo .rayignore para impedir que o Ray faça upload de diretórios grandes ou desnecessários:

cat << EOF > .rayignore
cluster-toolkit/
.git/
*.sh.log
EOF

Criar configuração do ambiente de execução

Crie um arquivo JSON para transmitir variáveis de ambiente ao job do Ray:

cat << EOF > ray_runtime_env_nemo.json
{
  "env_vars": {
    "HF_TOKEN": "${HF_TOKEN}",
    "WANDB_API_KEY": "${WANDB_API_KEY}",
    "HF_HOME": "/lustre/huggingface_cache",
    "GLOO_SOCKET_IFNAME": "eth0",
    "NCCL_SOCKET_IFNAME": "eth0"
  }
}
EOF

Enviar o job

Use a CLI do Ray para enviar o job ao endpoint do painel. Se o comando ray não for encontrado no Cloud Shell, instale-o com pip install ray:

ray job submit \
    --address="http://localhost:8265" \
    --working-dir . \
    --runtime-env ray_runtime_env_nemo.json \
    -- bash run_nemo_rl.sh

Os registros vão aparecer no terminal do Cloud Shell. O job vai carregar o modelo, inicializar os workers do Ray e iniciar o loop de treinamento do GRPO.

7. Monitorar a performance do treinamento

Nesta etapa, você vai observar o desempenho do sistema de arquivos Lustre durante o treinamento e o checkpointing.

Verificar os registros de treinamento

À medida que o treinamento avança, você vê registros indicando que os pontos de verificação estão sendo salvos em /lustre/nemo_rl_qwen_72b_ds_cp/nemo-rl-grpo-test1. Observe que a criação de pontos de verificação acontece de forma assíncrona e não bloqueia os trabalhadores do Ray por muito tempo.

Para conferir a velocidade do checkpointing, procure linhas de registro que indiquem pontos de verificação salvos.

Ver métricas do Lustre no console do Cloud

Para conferir as métricas da sua instância do Lustre:

1. No Console do Google Cloud, pesquise Serviço gerenciado para Lustre.
2. Clique no nome da instância (por exemplo, ${CLUSTER_NAME}-lustre ou rl-demo-gpu-lustre).
3. Clique na guia Monitoramento.

Aqui você pode observar:

• Capacidade (bytes/segundo): confira os picos durante a criação de pontos de verificação.
• Capacidade: monitore quanto espaço está sendo consumido pelos pontos de verificação.

O Lustre consegue gravar em alta velocidade, criando checkpoints em um tempo mínimo.

8. Limpar recursos

Essa etapa é essencial para evitar cobranças contínuas. Se deixados em execução, os clusters de GPU e o Managed Lustre podem gerar custos significativos.

Execute o seguinte comando no Cloud Shell para destruir toda a infraestrutura provisionada (cluster do GKE, pools de nós de GPU, instância do Managed Lustre e rede VPC) em uma única etapa:

gcluster destroy "${CLUSTER_NAME}"

Esse comando é executado de forma síncrona em primeiro plano, desativando toda a infraestrutura gerenciada pela implantação e gerando logs de progresso no terminal. Aguarde a conclusão total do comando antes de fechar a sessão do Cloud Shell.

9. Parabéns

Você concluiu o codelab Escalonar o aprendizado por reforço com o GKE e o Managed Lustre.

O que você aprendeu

• Como usar o Cluster Toolkit para provisionar um cluster de GPU do GKE com instâncias Spot e armazenamento Managed Lustre.
• Como implantar um cluster do KubeRay e fazer a montagem do armazenamento do Lustre.
• Como enviar uma carga de trabalho de treinamento do GRPO do NeMo-RL.
• Como observar a performance do armazenamento durante o treinamento.

Próximas etapas

• Conheça mais recursos do NVIDIA NeMo-RL.
• Saiba mais sobre o hipercomputador de IA do Google Cloud.
• Leia a documentação do Managed Service for Lustre.