Executar cargas de trabalho pré-criadas

Se você tiver interesse no treinamento gerenciado da Vertex AI, 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 uma seção dedicada:

  • NVIDIA NeMo: para o desenvolvimento de modelos fundamentais, 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 aprendizado por reforço (RL) e 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 vai orientar 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 nas etapas fundamentais do desenvolvimento de modelos. Ele fornece instruções detalhadas para usar o NeMo e 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. Verifique se você tem uma imagem de contêiner e os scripts de treinamento necessários.

Preparar uma imagem de contêiner

Há 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/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: autentique com Google Cloud.

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

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

Etapa 2: crie o script de conversão.

Crie um arquivo chamado enroot-convert.sh e adicione o seguinte conteúdo de script. 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: execute o script e verifique 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.

Baixar as receitas de treinamento

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

  1. Gere 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. Clone o repositório.

    Depois que as credenciais forem autenticadas, execute o comando a seguir para baixar as receitas.

    git clone 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 seu 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>/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. É possível 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.

Como monitorar o status e os registros de jobs

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, você pode 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 comando de inicialização. Para ver uma lista de partições disponíveis, execute o comando sinfo no nó de login do Slurm.

sinfo

A saída mostra os nomes de partição disponíveis, como a3u neste exemplo:

PARTITION AVAIL TIMELIMIT NÓS ESTADO NODELIST
a3u* cima infinito 2 idle~ 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. Executar o 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 nova execução do job falhar, siga estas etapas:
    • Faça o download dos dois arquivos a seguir:
      • gpt2-vocab.json
      • gpt2-merges.txt
    • Mova os arquivos baixados para o subdiretório torch/megatron/ no diretório de cache (por exemplo, <var>YOUR_CACHE_DIR</var>/torch/megatron/).
    • Renomeie os arquivos da seguinte forma:
      • gpt2-vocab.json foi renomeado como megatron-gpt-345m_vocab.
      • gpt2-merges.txt foi renomeado como megatron-gpt-345m_merges.

NVIDIA NeMo-RL

O framework NVIDIA NeMo-RL foi criado para alinhar modelos de linguagem grandes com preferências e instruções humanas. Esta seção mostra como usar o NeMo-RL em um cluster para realizar tarefas de alinhamento avançadas, incluindo ajuste fino supervisionado (SFT), ajuste de preferência (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 existente, se você 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 Máquina virtual do Google Compute Engine no console do Google Google Cloud 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 do Docker predefinida

Os arquivos .sqsh convertidos são fornecidos para imagens de contêiner pré-criadas. É possível 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 definir diretamente como o parâmetro de imagem do contêiner, use um dos seguintes caminhos. Substitua <region> pela sua região específica (por exemplo, europe, asia, us):

/gcs/vmds-containers-<region>/nemo_rl_squashfs/nemo_rl-h20250923.sqsh

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

gs://vmds-containers-<region>/nemo_rl_squashfs/nemo_rl-h20250923.sqsh DESTINATION

Baixar o código

Para acessar a receita de treinamento com a CLI git, acesse https://www.googlesource.com/new-password. É possível fazer o download da receita com o seguinte comando:

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

Iniciar jobs

Etapa 1: defina as variáveis de ambiente.

Para extrair modelos e dados do Hugging Face, talvez seja necessário definir o HF_TOKEN. Para usar o Weights & Biases na análise de experimentos, o WANDB_API_KEY precisa ser definido. Atualize estas variáveis no seguinte arquivo:

Arquivo para atualizar: $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: baixe ou copie o arquivo do contêiner para a pasta de inicialização.

Confira alguns exemplos.

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


# OR

/gcs/vmds-containers-<region>/vmds_nemo_rl_squashfs/nemo_rl-h20250923.sqsh \
  $HOME/vertex-oss-training/nemo_rl/nemo_rl-h20250923.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: inicie 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, e 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. Lá você encontra todos os registros e pontos de verificação pertencentes a esse job. Mais precisamente, dentro dele, você vai encontrar os seguintes diretórios e arquivos:

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

Desenvolvimento com o NeMo-RL

Configuração interativa

Há duas opções 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ó número 5) e leva você a um contêiner em execução do 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 para o trabalho no arquivo configs/auth.sh.

  2. Defina a variável de ambiente CLUSTER_TYPE de acordo com a seguinte diretriz:

    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 usando o 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

Início interativo

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 até que o <job_id>/attach.sh seja criado.

  • Para monitorar o progresso da verificação de lançamento, execute <job_id>/nemo_rl_output.log e confira o progresso do script de lançamento. Execute <job_id>/ray_logs/ e confira o progresso do lançamento do cabeçalho e dos workers do Ray.

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

bash <job_id>/attach.sh

A seguir

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