Ejecuta cargas de trabajo prediseñadas

Si te interesa el entrenamiento administrado de Gemini Enterprise Agent Platform, comunícate con tu representante de ventas para obtener acceso.

En esta guía, se muestra cómo usar el ecosistema de NVIDIA NeMo en un clúster de entrenamiento administrado para el desarrollo de modelos de IA generativa de extremo a extremo. Se proporcionan instrucciones paso a paso para los siguientes flujos de trabajo distintos, pero relacionados, cada uno de los cuales se aborda en su propia sección:

  • NVIDIA NeMo: Para el desarrollo de modelos de base, sigue estas instrucciones para realizar un entrenamiento previo a gran escala, un entrenamiento previo continuo (CPT) y un ajuste supervisado (SFT).
  • NVIDIA NeMo-RL: Para el ajuste de preferencias y la alineación de modelos, usa esta sección para aplicar técnicas avanzadas como el aprendizaje por refuerzo (RL) para alinear tu modelo con las instrucciones y preferencias humanas.

Ya sea que estés compilando un modelo desde cero o refinando uno existente, este documento te guiará para configurar tu entorno, administrar trabajos en contenedores y ejecutar secuencias de comandos de entrenamiento en el clúster.

NVIDIA NeMo

El framework de NVIDIA NeMo es una plataforma de extremo a extremo para compilar, personalizar e implementar modelos de IA generativa. Esta sección de la guía está destinada específicamente a desarrolladores e investigadores enfocados en las etapas fundamentales del desarrollo de modelos. Proporciona instrucciones paso a paso para usar NeMo para realizar un entrenamiento previo a gran escala, un entrenamiento previo continuo (CPT) y un ajuste supervisado (SFT) en un clúster de entrenamiento administrado.

Esta guía de clústeres de entrenamiento proporciona el flujo de trabajo completo para ejecutar un trabajo de entrenamiento con el framework de NeMo. El proceso se divide en dos partes principales: la configuración inicial única de tu entorno y los pasos recurrentes para iniciar un trabajo.

Configura tu entorno

Antes de iniciar un trabajo, debes preparar tu entorno. Para ello, asegúrate de tener una imagen de contenedor y las secuencias de comandos de entrenamiento necesarias.

Instala NeMo-Run

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

Prepara una imagen de contenedor

Tienes dos opciones para la imagen de contenedor: usar una imagen compilada previamente (recomendado) o compilar una personalizada.

Usa una imagen compilada previamente (recomendado)

Las imágenes de contenedor compiladas previamente se proporcionan en formato .squashfs. Copia la imagen adecuada para tu región en tu directorio de trabajo.

# Example for the US region
gcloud storage cp gs://vmds-containers-us/vmds_nemo_squashfs/nemo-20250721.sqsh .
Compila un contenedor personalizado (avanzado)

Sigue estos pasos solo si los contenedores compilados previamente no satisfacen tus necesidades. Este procedimiento te guía para convertir una imagen de contenedor personalizada al formato .squashfs con enroot.

Paso 1: Autentícate con Google Cloud.

Usa los siguientes comandos para asegurarte de que tu Google Cloud cuenta de usuario y el registro de Docker en el que se aloja tu imagen estén autenticados:

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

Paso 2: Crea la secuencia de comandos de conversión.

Crea un archivo llamado enroot-convert.sh y agrega el siguiente contenido de secuencia de comandos. Antes de ejecutar esta secuencia de comandos, debes actualizar las variables REMOTE_IMG y LOCAL_IMG para que apunten a tu imagen de contenedor y a la ruta de acceso de salida elegida.

#!/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};
"

Paso 3: Ejecuta la secuencia de comandos y verifica el resultado.

Ejecuta la secuencia de comandos en el nodo de acceso de Slurm.

sbatch -N 1 enroot-convert.sh

Una vez que se complete el trabajo, busca los registros de conversión en un archivo llamado slurm-<JOB_ID>.out y la imagen de contenedor final en la ruta de acceso que especificaste para LOCAL_IMG.

Descarga las recetas de entrenamiento

Las recetas de entrenamiento se almacenan en un repositorio privado googlesource.com. Para acceder a ellas con la línea de comandos de Git, primero debes generar credenciales de autenticación.

  1. Genera credenciales de autenticación.

    Visita la siguiente URL y sigue las instrucciones en pantalla. Esto configura tu entorno local para autenticarse con el repositorio. https://www.googlesource.com/new-password

  2. Clona el repositorio.

    Una vez que se autentiquen las credenciales, ejecuta el siguiente comando para descargar las recetas.

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

Inicia un trabajo de entrenamiento

Una vez que tu entorno esté configurado, puedes iniciar un trabajo de entrenamiento.

Paso 1: Establece variables de entorno

Es posible que se requieran las siguientes variables de entorno para tu trabajo:

  • Se requiere HF_TOKEN para descargar modelos y conjuntos de datos de Hugging Face.
  • Se requiere WANDB_API_KEY para usar Weights & Biases para el análisis de experimentos.
export HF_TOKEN=YOUR_HF_TOKEN
export WANDB_API_KEY=YOUR_WANDB_API_KEY

Paso 2: Ejecuta la secuencia de comandos de inicio

Navega hasta tu directorio de trabajo y ejecuta la secuencia de comandos run.py para iniciar un trabajo. En este ejemplo, se inicia un trabajo de entrenamiento de demostración con 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 inicio

  • --slurm-type se establece según el tipo de clúster (por ejemplo, hcc-a3m, hcc-a3u, hcc-a4).
  • --partition debe establecerse en una partición disponible. Puedes verificar los nombres de las particiones con el comando sinfo.
  • La secuencia de comandos run.py activa automáticamente varios directorios en el contenedor de Docker, incluidos --log-dir, --cache-dir y --data-dir, si están configurados.

Supervisa el estado y los registros del trabajo

Después de iniciar el trabajo, se muestra un bloque de estado:

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

Los registros de ejecución se escriben en la ruta de acceso que se muestra en el campo Local Directory del resultado de estado. Por ejemplo, puedes encontrar los archivos de registro en una ruta de acceso similar a la siguiente:

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

Errores comunes y soluciones

En esta sección, se describen los problemas comunes que pueden surgir durante la ejecución del trabajo y se proporcionan los pasos recomendados para resolverlos.

Error de partición no válida

De forma predeterminada, los trabajos intentan iniciarse en la partición general. Si la partición general no existe o no está disponible, el trabajo fallará con el siguiente error:

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

Solución:

Especifica una partición disponible con el argumento --partition o -p en tu comando de inicio. Para ver una lista de las particiones disponibles, ejecuta el comando sinfo en el nodo de acceso de Slurm.

sinfo

En el resultado, se muestran los nombres de las particiones disponibles, como a3u en este ejemplo:

PARTITION AVAIL TIMELIMIT NODES ESTADO NODELIST
a3u* up infinite 2 idle~ alice-a3u-[2-3]
a3u* up infinite 2 idle alice-a3u-[0-1]

Error de descarga del tokenizador

Es posible que encuentres un OSError relacionado con un vínculo multidispositivo cuando la secuencia de comandos intente descargar el tokenizador GPT2:

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

Soluciones:

Tienes dos opciones para resolver este problema:

  • Opción 1: Vuelve a ejecutar el trabajo. Este error suele ser transitorio. Volver a ejecutar el trabajo con el mismo --cache-dir puede resolver el problema.
  • Opción 2: Descarga manualmente los archivos del tokenizador. Si no se vuelve a ejecutar el trabajo, sigue estos pasos:
    • Descarga los siguientes dos archivos:
      • gpt2-vocab.json
      • gpt2-merges.txt
    • Mueve los archivos descargados al subdirectorio torch/megatron/ dentro de tu directorio de caché (por ejemplo, <var>YOUR_CACHE_DIR</var>/torch/megatron/).
    • Cambia el nombre de los archivos de la siguiente manera:
      • Cambia el nombre de gpt2-vocab.json a megatron-gpt-345m_vocab.
      • Cambia el nombre de gpt2-merges.txt a megatron-gpt-345m_merges.

NVIDIA NeMo-RL

El framework de NVIDIA NeMo-RL está diseñado para alinear modelos grandes de lenguaje con las instrucciones y preferencias humanas. En esta sección, se te guía para usar NeMo-RL en un clúster para realizar tareas de alineación avanzadas, incluido el ajuste supervisado (SFT), el ajuste de preferencias (como la optimización directa de preferencias o DPO) y el aprendizaje por refuerzo (RL).

La guía abarca dos flujos de trabajo principales: ejecutar un trabajo de entrenamiento por lotes estándar y usar el entorno de desarrollo interactivo para la depuración.

Requisitos previos

Antes de comenzar, crea un clúster siguiendo las instrucciones de la página Crea un clúster o usa un clúster de entrenamiento administrado existente, si tienes uno.

Conéctate al nodo de acceso del clúster

Para conectarte al nodo de acceso del clúster, busca el comando correcto de Google Cloud CLI . Para ello, navega a la página Máquina virtual de Google Compute Engine en la consola de Google Google Cloud y haz clic en SSH > Ver comando de Google Cloud CLI. Tendrán un aspecto similar a este:

ssh $USER_NAME@machine-addr

Ejemplo:

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

Usa la imagen de Docker compilada previamente

Se proporcionan archivos .sqsh convertidos para imágenes de contenedor compiladas previamente. Puedes seleccionar un contenedor para tu región y configurarlo directamente como el parámetro de imagen de contenedor o descargarlo en el sistema de archivos del clúster.

Para configurarlo directamente como el parámetro de imagen de contenedor, usa una de las siguientes rutas de acceso. Ten en cuenta que debes reemplazar <region> por tu región específica (por ejemplo, europe, asia, us):

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

Para descargar la imagen en el almacenamiento de Lustre del clúster, usa el siguiente comando:

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

Descargar el código

Para obtener acceso a la receta de entrenamiento con la CLI de git, visita https://www.googlesource.com/new-password. La receta se puede descargar con el siguiente comando:

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

Iniciar trabajos

Paso 1: Establece variables de entorno.

Para extraer modelos y datos de Hugging Face, es posible que se deba configurar HF_TOKEN. Para usar Weights &Biases para el análisis de experimentos, se debe configurar WANDB_API_KEY. Actualiza estas variables en el siguiente archivo:

Archivo para actualizar: $HOME/vertex-oss-training/nemo_rl/configs/auth.sh

Si no quieres usar Weights & Biases, establece logger.wandb_enabled en False en tu secuencia de comandos de inicio.

Paso 2: Descarga o copia el archivo del contenedor en tu carpeta de inicio.

Aquí tienes algunos ejemplos.

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`

Paso 3: Prepara o clona el repositorio de NeMo-RL.

Crea un clon del código de NeMo-RL si aún no está presente. Ten en cuenta que es posible que debas usar git submodule update --init --recursive si ya clonaste el repositorio sin la marca --recursive

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

Paso 4: Inicia el trabajo de entrenamiento.

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

Aquí:

  • --cluster-type se establece según el tipo de clúster:
    • A3-Mega: hcc-a3m
    • A3-Ultra: hcc-a3u
    • A4: hcc-a4
    • A3H: hcc-a3h
  • --partition debe establecerse en consecuencia, donde se puede usar sinfo para verificar las particiones de Slurm.

Después de que se inicie el trabajo, se creará un directorio nuevo con el nombre de su ID de trabajo de SLURM en tu ubicación actual. En el interior, encontrarás todos los registros y puntos de control que pertenecen a este trabajo. Más precisamente, dentro de él, encontrarás los siguientes directorios y archivos:

  • checkpoints/ : Este directorio se activa dentro del contenedor de NeMo-RL y contiene todos los puntos de control del entrenamiento.
  • ray-logs/: Este directorio contiene los registros del encabezado de Ray y los trabajadores de Ray.
  • nemo_rl_output.log: Este archivo contiene los registros de Slurm de tu trabajo enviado.
  • attach.sh (solo trabajos interactivos): Es una secuencia de comandos de Bash que te permite adjuntar a un trabajo interactivo. Si el trabajo se inicia correctamente, es posible que tarde unos minutos en crearse este archivo.

Desarrollo con NeMo-RL

Configuración interactiva

Hay dos opciones disponibles para el desarrollo interactivo rápido con NeMo-RL.

nemorlinteractive

Este es un comando auxiliar directo que te permite elegir un nodo de GPU del clúster (por ejemplo, el nodo número 5) y, luego, te lleva a un contenedor en ejecución para NeMo-RL dentro del nodo seleccionado. Este comando es útil para tus flujos de trabajo de un solo nodo.

Para usar nemorlinteractive, debes seguir estos pasos de requisitos previos:

  1. Proporciona todos los tokens de autenticación que desees (por ejemplo, HF y WandB) cargados en el trabajo en el archivo configs/auth.sh.

  2. Establece la variable de entorno CLUSTER_TYPE según la siguiente guía:

    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. Importa nemorlinteractive en tu terminal de Bash con la fuente bash_utils.sh:

    source bash_utils.sh

  4. Ejecuta el comando nemorlinteractive. Por ejemplo:

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

Inicio interactivo

Esta opción te permite ejecutar cargas de trabajo de forma interactiva en varios nodos de procesamiento. Los trabajos interactivos son más adecuados para casos de uso de depuración y verificación. Estas cargas de trabajo reservan el nodo de forma indefinida hasta que el desarrollador decide que finalizó la depuración y desea liberar los recursos.

Estos son los pasos que se deben seguir para esta opción:

Proporciona todos los tokens de autenticación que desees (por ejemplo, HF y WandB) cargados en el trabajo en el archivo configs/auth.sh.

sbatch -N <num_nodes> launch.sh --cluster_type hcc-a3m --interactive

  • Espera de 2 a 5 minutos y deberías ver que se creó <job_id>/attach.sh.

  • Para supervisar el progreso del inicio, consulta <job_id>/nemo_rl_output.log para ver el progreso de la secuencia de comandos de inicio y consulta <job_id>/ray_logs/ para ver el progreso del inicio del encabezado y los trabajadores de Ray.

  • Conéctate al trabajo interactivo. Esta secuencia de comandos te permite volver a conectarte incluso si pierdes la conexión:

bash <job_id>/attach.sh

¿Qué sigue?

La ejecución de una carga de trabajo compilada previamente verifica el estado operativo del clúster. El siguiente paso es ejecutar tu propia aplicación de entrenamiento personalizada.