Exécuter des charges de travail prédéfinies

Si vous êtes intéressé par Vertex AI Managed Training, contactez votre représentant commercial pour y accéder.

Ce guide vous explique comment utiliser l'écosystème NVIDIA NeMo sur un cluster d'entraînement géré pour développer des modèles d'IA générative de bout en bout. Il fournit des instructions détaillées pour les workflows distincts, mais associés, suivants, chacun étant traité dans sa propre section :

  • NVIDIA NeMo : pour développer des modèles de base, suivez ces instructions pour effectuer un pré-entraînement à grande échelle, un pré-entraînement continu (CPT) et un affinage supervisé (SFT).
  • NVIDIA NeMo-RL : pour l'alignement des modèles et le réglage des préférences, utilisez cette section pour appliquer des techniques avancées comme l'apprentissage par renforcement (RL) afin d'aligner votre modèle sur les instructions et les préférences humaines.

Que vous créiez un modèle de toutes pièces ou que vous en affiniez un existant, ce document vous guide dans la configuration de votre environnement, la gestion des jobs conteneurisés et le lancement des scripts d'entraînement sur le cluster.

NVIDIA NeMo

Le framework NVIDIA NeMo est une plate-forme de bout en bout permettant de créer, de personnaliser et de déployer des modèles d'IA générative. Cette section du guide s'adresse spécifiquement aux développeurs et aux chercheurs qui se concentrent sur les étapes fondamentales du développement de modèles. Il fournit des instructions détaillées pour utiliser NeMo afin d'effectuer un pré-entraînement à grande échelle, un pré-entraînement continu (CPT) et un affinage supervisé (SFT) sur un cluster d'entraînement géré.

Ce guide sur les clusters d'entraînement fournit le workflow complet pour exécuter une tâche d'entraînement avec le framework NeMo. Le processus se divise en deux parties principales : la configuration initiale unique de votre environnement et les étapes récurrentes pour lancer un job.

Configurer votre environnement

Avant de lancer un job, vous devez préparer votre environnement en vous assurant de disposer d'une image de conteneur et des scripts d'entraînement nécessaires.

Préparer une image de conteneur

Vous avez deux options pour l'image de conteneur : utiliser une image prédéfinie (recommandé) ou en créer une personnalisée.

Utiliser une image prédéfinie (recommandé)

Les images de conteneur prédéfinies sont fournies au format .squashfs. Copiez l'image appropriée pour votre région dans votre répertoire de travail.

# Example for the US region
gcloud storage cp gs://vmds-containers-us/nemo_squashfs/nemo-20250721.sqsh .
Créer un conteneur personnalisé (avancé)

Suivez ces étapes uniquement si les conteneurs prédéfinis ne répondent pas à vos besoins. Cette procédure vous guide tout au long de la conversion d'une image de conteneur personnalisée au format .squashfs à l'aide de enroot.

Étape 1 : S'authentifier avec Google Cloud.

Utilisez les commandes suivantes pour vous assurer que votre compte utilisateur Google Cloud et le registre Docker où votre image est hébergée sont authentifiés :

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

Étape 2 : Créez le script de conversion.

Créez un fichier nommé enroot-convert.sh et ajoutez-y le contenu du script suivant. Avant d'exécuter ce script, vous devez mettre à jour les variables REMOTE_IMG et LOCAL_IMG pour qu'elles pointent vers votre image de conteneur et le chemin d'accès de sortie de votre choix.

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

Étape 3 : Exécutez le script et vérifiez le résultat.

Exécutez le script sur le nœud de connexion Slurm.

sbatch -N 1 enroot-convert.sh

Une fois le job terminé, recherchez les journaux de conversion dans un fichier nommé slurm-<JOB_ID>.out et l'image de conteneur finale au chemin d'accès que vous avez spécifié pour LOCAL_IMG.

Télécharger les recettes d'entraînement

Les recettes d'entraînement sont stockées dans un dépôt googlesource.com privé. Pour y accéder avec la ligne de commande Git, vous devez d'abord générer des identifiants d'authentification.

  1. Générez des identifiants d'authentification.

    Accédez à l'URL suivante et suivez les instructions à l'écran. Cela configure votre environnement local pour l'authentification auprès du dépôt. https://www.googlesource.com/new-password

  2. Clonez le dépôt.

    Une fois les identifiants authentifiés, exécutez la commande suivante pour télécharger les recettes.

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

Lancer un job d'entraînement

Une fois votre environnement configuré, vous pouvez lancer un job d'entraînement.

Étape 1 : Définissez des variables d'environnement

Les variables d'environnement suivantes peuvent être requises pour votre tâche :

  • Le HF_TOKEN est nécessaire pour télécharger des modèles et des ensembles de données depuis Hugging Face.
  • Le WANDB_API_KEY est nécessaire pour utiliser Weights & Biases afin d'analyser les tests.
export HF_TOKEN=YOUR_HF_TOKEN
export WANDB_API_KEY=YOUR_WANDB_API_KEY

Étape 2 : Exécutez le script de lancement

Accédez à votre répertoire de travail et exécutez le script run.py pour démarrer un job. Cet exemple lance un job d'entraînement de démonstration avec 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

Paramètres de lancement

  • --slurm-type est défini en fonction du type de cluster (par exemple, hcc-a3m, hcc-a3u, hcc-a4).
  • --partition doit être défini sur une partition disponible. Vous pouvez vérifier les noms de partition avec la commande sinfo.
  • Le script run.py installe automatiquement plusieurs répertoires dans le conteneur Docker, y compris --log-dir, --cache-dir et --data-dir, s'ils sont définis.

Surveiller l'état et les journaux des jobs

Une fois la tâche lancée, un bloc d'état s'affiche :

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

Les journaux d'exécution sont écrits dans le chemin d'accès indiqué dans le champ Local Directory de la sortie d'état. Par exemple, vous pouvez trouver les fichiers journaux à un chemin d'accès semblable à celui-ci :

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

Erreurs courantes et solutions

Cette section décrit les problèmes courants qui peuvent survenir lors de l'exécution des jobs et fournit les étapes recommandées pour les résoudre.

Erreur de partition non valide

Par défaut, les jobs tentent de se lancer sur la partition générale. Si la partition générale n'existe pas ou n'est pas disponible, le job échouera et générera l'erreur suivante :

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

Solution :

Spécifiez une partition disponible à l'aide de l'argument --partition ou -p dans votre commande de lancement. Pour afficher la liste des partitions disponibles, exécutez la commande sinfo sur le nœud de connexion Slurm.

sinfo

Le résultat affiche les noms de partition disponibles, tels que a3u dans cet exemple :

PARTITION AVAIL TIMELIMIT NODES ÉTAT NODELIST
a3u* flèche vers le haut infini 2 idle~ alice-a3u-[2-3]
a3u* flèche vers le haut infini 2 inactif alice-a3u-[0-1]

Erreur de téléchargement du tokenizer

Vous pouvez rencontrer une erreur OS liée à un lien multi-appareils lorsque le script tente de télécharger le tokenizer GPT2 :

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

Solutions :

Pour résoudre ce problème, deux options s'offrent à vous :

  • Option 1 : Exécutez à nouveau le job. Cette erreur est souvent temporaire. Relancer le job en utilisant le même --cache-dir peut résoudre le problème.
  • Option 2 : Téléchargez manuellement les fichiers du tokenizer. Si l'exécution à nouveau du job échoue, procédez comme suit :
    • Téléchargez les deux fichiers suivants :
      • gpt2-vocab.json
      • gpt2-merges.txt
    • Déplacez les fichiers téléchargés dans le sous-répertoire torch/megatron/ de votre répertoire de cache (par exemple, <var>YOUR_CACHE_DIR</var>/torch/megatron/).
    • Renommez les fichiers comme suit :
      • gpt2-vocab.json a été renommé megatron-gpt-345m_vocab.
      • gpt2-merges.txt a été renommé megatron-gpt-345m_merges.

NVIDIA NeMo-RL

Le framework NVIDIA NeMo-RL est conçu pour aligner les grands modèles de langage sur les préférences et les instructions humaines. Cette section vous explique comment utiliser NeMo-RL sur un cluster pour effectuer des tâches d'alignement avancées, y compris le fine-tuning supervisé (SFT), le réglage des préférences (comme l'optimisation directe des préférences ou DPO) et l'apprentissage par renforcement (RL).

Ce guide couvre deux workflows principaux : l'exécution d'un job d'entraînement par lot standard et l'utilisation de l'environnement de développement interactif pour le débogage.

Prérequis

Avant de commencer, créez un cluster en suivant les instructions de la page Créer un cluster ou utilisez un cluster Managed Training existant, si vous en avez un.

Se connecter au nœud de connexion du cluster

Pour vous connecter au nœud de connexion du cluster, recherchez la commande Google Cloud CLI appropriée en accédant à la page Machine virtuelle Google Compute Engine dans la console Google Google Cloud , puis en cliquant sur SSH > Afficher la commande Google Cloud CLI. Cela devrait ressembler à ce qui suit :

ssh $USER_NAME@machine-addr

Exemple :

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

Utiliser l'image Docker prédéfinie

Des fichiers .sqsh convertis sont fournis pour les images de conteneurs prédéfinies. Vous pouvez sélectionner un conteneur pour votre région et le définir directement comme paramètre d'image de conteneur ou le télécharger sur le système de fichiers du cluster.

Pour le définir directement comme paramètre d'image de conteneur, utilisez l'un des chemins d'accès suivants. Notez que vous devez remplacer <region> par votre région spécifique (par exemple, europe, asia, us) :

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

Pour télécharger l'image sur le stockage Lustre du cluster, utilisez la commande suivante :

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

Télécharger le code

Pour accéder à la recette d'entraînement avec CLI de ligne de commande git, accédez à https://www.googlesource.com/new-password. Vous pouvez télécharger la recette à l'aide de la commande suivante :

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

Lancer des tâches

Étape 1 : Définissez les variables d'environnement.

Pour extraire des modèles et des données de Hugging Face, il peut être nécessaire de définir HF_TOKEN. Pour utiliser Weights & Biases pour l'analyse des tests, le WANDB_API_KEY doit être défini. Mettez à jour ces variables dans le fichier suivant :

Fichier à mettre à jour : $HOME/vertex-oss-training/nemo_rl/configs/auth.sh

Si vous ne souhaitez pas utiliser Weights & Biases, définissez logger.wandb_enabled sur False dans votre script de lancement.

Étape 2 : Téléchargez ou copiez le fichier de conteneur dans votre dossier de lancement.

Voici quelques exemples :

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`

Étape 3 : Préparez ou clonez le dépôt NeMo-RL.

Créez un clone du code NeMo-RL s'il n'est pas déjà présent. Notez que vous devrez peut-être utiliser git submodule update --init --recursive si vous avez déjà cloné le dépôt sans l'indicateur --recursive.

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

Étape 4 : Lancez le job d'entraînement.

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

Où :

  • --cluster-type est défini en fonction du type de cluster :
    • A3-Mega : hcc-a3m
    • A3-Ultra : hcc-a3u
    • A4 : hcc-a4
    • A3H : hcc-a3h
  • --partition doit être défini en conséquence, où sinfo peut être utilisé pour vérifier les partitions Slurm.

Une fois votre job démarré, un nouveau répertoire portant le nom de son ID de job SLURM est créé à votre emplacement actuel. Vous y trouverez tous les journaux et points de contrôle associés à ce job. Plus précisément, vous y trouverez les répertoires et fichiers suivants :

  • checkpoints/ : ce répertoire est monté dans le conteneur NeMo-RL et contient tous les points de contrôle de l'entraînement.
  • ray-logs/ : ce répertoire contient les journaux des nœuds principaux et de calcul Ray.
  • nemo_rl_output.log : ce fichier contient les journaux Slurm de votre tâche envoyée.
  • attach.sh (tâches interactives uniquement) : il s'agit d'un script Bash qui vous permet de vous associer à une tâche interactive. Si votre tâche est lancée correctement, la création de ce fichier peut prendre quelques minutes.

Développement avec NeMo-RL

Configuration interactive

Deux options sont disponibles pour un développement interactif rapide avec NeMo-RL.

nemorlinteractive

Il s'agit d'une commande d'assistance simple qui vous permet de choisir un nœud GPU dans le cluster (par exemple, le nœud 5), puis vous redirige vers un conteneur en cours d'exécution pour NeMo-RL dans le nœud sélectionné. Cette commande est utile pour vos workflows à nœud unique.

Pour utiliser nemorlinteractive, vous devez suivre les étapes préalables suivantes :

  1. Fournissez tous les jetons d'authentification souhaités (par exemple, HF et WandB) chargés dans le fichier configs/auth.sh du job.
  2. Définissez la variable d'environnement CLUSTER_TYPE en suivant les consignes ci-dessous :

    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. Importez nemorlinteractive dans votre terminal Bash en sourçant bash_utils.sh :

    source bash_utils.sh
    
  4. Exécutez la commande nemorlinteractive. Exemple :

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

Lancement interactif

Cette option vous permet d'exécuter des charges de travail de manière interactive sur plusieurs nœuds de calcul. Les jobs interactifs sont plus adaptés aux cas d'utilisation de débogage et de validation. Ces charges de travail réservent le nœud indéfiniment jusqu'à ce que le développeur décide que le débogage est terminé et qu'il souhaite libérer les ressources.

Voici les étapes à suivre pour cette option :

Fournissez tous les jetons d'authentification souhaités (par exemple, HF et WandB) chargés dans le job dans le fichier configs/auth.sh.

sbatch -N <num_nodes> launch.sh --cluster_type hcc-a3m --interactive
  • Patientez deux à cinq minutes. L'objet <job_id>/attach.sh devrait être créé.

  • Pour surveiller la progression du lancement, vérifiez <job_id>/nemo_rl_output.log pour voir la progression du script de lancement et <job_id>/ray_logs/ pour voir la progression du lancement des nœuds principaux et des nœuds de calcul Ray.

  • Connectez-vous à la tâche interactive. Ce script vous permet de vous reconnecter même si vous perdez la connexion :

bash <job_id>/attach.sh

Étapes suivantes

L'exécution d'une charge de travail prédéfinie permet de vérifier l'état opérationnel du cluster. L'étape suivante consiste à exécuter votre propre application d'entraînement personnalisé.