Ejecutar cargas de trabajo prediseñadas

Si le interesa Vertex AI Managed Training, póngase en contacto con su representante de ventas para obtener acceso.

En esta guía se explica cómo usar el ecosistema NVIDIA NeMo en un clúster de entrenamiento gestionado para desarrollar modelos de IA generativa de principio a fin. Proporciona instrucciones detalladas para los siguientes flujos de trabajo, que son distintos pero están relacionados entre sí. Cada uno de ellos se explica en su propia sección:

  • NVIDIA NeMo: para el desarrollo de modelos fundacionales, sigue estas instrucciones para llevar a cabo el preentrenamiento a gran escala, el preentrenamiento continuo (CPT) y el ajuste fino supervisado (SFT).
  • NVIDIA NeMo-RL: para la alineación de modelos y el ajuste de preferencias, usa esta sección para aplicar técnicas avanzadas como el aprendizaje por refuerzo (RL) y alinear tu modelo con las instrucciones y preferencias de los usuarios.

Tanto si vas a crear un modelo desde cero como si quieres perfeccionar uno que ya tengas, este documento te explica cómo configurar tu entorno, gestionar trabajos en contenedores y lanzar secuencias de comandos de entrenamiento en el clúster.

NVIDIA NeMo

El framework NVIDIA NeMo es una plataforma integral para crear, personalizar y desplegar modelos de IA generativa. Esta sección de la guía está dirigida específicamente a desarrolladores e investigadores que se centran en las fases iniciales del desarrollo de modelos. En ella se proporcionan instrucciones detalladas para usar NeMo y llevar a cabo el preentrenamiento a gran escala, el preentrenamiento continuo (CPT) y el ajuste fino supervisado (SFT) en un clúster de entrenamiento gestionado.

En esta guía de clústeres de entrenamiento se describe el flujo de trabajo completo para ejecutar un trabajo de entrenamiento con el framework 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.

Configurar un entorno

Antes de iniciar un trabajo, debe preparar su entorno. Para ello, asegúrese de que tiene una imagen de contenedor y las secuencias de comandos de entrenamiento necesarias.

Preparar una imagen de contenedor

Tienes dos opciones para la imagen de contenedor: usar una imagen prediseñada (opción recomendada) o crear una personalizada.

Usar una imagen prediseñada (recomendado)

Las imágenes de contenedor prediseñadas 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/nemo_squashfs/nemo-20250721.sqsh .
Crear un contenedor personalizado (opción avanzada)

Siga estos pasos solo si los contenedores precompilados no satisfacen sus 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 tanto tu Google Cloud cuenta de usuario como 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 añade el siguiente contenido de la secuencia de comandos. Antes de ejecutar esta secuencia de comandos, debe actualizar las variables REMOTE_IMG y LOCAL_IMG para que apunten a la imagen de su contenedor y a la ruta de salida que haya elegido.

#!/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 inicio de sesión de Slurm.

sbatch -N 1 enroot-convert.sh

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

Descargar 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. Generar credenciales de autenticación.

    Visita la siguiente URL y sigue las instrucciones que aparecen en pantalla. De esta forma, se configura tu entorno local para autenticarte en el repositorio. https://www.googlesource.com/new-password

  2. Clona el repositorio.

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

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

Iniciar una tarea de entrenamiento

Una vez que hayas configurado tu entorno, podrás iniciar un trabajo de entrenamiento.

Paso 1: Define las variables de entorno

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

  • La HF_TOKEN es necesaria para descargar modelos y conjuntos de datos de Hugging Face.
  • El WANDB_API_KEY es necesario para usar Weights & Biases para analizar experimentos.
export HF_TOKEN=YOUR_HF_TOKEN
export WANDB_API_KEY=YOUR_WANDB_API_KEY

Paso 2: Ejecuta la secuencia de comandos de inicio

Ve a tu directorio de trabajo y ejecuta la secuencia de comandos run.py para iniciar un trabajo. En este ejemplo se inicia una tarea 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>/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 lanzamiento

  • --slurm-type se define en función del tipo de clúster (por ejemplo, hcc-a3m, hcc-a3u o hcc-a4).
  • --partition debe asignarse a una partición disponible. Puedes consultar los nombres de las particiones con el comando sinfo.
  • La secuencia de comandos run.py monta automáticamente varios directorios en el contenedor de Docker, incluidos --log-dir, --cache-dir y --data-dir, si están definidos.

Monitorizar 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 que se muestra en el campo Local Directory de la salida de estado. Por ejemplo, puedes encontrar los archivos de registro en una ruta similar a esta:

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

Errores habituales y soluciones

En esta sección se describen los problemas habituales que pueden surgir durante la ejecución de los trabajos y se indican los pasos recomendados para resolverlos.

Error de partición no válida

De forma predeterminada, las tareas intentan iniciarse en la partición general. Si la partición general no existe o no está disponible, la tarea fallará y se 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 inicio de sesión de Slurm.

sinfo

El resultado muestra los nombres de las particiones disponibles, como a3u en este ejemplo:

PARTITION AVAIL TIMELIMIT NODES ESTADO NODELIST
a3u* arriba infinito 2 Inactivo~ alice-a3u-[2-3]
a3u* arriba infinito 2 inactivo alice-a3u-[0-1]

Error al descargar el tokenizador

Puede que se produzca un error OSError relacionado con un enlace 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 solucionar este problema:

  • Opción 1: Vuelve a ejecutar el trabajo. Este error suele ser temporal. Puede que el problema se solucione si vuelves a ejecutar el trabajo con el mismo --cache-dir.
  • Opción 2: Descarga manualmente los archivos del tokenizador. Si vuelve a ejecutar el trabajo y falla, siga estos pasos:
    • Descarga los dos archivos siguientes:
      • gpt2-vocab.json
      • gpt2-merges.txt
    • Mueve los archivos descargados al subdirectorio torch/megatron/ 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 NVIDIA NeMo-RL se ha diseñado para alinear los modelos de lenguaje extensos con las preferencias e instrucciones humanas. En esta sección se explica cómo usar NeMo-RL en un clúster para llevar a cabo 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 esta guía se describen dos flujos de trabajo principales: ejecutar una tarea de entrenamiento por lotes estándar y usar el entorno de desarrollo interactivo para depurar.

Requisitos previos

Antes de empezar, crea un clúster siguiendo las instrucciones de la página Crear clúster o usa un clúster de entrenamiento gestionado que ya tengas.

Conectarse al nodo de inicio de sesión del clúster

Para conectarte al nodo de inicio de sesión del clúster, busca el comando correcto de la CLI de Google Cloud. Para ello, ve a la página Máquina virtual de Google Compute Engine en la consola de Google Cloud Google Cloud y haz clic en SSH > Ver comando de la CLI de Google Cloud. Tendrá 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

Usar la imagen de Docker prediseñada

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

Para definirlo directamente como parámetro de imagen del contenedor, utilice una de las siguientes rutas. Ten en cuenta que debes sustituir <region> por tu región específica (por ejemplo, europe, asia o us):

/gcs/vmds-containers-<region>/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>/nemo_rl_squashfs/nemo_rl-h20250923.sqsh DESTINATION

Descargar código

Para acceder a la receta de formación con la interfaz de línea de comandos 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 tareas

Paso 1: Define las variables de entorno.

Para extraer modelos y datos de Hugging Face, puede que sea necesario definir el HF_TOKEN. Para usar Weights & Biases en el análisis de experimentos, debe definir el WANDB_API_KEY. Actualiza estas variables en el siguiente archivo:

Archivo que se va a actualizar: $HOME/vertex-oss-training/nemo_rl/configs/auth.sh

Si no quieres usar Weights & Biases, define logger.wandb_enabled como False en tu secuencia de comandos de lanzamiento.

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

A continuación, te ofrecemos 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 lo has hecho. Ten en cuenta que es posible que tengas que usar git submodule update --init --recursive si ya has clonado el repositorio sin la marca --recursive.

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

Paso 4: Inicia la tarea de entrenamiento.

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

Donde:

  • --cluster-type se define en función del tipo de clúster:
    • A3-Mega: hcc-a3m
    • A3 Ultra: hcc-a3u
    • A4: hcc-a4
    • A3H: hcc-a3h
  • --partition debe configurarse en consecuencia, donde sinfo se puede usar para comprobar las particiones de Slurm.

Una vez que se inicie el trabajo, se creará un directorio con el nombre del ID de trabajo de SLURM en tu ubicación actual. En él encontrarás todos los registros y puntos de control de este trabajo. Más concretamente, en ella encontrarás los siguientes directorios y archivos:

  • checkpoints/ → Este directorio se monta en el contenedor de NeMo-RL y contiene todos los puntos de control del entrenamiento.
  • ray-logs/: este directorio contiene los registros de ray head y ray workers.
  • nemo_rl_output.log → Este archivo contiene los registros de Slurm del trabajo que has enviado.
  • attach.sh (Solo trabajos interactivos) → Se trata de una secuencia de comandos bash que te permite adjuntar un trabajo interactivo. Si el trabajo se inicia correctamente, este archivo puede tardar un par de minutos en crearse.

Desarrollo con NeMo-RL

Configuración interactiva

Hay dos opciones disponibles para desarrollar de forma rápida e interactiva con NeMo-RL.

nemorlinteractive

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

Para usar nemorlinteractive, debes seguir estos pasos:

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

  2. Define la variable de entorno CLUSTER_TYPE según las siguientes directrices:

    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 bash obteniendo el 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 computación. Los trabajos interactivos son los más adecuados para depurar y verificar casos prácticos. Estas cargas de trabajo reservan el nodo indefinidamente hasta que el desarrollador decide que la depuración ha finalizado y quiere liberar los recursos.

Estos son los pasos que debes seguir para usar esta opción:

Proporciona todos los tokens de autenticación que quieras (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 para que se cree <job_id>/attach.sh.

  • Para monitorizar el progreso de la comprobación de 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 encabezados y los trabajadores de Ray.

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

bash <job_id>/attach.sh

Siguientes pasos

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