Exécuter Megatron-LM sur SLURM

Prérequis

Un accès SLURM avec droits de soumission à une partition GPU.
Megatron-LM cloné sur un système de fichiers visible par tous les nœuds de l'allocation (NFS, Lustre ou similaire). Tous les nœuds doivent accéder aux mêmes chemins pour le code, les données, les checkpoints et la sortie.
uv installé ; exécutez uv sync --extra training --extra dev (ou --extra lts) sur le worktree une fois avant la soumission pour que le .venv soit matérialisé et visible à chaque nœud.

Script sbatch minimal

Enregistrez en tant que run_megatron.slurm dans le worktree :

#!/bin/bash
#SBATCH --job-name=megatron
#SBATCH --account=<SLURM_ACCOUNT>
#SBATCH --partition=<SLURM_PARTITION>
#SBATCH --nodes=<NODES>
#SBATCH --ntasks-per-node=1
#SBATCH --gpus-per-node=<GPUS_PER_NODE>
#SBATCH --time=<HH:MM:SS>
#SBATCH --output=logs/%x-%j.out
#SBATCH --error=logs/%x-%j.err

set -euo pipefail
cd <MEGATRON_WORKTREE>

export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n1)
export MASTER_PORT=${MASTER_PORT:-29500}
export NNODES=${SLURM_NNODES}
export GPUS_PER_NODE=<GPUS_PER_NODE>
export WORLD_SIZE=$((NNODES * GPUS_PER_NODE))

# Set CUDA_DEVICE_MAX_CONNECTIONS only when your configuration requires it
# (see the section below). Example for pre-Blackwell with TP>1 or CP>1
# (non-FSDP):
#   export CUDA_DEVICE_MAX_CONNECTIONS=1

srun --ntasks=${NNODES} --ntasks-per-node=1 bash -c '
  # NODE_RANK comes from SLURM_NODEID with one task per node.
  NODE_RANK=${SLURM_NODEID}
  uv run python -m torch.distributed.run \
    --nnodes='"${NNODES}"' \
    --nproc-per-node='"${GPUS_PER_NODE}"' \
    --node-rank=${NODE_RANK} \
    --master-addr='"${MASTER_ADDR}"' \
    --master-port='"${MASTER_PORT}"' \
    pretrain_gpt.py \
      <MEGATRON_ARGS>
'

Soumettez :

mkdir -p logs && JOB_ID=$(sbatch --parsable run_megatron.slurm)
echo "Submitted ${JOB_ID}"

Règles multi-nœuds

Soumettez depuis le worktree que vous avez l'intention d'exécuter, ou placez-vous-y dans le script. Tous les nœuds doivent accéder au même chemin sur un système de fichiers partagé (NFS, Lustre ou similaire) — les chemins locaux aux nœuds ne seront pas visibles aux autres ranks.
Utilisez un groupe de workers torchrun sur tous les nœuds ; ne lancez pas des jobs indépendants à nœud unique.
--nproc-per-node doit être égal au nombre de GPUs visibles par nœud.
Écrivez les checkpoints, les données tensorboard et les logs structurés dans le stockage partagé.

CUDA_DEVICE_MAX_CONNECTIONS

La bonne valeur dépend de votre matériel et du mode de parallélisme. Ne l'exportez pas inconditionnellement :

Pré-Blackwell (Hopper, Ampere) avec TP>1 ou CP>1, non-FSDP : définissez à 1. Le code correspondant utilise une assertion — vous obtiendrez une erreur d'assertion si ce n'est pas 1, pas un deadlock silencieux.
Blackwell : non requis ; le définir n'a aucun effet.
Torch-FSDP2 ou Megatron-FSDP : ne doit PAS être 1. Laissez la variable d'environnement non définie, ou définissez-la à une valeur supérieure à 1.
overlap_moe_expert_parallel_comm activé : définissez à 32.

Définissez-le explicitement dans le script sbatch quand votre configuration l'exige.

Conteneurs

De nombreux sites exécutent Megatron-LM dans un conteneur (enroot/pyxis sur certains clusters, singularity sur d'autres). Si c'est le cas, le .venv géré par uv doit résider sur un chemin visible depuis l'intérieur du conteneur, et l'image de conteneur doit fournir les versions CUDA / NCCL / torch attendues par le repo (voir docker/.ngc_version.dev et .ngc_version.lts). Le squelette ci-dessus reste identique ; enrobez l'invocation srun avec les flags conteneur de votre ordonnanceur (--container-image=…, --container-mounts=…, etc.).

Surveiller et collecter

squeue -j "$JOB_ID" -o "%.10i %.8T %.10M %.6D %R"
sacct -j "$JOB_ID" --format=JobID,State,ExitCode,Elapsed
scancel "$JOB_ID"

Si votre script d'entraînement écrit un artefact de résultat (un fichier de métriques JSON du rank 0, un checkpoint final, etc.), interrogez l'artefact plutôt que d'attendre uniquement l'état squeue. Les résultats utiles apparaissent généralement avant que SLURM marque le job comme terminé, et interroger l'artefact vous permet d'annuler le job dès qu'il arrive au lieu de conserver l'allocation jusqu'au timeout.

Diagnostic des défaillances

Scannez stderr de chaque rank, pas seulement du rank 0. La première pile d'appels Python non-NCCL est généralement la cause racine ; les timeouts NCCL ultérieurs sur d'autres ranks sont des symptômes en aval du premier crash.

Classifiez rapidement :

OOM : enregistrez le rank, la phase (forward / backward / optimizer), la taille du batch, la longueur de séquence, le parallélisme (TP/DP/CP/PP) et la mémoire de pointe avant d'ajuster.
Erreur de forme / divisibilité : vérifiez WORLD_SIZE = TP × DP × CP × PP et la divisibilité du nombre de têtes (num_attention_heads % TP == 0).
Erreur d'import : mauvais worktree, uv sync manquant ou PYTHONPATH obsolète. Confirmez cd <MEGATRON_WORKTREE> avant le lancement.
Défaillance NCCL sans pile d'appels Python : vérifiez l'allocation, la joignabilité du port, la résolution de MASTER_ADDR et la cohérence des commandes entre les ranks.

Pièges courants

Oublier uv sync avant la première soumission. Si le venv est manquant, chaque job le reconstruit depuis l'intérieur de srun, coûtant des minutes par job.
Écrire les logs dans un chemin local au nœud qui disparaît à la sortie du job. Écrivez toujours sur le système de fichiers partagé.
Définir CUDA_DEVICE_MAX_CONNECTIONS=1 aveuglément. La bonne valeur dépend du matériel et du mode de parallélisme (voir la section dédiée ci-dessus). La définir à 1 avec FSDP cause un problème différent ; sur Blackwell elle n'a aucun effet ; sur pré-Blackwell avec TP>1 ou CP>1 (non-FSDP) le code utilise une assertion, ce n'est pas un deadlock.
Exécuter torchrun nu au lieu de uv run python -m torch.distributed.run. torchrun nu peut être dispatché via un interpréteur Python qui ne voit pas les packages du venv, selon la façon dont le venv est configuré.

run-on-slurm