Multi-Node Slurm

Convertissez les commandes uv run python -m torch.distributed.run sur nœud unique en scripts sbatch multi-nœud Slurm avec support de conteneur Enroot, et déboguez les défaillances multi-nœud courantes.

Deux approches : srun-native vs uv run torch.distributed

Préférez srun-native — plus simple, évite les problèmes d'échappement shell avec TRAIN_CMD. Megatron Bridge dérive automatiquement RANK, WORLD_SIZE, LOCAL_RANK, MASTER_ADDR, MASTER_PORT des variables d'environnement SLURM (SLURM_PROCID, SLURM_NTASKS, SLURM_LOCALID, SLURM_NODELIST) via les helpers de common_utils.py appelés lors de l'initialisation distribuée dans initialize.py, donc vous n'avez jamais besoin de les définir manuellement.

Environnement du cluster

Conteneur

CONTAINER_IMAGE="<PATH_TO_YOUR_CONTAINER>.sqsh"
CONTAINER_MOUNTS="<SHARED_FS>:<SHARED_FS>,<PATH_TO_MEGATRON_BRIDGE>:/opt/Megatron-Bridge,<PATH_TO_DATA>:/opt/data"

Chemins standards

WORKDIR="/opt/Megatron-Bridge"
DATA_PATH="<PATH_TO_PREPROCESSED_DATA>/dclm_01_01_text_document"

Tokens / Caches

export GH_TOKEN=<YOUR_GITHUB_TOKEN>
export HF_TOKEN=<YOUR_HF_TOKEN>
export HF_HOME=<SHARED_FS>/HF_HOME
export UV_CACHE_DIR="<SHARED_FS>/uv_cache"
export NEMO_HOME="<SHARED_FS>/cache/nemo"

Important : NEMO_HOME doit pointer vers un système de fichiers partagé (ex. Lustre) pour les jobs multi-nœud SFT/PEFT. Le défaut (/root/.cache/nemo) est local au conteneur et non partagé entre nœuds. Sans cela, les fichiers de données avec séquences compactées préparés sur le nœud 0 sont invisibles aux autres nœuds, causant TypeError: 'NoneType' object is not an iterator.

Répertoire de journaux

<SHARED_FS>/logs/<job_name>_<suffix>

Approche srun-native (Préférée)

Slurm génère tous les processus directement. Pas de torch.distributed.run, pas d'échappement TRAIN_CMD.

En-têtes SBATCH

#SBATCH --job-name=<model>-<task>
#SBATCH --nodes=<NNODES>
#SBATCH --ntasks-per-node=8          # Slurm génère 8 tâches par nœud
#SBATCH --gpus-per-node=8
#SBATCH --time=00:30:00
#SBATCH --account=<YOUR_ACCOUNT>
#SBATCH --partition=batch
#SBATCH --output=<SHARED_FS>/logs/<job_name>_%j.log
#SBATCH --exclusive

Construction et lancement

Srun en deux phases : d'abord un srun mono-processus pour remplir le cache uv, puis le srun multi-nœud complet.

# Exports d'env au niveau sbatch (avant srun)
export TORCH_NCCL_AVOID_RECORD_STREAMS=1
export NCCL_NVLS_ENABLE=0

# Phase 1 : uv sync mono-processus pour compiler/remplir le cache partagé
srun --mpi=pmix -N 1 --ntasks=1 \
  --container-image="$CONTAINER_IMAGE" \
  --container-mounts="$CONTAINER_MOUNTS" \
  --no-container-mount-home \
  bash -c "cd $WORKDIR && uv sync"

# Phase 2 : Exécution multi-nœud complète (uv sync est un no-op rapide avec cache chaud)
srun --mpi=pmix \
  --container-image="$CONTAINER_IMAGE" \
  --container-mounts="$CONTAINER_MOUNTS" \
  --no-container-mount-home \
  bash -c "cd $WORKDIR && uv sync && uv run --no-sync python <script.py> <args>"

Points clés de srun-native

• La phase 1 exécute uv sync une fois sur un nœud/processus unique, construisant tous les wheels dans le cache partagé sur Lustre
• Le uv sync de la phase 2 est un no-op rapide (tout est en cache) — sûr d'exécuter sur tous les rangs sans gardes de sommeil
• initialize.py + common_utils.py définissent automatiquement RANK, WORLD_SIZE, LOCAL_RANK, MASTER_ADDR, MASTER_PORT à partir des variables d'environnement SLURM
• Les variables d'env comme HF_TOKEN, HF_HOME, UV_CACHE_DIR exportées au niveau sbatch sont héritées par les tâches srun
• Référence : examples/models/vlm/glm_45v/slurm_sft.sh, examples/models/minimax_m2/slurm_conversion.sh

Approche uv run torch.distributed (Legacy)

Utilisez quand le script nécessite torch.distributed.run (ex. MLM pretrain_gpt.py) ou quand Bridge's initialize.py n'est pas dans le chemin d'appel.

1. Ajoutez les en-têtes SBATCH

#SBATCH --job-name=<model>-<framework>
#SBATCH --nodes=<NNODES>
#SBATCH --ntasks-per-node=1          # TOUJOURS 1 — torchrun gère la génération par nœud
#SBATCH --gpus-per-node=8
#SBATCH --time=00:30:00
#SBATCH --account=<YOUR_ACCOUNT>
#SBATCH --partition=batch
#SBATCH --output=<SHARED_FS>/logs/<job_name>_%j.log
#SBATCH --exclusive

Critique : --ntasks-per-node=1, PAS 8. uv run python -m torch.distributed.run --nproc_per_node=8 génère 8 processus par nœud. Utiliser ntasks-per-node=8 cause des collisions de port EADDRINUSE (8 tâches x 8 procs = 64 par nœud).

2. Convertissez en multi-nœud

Remplacez mono-nœud :

uv run python -m torch.distributed.run --nproc_per_node=8 \
  <script> <args>

Avec multi-nœud (à l'intérieur de la chaîne TRAIN_CMD) :

uv run python -m torch.distributed.run \
  --nproc_per_node=8 \
  --nnodes=\${SLURM_JOB_NUM_NODES} \
  --node_rank=\${SLURM_NODEID} \
  <script> <args>

MASTER_ADDR et MASTER_PORT sont dérivés automatiquement des variables d'env SLURM par initialize.py / common_utils.py — pas besoin de les définir.

3. Enveloppez dans TRAIN_CMD + srun en deux phases

Utilisez le même motif en deux phases : d'abord un srun mono-processus pour réchauffer le cache uv, puis l'exécution complète.

Les exports d'environnement vont à l'intérieur de TRAIN_CMD (ils doivent être définis à l'intérieur du conteneur) :

TRAIN_CMD="
export CUDA_DEVICE_MAX_CONNECTIONS=1 && \
export NVTE_ALLOW_NONDETERMINISTIC_ALGO=1 && \
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True && \
export NCCL_NVLS_ENABLE=0 && \
export GH_TOKEN=$GH_TOKEN && \
export HF_TOKEN=$HF_TOKEN && \
export HF_HOME=$HF_HOME && \
export UV_CACHE_DIR=$UV_CACHE_DIR && \
wandb login \$WANDB_API_KEY && \
mkdir -p $LOGDIR && \
cd $WORKDIR && \
uv sync && \
<training command here>
"

4. Lancez (deux phases)

# Phase 1 : uv sync mono-processus pour compiler/remplir le cache partagé
srun --mpi=pmix -N 1 --ntasks=1 \
  --container-image="$CONTAINER_IMAGE" \
  --container-mounts="$CONTAINER_MOUNTS" \
  --no-container-mount-home \
  bash -c "cd $WORKDIR && uv sync"

# Phase 2 : Exécution multi-nœud complète (uv sync dans TRAIN_CMD est un no-op rapide)
srun --mpi=pmix --no-kill \
  --container-image="$CONTAINER_IMAGE" \
  --container-mounts="$CONTAINER_MOUNTS" \
  --no-container-mount-home \
  bash -c "$TRAIN_CMD" 2>&1 | tee "$LOGDIR/<prefix>_${SLURM_JOB_ID}.log"

5. (Optionnel) Ajoutez un pied de page d'extraction de perte

echo "======================================"
echo "Done. Losses:"
echo "======================================"
grep -E "iteration\s+" "$LOGDIR/<prefix>_${SLURM_JOB_ID}.log" | grep -iE "lm loss|reduced_train_loss" | head -25

Allocation GPU interactive (salloc + srun)

Pour les tests ad-hoc (inférence, débogage de conversion), suivez toujours ces 3 étapes :

Étape 1 : Allouez le nœud

salloc --account <YOUR_ACCOUNT> -N 1 \
  -J <YOUR_ACCOUNT>-debug \
  -p interactive --gpus-per-node=8 -t 240

Étape 2 : Lancez le shell du conteneur

srun --mpi=pmix --no-kill \
  --container-image $CONTAINER_IMAGE \
  --container-mounts $CONTAINER_MOUNTS \
  --account <YOUR_ACCOUNT> -N 1 \
  -J <YOUR_ACCOUNT>-debug \
  --no-container-mount-home --gpus-per-node=8 \
  -p interactive --pty bash

Étape 3 : Configurez l'environnement à l'intérieur du conteneur

export GH_TOKEN=<YOUR_GITHUB_TOKEN>
wandb login <YOUR_WANDB_KEY>
export HF_TOKEN=<YOUR_HF_TOKEN>
export HF_HOME=<SHARED_FS>/HF_HOME
export UV_CACHE_DIR="<SHARED_FS>/uv_cache"
export NEMO_HOME="<SHARED_FS>/cache/nemo"
uv sync

Puis exécutez les commandes avec uv run (utilise l'env virtuel synchronisé) :

uv run python -m torch.distributed.run --nproc_per_node=8 \
  examples/conversion/hf_to_megatron_generate_text.py \
  --hf_model_path <org>/<model> --prompt "What is AI?" --max_new_tokens 50 --ep 8

Pièges avec l'allocation interactive :

Débogage des défaillances multi-nœud

Diagnostic rapide

Recherchez ces motifs dans le journal (dans l'ordre) :

# 1. Trouvez l'erreur réelle (filtrez le bruit)
grep -a 'Error\|OOM\|CUDA out of memory\|FAILED\|Killed' job.log \
  | grep -v 'UserWarning\|AllocatorConfig\|transformer_engine\|frame\|srun: error'

# 2. Vérifiez quel rang a échoué en premier
grep -a 'Failures:' -A 20 job.log | head -25

# 3. Vérifiez le timeout NCCL
grep -a 'ncclUniqueId\|timeout\|crash on rank 0' job.log | head -5

Checklist de débogage

Quand un job multi-nœud échoue :

1. Vérifiez le code de sortie : 1 = erreur Python, 9 = OOM tué, 143 = SIGTERM (timeout ou cascade)
2. Trouvez la première défaillance : Quelle tâche/nœud a échoué en premier ? Les autres reçoivent SIGTERM (143) en cascade
3. grep l'erreur réelle : Filtrez les UserWarnings et vidages de trames NCCL
4. Vérifiez le rang 0 spécifiquement : La plupart des erreurs de sauvegarde/export se produisent sur le rang 0
5. Vérifiez le dimensionnement EP : Pour les modèles MoE, vérifiez que num_experts / EP s'ajuste dans la mémoire GPU avec marge
6. Essayez interactive d'abord : Utilisez salloc -N 2 -p interactive pour itérer plus vite que la queue sbatch

Timeout NCCL à dist.barrier() — "crash on rank 0"

Symptôme : Tous les rangs sur nœud 2+ affichent :

[rank8] is setting up NCCL communicator and retrieving ncclUniqueId from [0]
... wait timeout after 600000ms
This may indicate a possible application crash on rank 0

Causes racines (vérifiez dans l'ordre) :

Le problème de save_artifacts : Quand trust_remote_code=True, le rang 0 exécute save_artifacts() (télécharge tokenizer, config, code de modélisation personnalisé) tandis que tous les autres rangs passent directement à dist.barrier(). Si save_artifacts est lent ou échoue, les autres rangs timeout.

Correction pour ImportError dans save_artifacts (hf_pretrained/base.py) :

# Changez :
except OSError:
    pass
# En :
except (OSError, ImportError):
    pass

OOM pour les modèles MoE

Symptôme : torch.OutOfMemoryError: CUDA out of memory lors du chargement du modèle ou de la passe avant.

Point clé : TP ne réduit PAS la mémoire des experts. Seul EP divise les experts entre GPUs.

Formule de dimensionnement :

experts_per_gpu = num_experts / EP
expert_memory_gb ≈ experts_per_gpu * expert_params * 2 / 1e9  (bf16)
total_per_gpu ≈ expert_memory_gb + attention_memory_gb + kv_cache_gb

Exemple MiniMax-M2 (256 experts, ~230GB fp8 → ~460GB bf16) :

Règles empiriques :

• Roundtrip (weight-only) : peut utiliser plus d'experts par GPU (~60GB params modèle OK)
• Inférence (passe avant + cache KV) : a besoin de marge (~40GB params modèle max)
• Entraînement (activations + optimiseur) : a besoin de même plus de marge (~30GB params modèle max)

ModuleNotFoundError: No module named 'megatron.core.tensor_parallel'

Cause : Le megatron-core pré-installé du conteneur entre en conflit avec 3rdparty/Megatron-LM local.

Correction : Ajouter uv sync avant d'exécuter :

CMD="if [ \"\$SLURM_LOCALID\" -eq 0 ]; then uv sync; else sleep 10; fi && "
CMD="${CMD}uv run --no-sync python <script> <args>"

Mismatch de poids FP8 en Roundtrip

Symptôme : Roundtrip complète mais affiche ❌ pour tous les poids des experts et lève ValueError: Weight mismatch detected.

Cause : Les poids HF originaux sont FP8, Megatron stocke en BF16. Les poids exportés sont BF16. La comparaison avec l'original FP8 dépasse atol=1e-1.

C'est attendu pour les modèles FP8. La conversion est correcte ; la tolérance de comparaison est insuffisante pour l'écart de précision FP8→BF16.

WORLD_SIZE non défini avec srun

Symptôme : Le script se ferme avec "must be launched with torchrun".

Cause : Les scripts vérifient os.environ.get("WORLD_SIZE") que torchrun définit mais srun ne le fait pas.

Correction : Vérifier aussi SLURM_NTASKS :

if os.environ.get("WORLD_SIZE") is None and os.environ.get("SLURM_NTASKS") is None:
    sys.exit(1)

Les helpers de common_utils.py de Bridge (appelés par initialize.py) remplissent les variables d'env à partir de SLURM :

if "RANK" not in os.environ:
    os.environ["RANK"] = str(get_rank_safe())          # utilise SLURM_PROCID
if "WORLD_SIZE" not in os.environ:
    os.environ["WORLD_SIZE"] = str(get_world_size_safe())  # utilise SLURM_NTASKS
if "MASTER_ADDR" not in os.environ:
    os.environ["MASTER_ADDR"] = get_master_addr_safe()     # analyse SLURM_NODELIST
if "MASTER_PORT" not in os.environ:
    os.environ["MASTER_PORT"] = str(get_master_port_safe()) # dérive de SLURM_JOB_ID

Pièges clés

1. Srun en deux phases pour uv sync : Exécutez d'abord un srun mono-processus pour réchauffer le cache, puis le srun multi-nœud complet. Le second uv sync est un no-op rapide puisque tout est déjà en cache sur le système de fichiers partagé.

2. --no-container-mount-home est un flag srun, PAS une directive #SBATCH.

3. Échappement à l'intérieur de TRAIN_CMD : Puisque TRAIN_CMD est une chaîne entre guillemets doubles, échappez les $ intérieurs pour les variables Slurm qui doivent se développer au runtime (pas au moment sbatch) :

   • \${SLURM_PROCID}, \${SLURM_JOB_NUM_NODES}, \${SLURM_NODEID}
   • Les variables côté hôte comme $GH_TOKEN, $LOGDIR, $WORKDIR se développent au moment sbatch — pas d'échappement nécessaire.

4. Bridge rm -rf nemo_experiments : Ajouter avant l'entraînement pour éviter la reprise automatique de checkpoint obsolète.

5. MLM a besoin de PYTHONPATH : Pour les scripts pretrain_gpt.py, ajouter à l'intérieur de TRAIN_CMD :

   PYTHONPATH=${WORKDIR}/3rdparty/Megatron-LM:\${PYTHONPATH:-} \

6. Heuristique de nombre de nœuds : Total GPUs = NNODES * 8. Doit satisfaire : TP * PP * EP * DP >= total_GPUs où DP = total_GPUs / (TP * PP * EP).

7. NEMO_HOME sur système de fichiers partagé pour SFT multi-nœud : Le cache nemo par défaut (/root/.cache/nemo) est local au conteneur. SFT multi-nœud avec séquences compactées prépare des fichiers .npy sur un nœud qui sont invisibles aux autres. Définir export NEMO_HOME=<SHARED_FS>/cache/nemo pour que les données compactées soient partagées. Sans cela, les rangs sur d'autres nœuds échouent avec TypeError: 'NoneType' object is not an iterator.

Modèle complet

#!/bin/bash
# ==============================================================================
# <MODEL_NAME> <pretrain|sft> — <Framework: MLM | Megatron Bridge>
#
# Défaut : TP<X> PP<Y> EP<Z>, NNODES=<N> (<N*8> GPUs), MBS=<M>, GBS=<G>
#
# Usage:
#   sbatch <script_name>.sh
# ==============================================================================

#SBATCH --job-name=<job-name>
#SBATCH --nodes=<NNODES>
#SBATCH --ntasks-per-node=1
#SBATCH --gpus-per-node=8
#SBATCH --time=00:30:00
#SBATCH --account=<YOUR_ACCOUNT>
#SBATCH --partition=batch
#SBATCH --output=<SHARED_FS>/logs/<job_name>_%j.log
#SBATCH --exclusive

# ── Conteneur ─────────────────────────────────────────────────────────────
CONTAINER_IMAGE="<PATH_TO_YOUR_CONTAINER>.sqsh"
CONTAINER_MOUNTS="<SHARED_FS>:<SHARED_FS>,<PATH_TO_MEGATRON_BRIDGE>:/opt/Megatron-Bridge,<PATH_TO_DATA>:/opt/data"

# ── Chemins ───────────────────────────────────────────────────────────────
WORKDIR="/opt/Megatron-Bridge"
LOGDIR="<SHARED_FS>/logs/<logdir_name>"
DATA_PATH="<PATH_TO_PREPROCESSED_DATA>/dclm_01_01_text_document"

# ── Parallélisme ──────────────────────────────────────────────────────────
TP=1; PP=1; EP=1

# ── Entraînement ──────────────────────────────────────────────────────────
MBS=1; GBS=256
SEQ=4096
SEED=1234
TRAIN_ITERS=20

# ── Tokens / Caches ───────────────────────────────────────────────────────
export GH_TOKEN=<YOUR_GITHUB_TOKEN>
export HF_TOKEN=<YOUR_HF_TOKEN>
export HF_HOME=<SHARED_FS>/HF_HOME
export UV_CACHE_DIR="<SHARED_FS>/uv_cache"
export NEMO_HOME="<SHARED_FS>/cache/nemo"

# ── Construisez la commande d'entraînement ────────────────────────────────
TRAIN_CMD="
export CUDA_DEVICE_MAX_CONNECTIONS=1 && \
export NVTE_ALLOW_NONDETERMINISTIC_ALGO=1 && \
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True && \
export NCCL_NVLS_ENABLE=0 && \
export GH_TOKEN=$GH_TOKEN && \
export HF_TOKEN=$HF_TOKEN && \
export HF_HOME=$HF_HOME && \
export UV_CACHE_DIR=$UV_CACHE_DIR && \
export NEMO_HOME=$NEMO_HOME && \
wandb login \$WANDB_API_KEY && \
mkdir -p $LOGDIR && \
cd $WORKDIR && \
uv sync && \
<TRAINING_COMMAND_HERE>
"

echo \"======================================\"
echo \"<MODEL_NAME> <Framework> Pretrain\"
echo \"Job: \$SLURM_JOB_ID | Nodes: \$SLURM_JOB_NUM_NODES\"
echo \"TP=\$TP PP=\$PP EP=\$EP MBS=\$MBS GBS=\$GBS\"
echo \"======================================\"

# Phase 1 : uv sync mono-processus pour compiler/remplir le cache partagé
srun --mpi=pmix -N 1 --ntasks=1 \
  --container-image="$CONTAINER_IMAGE" \
  --container-mounts="$CONTAINER_MOUNTS" \
  --no-container-mount-home \
  bash -c "cd $WORKDIR && uv sync"

# Phase 2 : Exécution multi-nœud complète (uv sync dans TRAIN_CMD est un no-op rapide)
srun --mpi=pmix --no-kill \
  --container-image="$CONTAINER_IMAGE" \
  --container-mounts="$CONTAINER_MOUNTS" \
  --no-container-mount-home \
  bash -c "$TRAIN_CMD" 2>&1 | tee "$LOGDIR/<prefix>_${SLURM_JOB_ID}.log"

echo ""
echo "======================================"
echo "Done. Losses:"
echo "======================================"
grep -E "iteration\s+" "$LOGDIR/<prefix>_${SLURM_JOB_ID}.log" | grep -iE "lm loss|reduced_train_loss" | head -25

Corps TRAIN_CMD spécifique à Bridge

rm -rf nemo_experiments && \
uv run python -m torch.distributed.run \
  --nproc_per_node=8 \
  --nnodes=\${SLURM_JOB_NUM_NODES} \
  --node_rank=\${SLURM_NODEID} \
  scripts/training/run_recipe.py \
  --recipe <recipe_name> \
  model.tensor_model_parallel_size=$TP \
  model.pipeline_model_parallel_size=$PP \
  ...overrides...

Corps TRAIN_CMD spécifique à MLM

PYTHONPATH=${WORKDIR}/3rdparty/Megatron-LM:\${PYTHONPATH:-} \
uv run python -m torch.distributed.run \
  --nproc_per_node=8 \
  --nnodes=\${SLURM_JOB_NUM_NODES} \
  --node_rank=\${SLURM_NODEID} \
  3rdparty/Megatron-LM/pretrain_gpt.py \
  --tensor-model-parallel-size $TP \
  --pipeline-model-parallel-size $PP \
  ...args...