Executar cargas de trabalho pré-criadas

Se você tiver interesse no treinamento gerenciado da Gemini Enterprise Agent Platform, entre em contato com seu representante de vendas para ter acesso.

Este guia mostra como usar o ecossistema NVIDIA NeMo em um cluster de treinamento gerenciado para o desenvolvimento de modelos de IA generativa de ponta a ponta. Ele fornece instruções detalhadas para os seguintes fluxos de trabalho distintos, mas relacionados, cada um abordado em sua própria seção:

  • NVIDIA NeMo: para o desenvolvimento de modelos de fundação, siga estas instruções para realizar pré-treinamento em grande escala, pré-treinamento contínuo (CPT) e ajuste supervisionado (SFT).
  • NVIDIA NeMo-RL: para alinhamento de modelos e ajuste de preferências, use esta seção para aplicar técnicas avançadas, como o aprendizado por reforço (RL), para alinhar seu modelo com instruções e preferências humanas.

Se você estiver criando um modelo do zero ou refinando um já existente, este documento orienta você na configuração do ambiente, no gerenciamento de jobs em contêineres e na execução de scripts de treinamento no cluster.

NVIDIA NeMo

O framework NVIDIA NeMo é uma plataforma completa para criar, personalizar, e implantar modelos de IA generativa. Esta seção do guia é especificamente para desenvolvedores e pesquisadores focados nos estágios fundamentais do desenvolvimento de modelos. Ela fornece instruções detalhadas para usar o NeMo para realizar pré-treinamento em grande escala, pré-treinamento contínuo (CPT) e ajuste fino supervisionado (SFT) em um cluster de treinamento gerenciado.

Este guia de clusters de treinamento fornece o fluxo de trabalho completo para executar um job de treinamento com o framework NeMo. O processo é dividido em duas partes principais: a configuração inicial única do ambiente e as etapas recorrentes para iniciar um job.

Configurar o ambiente

Antes de iniciar um job, prepare o ambiente. Para isso, verifique se você tem uma imagem do contêiner e os scripts de treinamento necessários.

Instalar o NeMo-Run

pip install git+https://github.com/NVIDIA/NeMo-Run.git

Preparar uma imagem do contêiner

Você tem duas opções para a imagem do contêiner: usar uma imagem pré-criada (recomendado) ou criar uma personalizada.

Usar uma imagem pré-criada (recomendado)

As imagens de contêiner pré-criadas são fornecidas no formato .squashfs. Copie a imagem adequada para sua região no diretório de trabalho.

# Example for the US region
gcloud storage cp gs://vmds-containers-us/vmds_nemo_squashfs/nemo-20250721.sqsh .
Criar um contêiner personalizado (avançado)

Siga estas etapas apenas se os contêineres pré-criados não atenderem às suas necessidades. Este procedimento orienta você na conversão de uma imagem de contêiner personalizada para o formato .squashfs usando enroot.

Etapa 1: autenticar com Google Cloud.

Use os comandos a seguir para garantir que sua conta de usuário Google Cloud e o registro do Docker em que a imagem está hospedada sejam autenticados:

gcloud auth login
gcloud auth configure-docker us-docker.pkg.dev

Etapa 2: criar o script de conversão.

Crie um arquivo chamado enroot-convert.sh e adicione o conteúdo do script a seguir. Antes de executar esse script, atualize as variáveis REMOTE_IMG e LOCAL_IMG para apontar para a imagem do contêiner e o caminho de saída escolhido.

#!/bin/bash

#SBATCH --gpus-per-node=8
#SBATCH --exclusive
#SBATCH --mem=0
#SBATCH --ntasks-per-node=1

# Run this script on the slurm login node:
# sbatch -N 1 enroot-convert.sh

set -x
set -e

# The remote docker image URI.
REMOTE_IMG="docker://us-docker.pkg.dev/{YOUR_CONTAINER_IMG_URI}:{YOUR_CONTAINER_IMAGE_TAG}"

# The local path to the to be imported enroot squash file.
LOCAL_IMG="${HOME}/my_nemo.sqsh"

# The path to the enroot config file.
TMP_ENROOT_CONFIG_PATH="/tmp/\$(id -u --name)/config/enroot"

# Download the docker image to each node.
srun -l -N "${SLURM_NNODES}" \
bash -c "
mkdir -p ${TMP_ENROOT_CONFIG_PATH};
echo 'machine us-docker.pkg.dev login oauth2accesstoken password $(gcloud auth print-access-token)' > ${TMP_ENROOT_CONFIG_PATH}/.credentials;
rm -f ${LOCAL_IMG};
ENROOT_CONFIG_PATH=${TMP_ENROOT_CONFIG_PATH} ENROOT_MAX_PROCESSORS=$(( $(nproc) / 2 )) enroot import -o ${LOCAL_IMG} ${REMOTE_IMG};
"

Etapa 3: executar o script e verificar a saída.

Execute o script no nó de login do Slurm.

sbatch -N 1 enroot-convert.sh

Depois que o job for concluído, encontre os registros de conversão em um arquivo chamado slurm-<JOB_ID>.out e a imagem do contêiner final no caminho especificado para LOCAL_IMG.

Fazer o download das receitas de treinamento

As receitas de treinamento são armazenadas em um repositório particular googlesource.com. Para acessá-las com a linha de comando do Git, primeiro gere credenciais de autenticação.

  1. Gerar credenciais de autenticação.

    Acesse o URL a seguir e siga as instruções na tela. Isso configura seu ambiente local para autenticar com o repositório. https://www.googlesource.com/new-password

  2. Clonar o repositório.

    Depois que as credenciais forem autenticadas, execute o comando a seguir para fazer o download das receitas.

    git clone -b release-allowlist-ga https://vertex-model-garden.googlesource.com/vertex-oss-training
    

Iniciar um job de treinamento

Depois que o ambiente estiver configurado, você poderá iniciar um job de treinamento.

Etapa 1: definir variáveis de ambiente

As seguintes variáveis de ambiente podem ser necessárias para o job:

  • O HF_TOKEN é necessário para fazer o download de modelos e conjuntos de dados do Hugging Face.
  • O WANDB_API_KEY é necessário para usar o Weights & Biases na análise de experimentos.
export HF_TOKEN=YOUR_HF_TOKEN
export WANDB_API_KEY=YOUR_WANDB_API_KEY

Etapa 2: executar o script de inicialização

Navegue até o diretório de trabalho e execute o script run.py para iniciar um job. Este exemplo inicia um job de treinamento de demonstração com o Llama 3.1-2b.

# Set the working directory
export WORK_DIR=$HOME/vertex-oss-training/nemo
cd $WORK_DIR

gcloud storage cp
gs://vmds-containers-<region>/vmds_nemo_squashfs/nemo-20250721.sqsh nemo-demo.sqsh

# Launch the training job
export NEMORUN_HOME=$WORK_DIR && \
python3 run.py -e slurm --slurm-type hcc-a3m --partition a3m \
  -d $WORK_DIR -i $WORK_DIR/nemo-demo.sqsh \
  -s pretrain/llama3p1_2b_pt.py -n 2 \
  --experiment-name nemo-demo-run

Parâmetros de inicialização

  • --slurm-type é definido com base no tipo de cluster (por exemplo, hcc-a3m, hcc-a3u, hcc-a4).
  • --partition precisa ser definido como uma partição disponível. Você pode verificar os nomes das partições com o comando sinfo.
  • O script run.py monta automaticamente vários diretórios no contêiner do Docker, incluindo --log-dir, --cache-dir e --data-dir, se eles estiverem definidos.

Monitorar o status e os registros do job

Depois de iniciar o job, um bloco de status será exibido:

Experiment Status for nemo-demo-run_1753123402

Task 0: nemo-demo-run
- Status: RUNNING
- Executor: SlurmExecutor on @localhost
- Job id: 75
- Local Directory: $NEMORUN_HOME/experiments/nemo-demo-run/nemo-demo-run_1753123402/nemo-demo-run

Os registros de execução são gravados no caminho mostrado no campo Local Directory da saída de status. Por exemplo, é possível encontrar os arquivos de registro em um caminho semelhante a este:

$NEMORUN_HOME/experiments/nemo-demo-run/nemo-demo-run_1753123402/nemo-demo-run/<JOB_ID>.log

Erros comuns e soluções

Esta seção descreve problemas comuns que podem surgir durante a execução do job e fornece etapas recomendadas para resolvê-los.

Erro de partição inválida

Por padrão, os jobs tentam ser iniciados na partição geral. Se a partição geral não existir ou não estiver disponível, o job vai falhar com o seguinte erro:

sbatch: error: invalid partition specified: general
sbatch: error: Batch job submission failed: Invalid partition name specified

Solução:

Especifique uma partição disponível usando o argumento --partition ou -p no seu comando de inicialização. Para conferir uma lista de partições disponíveis, execute o comando sinfo no nó de login do Slurm.

sinfo

A saída mostra os nomes das partições disponíveis, como a3u neste exemplo:

PARTITION AVAIL TIMELIMIT NODES ESTADO NODELIST
a3u* cima infinito 2 inativo~ alice-a3u-[2-3]
a3u* cima infinito 2 inativo alice-a3u-[0-1]

Erro de download do tokenizador

Você pode encontrar um OSError relacionado a um link entre dispositivos quando o script tenta fazer o download do tokenizador GPT2:

OSError: [Errno 18] Invalid cross-device link: 'gpt2-vocab.json' -> '/root/.cache/torch/megatron/megatron-gpt-345m_vocab'

Soluções:

Você tem duas opções para resolver esse problema:

  • Opção 1: execute o job novamente. Esse erro geralmente é temporário. A execução do job novamente usando o mesmo --cache-dir pode resolver o problema.
  • Opção 2: faça o download manual dos arquivos do tokenizador. Se a execução do job falhar novamente, siga estas etapas:
    • Faça o download dos dois arquivos a seguir:
      • gpt2-vocab.json
      • gpt2-merges.txt
    • Mova os arquivos salvos para o torch/megatron/ subdiretório dentro do seu diretório de cache (por exemplo, <var>YOUR_CACHE_DIR</var>/torch/megatron/).
    • Renomeie os arquivos da seguinte maneira:
      • Renomeie gpt2-vocab.json para megatron-gpt-345m_vocab.
      • Renomeie gpt2-merges.txt para megatron-gpt-345m_merges.

NVIDIA NeMo-RL

O framework NVIDIA NeMo-RL foi projetado para alinhar modelos de linguagem grandes com preferências e instruções humanas. Esta seção orienta você no uso do NeMo-RL em um cluster para realizar tarefas de alinhamento avançadas, incluindo ajuste fino supervisionado (SFT), ajuste de preferências (como otimização direta de preferências ou DPO) e aprendizado por reforço (RL).

O guia aborda dois fluxos de trabalho principais: executar um job de treinamento em lote padrão e usar o ambiente de desenvolvimento interativo para depuração.

Pré-requisitos

Antes de começar, crie um cluster seguindo as instruções na página Criar cluster ou use um cluster de treinamento gerenciado já existente, se tiver um.

Conectar-se ao nó de login do cluster

Para se conectar ao nó de login do cluster, encontre o comando correto da Google Cloud CLI navegando até a página da máquina virtual do Google Compute Engine no Google Google Cloud console e clicando em SSH > Ver comando da Google Cloud CLI. Será semelhante a:

ssh $USER_NAME@machine-addr

Exemplo:

ssh $USER_NAME@nic0.sliua3m1-login-001.europe-north1-c.c.infinipod-shared-dev.internal.gcpnode.com

Usar a imagem Docker pré-criada

Os arquivos .sqsh convertidos são fornecidos para imagens de contêiner pré-criadas. Você pode selecionar um contêiner para sua região e defini-lo diretamente como o parâmetro de imagem do contêiner ou fazer o download dele para o sistema de arquivos do cluster.

Para defini-lo diretamente como o parâmetro de imagem do contêiner, use um dos caminhos a seguir. Observe que você deve substituir <region> pela sua região específica (por exemplo, europe, asia, us):

/gcs/vmds-containers-<region>/vmds_nemo_rl_squashfs/nemo_rl-h20260302.sqsh

Para fazer o download da imagem para o armazenamento lustre do cluster, use o seguinte comando:

gs://vmds-containers-<region>/vmds_nemo_rl_squashfs/nemo_rl-h20260302.sqsh DESTINATION

Baixar o código

Para ter acesso à receita de treinamento com a CLI do Git, acesse https://www.googlesource.com/new-password. A receita pode ser baixada com o seguinte comando:

cd $HOME
git clone -b release-allowlist-ga https://vertex-model-garden.googlesource.com/vertex-oss-training

Iniciar jobs

Etapa 1: definir variáveis de ambiente.

Para extrair modelos e dados do Hugging Face, o HF_TOKEN pode precisar ser definido. Para usar o Weights &Biases na análise de experimentos, o WANDB_API_KEY precisa ser definido. Atualize essas variáveis no arquivo a seguir:

Arquivo a ser atualizado: $HOME/vertex-oss-training/nemo_rl/configs/auth.sh

Se você não quiser usar o Weights & Biases, defina logger.wandb_enabled como False no script de inicialização.

Etapa 2: faça o download ou copie o arquivo de contêiner para a pasta de inicialização.

Confira alguns exemplos.

gcloud storage cp \
  gs://vmds-containers-<region>/vmds_nemo_rl_squashfs/nemo_rl-h20260302.sqsh \
  $HOME/vertex-oss-training/nemo_rl/nemo_rl-h20260302.sqsh


# OR

/gcs/vmds-containers-<region>/vmds_nemo_rl_squashfs/nemo_rl-h20260302.sqsh \
  $HOME/vertex-oss-training/nemo_rl/nemo_rl-h20260302.sqsh

cd $HOME/vertex-oss-training/nemo_rl/

# Where region is either `us`, `asia`, or `europe`

Etapa 3: prepare ou clone o repositório NeMo-RL.

Crie um clone do código NeMo-RL se ele ainda não estiver presente. Talvez seja necessário usar git submodule update --init --recursive se você já tiver clonado o repositório sem a flag --recursive.

git clone https://github.com/NVIDIA-NeMo/RL --recursive

Etapa 4: iniciar o job de treinamento.

sbatch -N <num_nodes> launch.sh --cluster_type hcc-a3m --job_script algorithms/dpo.sh

Em que:

  • --cluster-type é definido com base no tipo de cluster:
    • A3-Mega: hcc-a3m
    • A3-Ultra: hcc-a3u
    • A4: hcc-a4
    • A3H: hcc-a3h
  • --partition precisa ser definido de acordo com o caso, em que sinfo pode ser usado para verificar as partições do Slurm.

Depois que o job for iniciado, um novo diretório com o nome do ID do job do SLURM será criado no local atual. Dentro dele, você encontrará todos os registros e checkpoints pertencentes a esse job. Mais precisamente, dentro dele, você encontrará os seguintes diretórios e arquivos:

  • checkpoints/ → Esse diretório é montado dentro do contêiner NeMo-RL e contém todos os checkpoints do treinamento.
  • ray-logs/ → Esse diretório contém os registros do cabeçalho e dos workers do Ray.
  • nemo_rl_output.log → Esse arquivo contém os registros do Slurm do job enviado.
  • attach.sh (somente jobs interativos) → Esse é um script bash que permite anexar a um job interativo. Se o job for iniciado com sucesso, pode levar alguns minutos para que esse arquivo seja criado.

Desenvolvimento com NeMo-RL

Configuração interativa

Duas opções estão disponíveis para desenvolvimento interativo rápido com o NeMo-RL.

nemorlinteractive

Esse é um comando auxiliar simples que permite escolher um nó de GPU no cluster (digamos, o número 5) e, em seguida, leva você a um contêiner em execução para o NeMo-RL dentro do nó selecionado. Esse comando é útil para fluxos de trabalho de nó único.

Para usar nemorlinteractive, siga estas etapas de pré-requisito:

  1. Forneça todos os tokens de autenticação que você quer (por exemplo, HF e WandB) carregados no job no arquivo configs/auth.sh.
  2. Defina a variável de ambiente CLUSTER_TYPE de acordo com a diretriz a seguir:

    export CLUSTER_TYPE="hcc-a3m" # --> if you have A3-Mega cluster
    export CLUSTER_TYPE="hcc-a3u" # --> if you have A3-Ultra cluster
    export CLUSTER_TYPE="hcc-a4"  # --> If you have A4 cluster
    export CLUSTER_TYPE="hcc-a3h" # --> If you have A3H cluster
    
  3. Importe nemorlinteractive no terminal bash de origem do bash_utils.sh:

    source bash_utils.sh
    
  4. Execute o comando nemorlinteractive. Exemplo:

    # Assuming you want to take the compute node number 5.
    nemorlinteractive 5
    

Inicialização interativa

Essa opção permite executar cargas de trabalho de forma interativa em vários nós de computação. Os jobs interativos são mais adequados para casos de uso de depuração e verificação. Essas cargas de trabalho reservam o nó indefinidamente até que o desenvolvedor decida que a depuração foi concluída e queira liberar os recursos.

Estas são as etapas que precisam ser seguidas para essa opção:

Forneça todos os tokens de autenticação que você quer (por exemplo, HF e WandB) carregados no job no arquivo configs/auth.sh.

sbatch -N <num_nodes> launch.sh --cluster_type hcc-a3m --interactive
  • Aguarde de 2 a 5 minutos e você verá <job_id>/attach.sh criado.

  • Para monitorar o progresso da inicialização, verifique <job_id>/nemo_rl_output.log para ver o progresso do script de inicialização e verifique <job_id>/ray_logs/ para ver o progresso do cabeçalho e dos workers do Ray.

  • Conecte-se ao job interativo. Esse script permite que você se conecte novamente mesmo se perder a conexão:

bash <job_id>/attach.sh

A seguir

A execução de uma carga de trabalho pré-criada verifica o status operacional do cluster. A próxima etapa é executar seu próprio aplicativo de treinamento personalizado.