Ejecuta cargas de trabajo prediseñadas

Si te interesa el entrenamiento administrado de Vertex AI, 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. Proporciona 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 dedicada:

  • NVIDIA NeMo: Para el desarrollo de modelos básicos, sigue estas instrucciones para realizar el entrenamiento previo a gran escala, el entrenamiento previo continuo (CPT) y el ajuste supervisado (SFT).
  • NVIDIA NeMo-RL: Para el ajuste de preferencias y la alineación del modelo, 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 creando un modelo desde cero o perfeccionando uno existente, este documento te guía para configurar tu entorno, administrar trabajos en contenedores y lanzar secuencias de comandos de entrenamiento en el clúster.

NVIDIA NeMo

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

En esta guía sobre clústeres de entrenamiento, se 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 del 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. En este procedimiento, se explica cómo 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 se autentiquen tanto tu cuenta de usuario de Google Cloud como el registro de Docker en el que se aloja tu imagen:

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 agrégale 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 que elegiste.

#!/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 de googlesource.com. Para acceder a ellos 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 que se autentique 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 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: Configura las variables de entorno

Es posible que tu trabajo requiera las siguientes variables de entorno:

  • El HF_TOKEN es necesario para descargar modelos y conjuntos de datos de Hugging Face.
  • Se requiere WANDB_API_KEY para usar Weights & Biases en 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 el script 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 de los trabajos

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 del 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á y mostrará 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 el 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 NODOS ESTADO NODELIST
a3u* up infinito 2 idle~ alice-a3u-[2-3]
a3u* up infinito 2 Inactivo alice-a3u-[0-1]

Error de descarga del tokenizador

Es posible que encuentres un OSError relacionado con una vinculación entre dispositivos 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 puede volver 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:
      • Se cambió el nombre de gpt2-vocab.json a megatron-gpt-345m_vocab.
      • Se cambió 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 los modelos de lenguaje grandes con las preferencias y las instrucciones humanas. En esta sección, se explica cómo usar NeMo-RL en un clúster para realizar tareas de alineación avanzadas, como el ajuste fino supervisado (SFT), el ajuste de preferencias (como la optimización directa de preferencias o DPO) y el aprendizaje por refuerzo (RL).

En la guía, se abordan 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 Crear 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

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

Para establecerlo directamente como el parámetro de imagen del 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 o us):

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

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

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

Descargar el código

Para acceder 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 https://vertex-model-garden.googlesource.com/vertex-oss-training

Iniciar trabajos

Paso 1: Configura las variables de entorno.

Es posible que se deba configurar HF_TOKEN para extraer modelos y datos de Hugging Face. Para usar Weights & Biases en 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, configura logger.wandb_enabled como False en tu secuencia de comandos de inicio.

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

Aquí tienes algunos ejemplos.

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`

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
    • R.: hcc-a4
    • A3H: hcc-a3h
  • --partition se debe configurar según corresponda, donde sinfo se puede usar 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 la ubicación actual. Allí, encontrarás todos los registros y puntos de control correspondientes 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 nodo principal y de los nodos trabajadores de Ray.
  • nemo_rl_output.log: Este archivo contiene los registros de Slurm del trabajo que enviaste.
  • attach.sh (solo para trabajos interactivos) → Es una secuencia de comandos de Bash que te permite adjuntar un trabajo interactivo. Si tu trabajo se inicia correctamente, es posible que la creación de este archivo tarde unos minutos.

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 de ayuda sencillo 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 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. Configura la variable de entorno CLUSTER_TYPE según el siguiente lineamiento:

    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 de 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

Lanzamiento 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 los 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 quiere 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 entre 2 y 5 minutos, y deberías ver que se creó <job_id>/attach.sh.

  • Para supervisar el progreso de la verificación del lanzamiento, consulta <job_id>/nemo_rl_output.log para ver el progreso de la secuencia de comandos de lanzamiento y <job_id>/ray_logs/ para ver el progreso del lanzamiento de los trabajadores y el cabezal de rayos.

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

Ejecutar una carga de trabajo prediseñada verifica el estado operativo del clúster. El siguiente paso es ejecutar tu propia aplicación de entrenamiento personalizada.