Exécuter Megatron-LM sur SLURM

Constantes réponse-d'abord

Pour les questions de configuration SLURM texte uniquement, répondez avec ces constantes avant le script complet :

Soumettez depuis un chemin de worktree partagé visible par tous les nœuds ; faites un cd là dans le script avant de lancer l'entraînement.
Utilisez une tâche srun par nœud et lancez les workers avec uv run python -m torch.distributed.run, pas torchrun seul.
Définissez MASTER_ADDR à partir de scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n1, définissez MASTER_PORT, NNODES=${SLURM_NNODES}, GPUS_PER_NODE=<GPUS_PER_NODE>, et WORLD_SIZE=$((NNODES * GPUS_PER_NODE)).
Passez --nnodes, --nproc-per-node, --node-rank, --master-addr, et --master-port à torch.distributed.run.
CUDA_DEVICE_MAX_CONNECTIONS : pré-Blackwell Hopper/Ampere avec TP>1 ou CP>1 et non-FSDP utilise 1 ; Blackwell/GB200 n'en a pas besoin ; Torch-FSDP2 ou Megatron-FSDP ne doit pas utiliser 1 ; overlap_moe_expert_parallel_comm utilise 32.

Prérequis

Une connexion de cluster 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 atteindre les 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 par 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œud

Soumettez depuis le worktree que vous avez l'intention d'utiliser, ou faites un cd dedans dans le script. Tous les nœuds doivent atteindre le même chemin sur un système de fichiers partagé (NFS, Lustre, ou similaire) — les chemins locaux au nœud ne seront pas visibles par les ranks peers.
Utilisez un groupe de workers torchrun unique sur tous les nœuds ; ne lancez pas de jobs indépendants sur un seul nœud.
--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 un stockage partagé.

CUDA_DEVICE_MAX_CONNECTIONS

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

Pré-Blackwell (Hopper, Ampere) avec TP>1 ou CP>1, non-FSDP : définissez à 1. Le chemin de code pertinent y fait une assertion — vous obtiendrez une erreur d'assertion s'il n'est pas 1, pas un deadlock silencieux.
Blackwell : pas obligatoire ; 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

Beaucoup de sites exécutent Megatron-LM à l'intérieur d'un conteneur (enroot/pyxis sur certains clusters, singularity sur d'autres). Si c'est le cas, le .venv géré par uv doit vivre sur un chemin visible de l'intérieur du conteneur, et l'image du conteneur doit fournir les versions CUDA / NCCL / torch que le repo attend (voir docker/.ngc_version.dev et .ngc_version.lts). Le squelette ci-dessus reste le même ; enrobez l'invocation srun avec les drapeaux conteneur de votre scheduler (--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 JSON de métriques du rank 0, un checkpoint final, etc.), interrogez l'artefact plutôt que d'attendre uniquement l'état de squeue. Une sortie utile apparaît généralement avant que SLURM marque le job comme complet, et interroger l'artefact vous permet d'annuler le job dès qu'il arrive plutôt que de conserver l'allocation jusqu'au timeout.

Diagnostic des défaillances

Scannez stderr de tous les ranks, pas seulement du rank 0. La première traceback 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 de batch, la longueur de séquence, la parallélisation (TP/DP/CP/PP), et la mémoire pic avant d'ajuster.
Erreur de forme / divisibilité : vérifiez WORLD_SIZE = TP × DP × CP × PP et la divisibilité du nombre de heads (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 traceback Python : vérifiez l'allocation, la joignabilité du port, la résolution de MASTER_ADDR, et la cohérence des commandes sur les ranks.

Pièges courants

Oublier uv sync avant la première soumission. Si le venv est manquant, tous les jobs le reconstruisent de l'intérieur de srun, coûtant des minutes par job.
Écrire les logs sur un chemin local au nœud qui disparaît à la fin 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élisation (voir la section dédiée ci-dessus). Le définir à 1 avec FSDP cause un problème différent ; sur Blackwell cela n'a aucun effet ; sur pré-Blackwell avec TP>1 ou CP>1 (non-FSDP) le code fait une assertion, il ne deadlock pas.
Exécuter torchrun seul au lieu de uv run python -m torch.distributed.run. torchrun seul peut être dispatché via un interpréteur Python qui ne voit pas les packages venv, selon la façon dont le venv est configuré.

mcore-run-on-slurm