Slurm Multi-Nœud

Convertir les commandes uv run python -m torch.distributed.run simple-nœud en scripts sbatch Slurm multi-nœud avec support de conteneurs Enroot, et déboguer les défaillances multi-nœud courantes.

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

Préférer srun-native — plus simple, évite les problèmes d'échappement shell avec TRAIN_CMD. Megatron Bridge auto-dérive RANK, WORLD_SIZE, LOCAL_RANK, MASTER_ADDR, MASTER_PORT à partir des variables d'env Slurm (SLURM_PROCID, SLURM_NTASKS, SLURM_LOCALID, SLURM_NODELIST) via les helpers common_utils.py appelés lors de l'init distribué initialize.py, donc tu n'as 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 SFT/PEFT multi-nœud. Le défaut (/root/.cache/nemo) est local au conteneur et non partagé entre nœuds. Sans cela, les fichiers de données packed-sequence préparés sur le nœud 0 sont invisibles aux autres nœuds, causant une TypeError: 'NoneType' object is not an iterator.

Répertoire de logs

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

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

Slurm crée 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 crée 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

Deux phases de srun : d'abord un srun simple-processus pour peupler 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 simple-processus pour construire/peupler 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 puisque le cache est 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 seul nœud/processus, construisant toutes les wheels dans le cache partagé sur Lustre
• Le uv sync de la phase 2 est un no-op rapide (tout est en cache) — safe d'exécuter sur tous les ranks sans gardes sleep
• initialize.py + common_utils.py définissent automatiquement RANK, WORLD_SIZE, LOCAL_RANK, MASTER_ADDR, MASTER_PORT à partir des variables d'env 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/glm/glm_45v/slurm_sft.sh, examples/models/minimax/minimax_m2/slurm_conversion.sh

Approche uv run torch.distributed (Legacy)

À utiliser 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. Ajouter les en-têtes SBATCH

#SBATCH --job-name=<model>-<framework>
#SBATCH --nodes=<NNODES>
#SBATCH --ntasks-per-node=1          # TOUJOURS 1 — torchrun gère la création 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 crée 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. Convertir en multi-nœud

Remplacer simple-nœud :

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

Par 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 auto-dérivés à partir des variables d'env Slurm par initialize.py / common_utils.py — pas besoin de les définir.

3. Envelopper en TRAIN_CMD + deux phases srun

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

Les exports d'env 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. Lancer (deux phases)

# Phase 1 : uv sync simple-processus pour construire/peupler 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) Ajouter un pied de page d'extraction de losses

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), toujours suivre ces 3 étapes :

Étape 1 : Allouer le nœud

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

Étape 2 : Lancer 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 : Configurer 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écuter les commandes avec uv run (utilise l'env virtuel synced) :

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 allocation interactive :

Déboguer les défaillances multi-nœud

Diagnostic rapide

Vérifier le log pour ces motifs (dans cet ordre) :

# 1. Trouver l'erreur réelle (filtrer 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érifier quel rank a crashé en premier
grep -a 'Failures:' -A 20 job.log | head -25

# 3. Vérifier les timeouts 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érifier le code de sortie : 1 = erreur Python, 9 = OOM tué, 143 = SIGTERM (timeout ou cascade)
2. Trouver la première défaillance : Quel rank/nœud a crashé en premier ? Les autres reçoivent SIGTERM (143) en cascade
3. grep l'erreur réelle : Filtrer les UserWarnings, dumps de frames NCCL
4. Vérifier le rank 0 spécifiquement : La plupart des erreurs save/export se produisent sur le rank 0
5. Vérifier le dimensionnement EP : Pour les modèles MoE, s'assurer que num_experts / EP rentre en mémoire GPU avec marge
6. Essayer interactive d'abord : Utiliser 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 ranks 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érifier dans cet ordre) :

Le problème save_artifacts : Quand trust_remote_code=True, rank 0 exécute save_artifacts() (télécharge tokenizer, config, code de modélisation custom) tandis que tous les autres ranks sautent directement à dist.barrier(). Si save_artifacts est lent ou crash, les autres ranks timeout.

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

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

OOM pour modèles MoE

Symptôme : torch.OutOfMemoryError: CUDA out of memory lors du chargement de modèle ou forward pass.

Insight clé : TP ne réduit PAS la mémoire d'expert. 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 de thumb :

• Roundtrip (weight-only) : peut utiliser plus d'experts par GPU (~60GB params de modèle OK)
• Inférence (forward pass + KV cache) : nécessite de la marge (~40GB params de modèle max)
• Training (activations + optimizer) : nécessite encore plus de marge (~30GB params de 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 dans roundtrip

Symptôme : Roundtrip se complète mais affiche ❌ pour tous les poids d'expert et lève une ValueError: Weight mismatch detected.

Cause : Les poids HF originaux sont FP8, Megatron stocke en BF16. Les poids exportés sont BF16. La comparaison contre 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 : Script quitte avec « must be launched with torchrun ».

Cause : Les scripts vérifient os.environ.get("WORLD_SIZE") que torchrun définit mais srun ne définit 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 Bridge common_utils.py (appelés par initialize.py) peuplent les variables d'env à partir de Slurm :

if "RANK" not in os.environ:
    os.environ["RANK"] = str(get_rank_safe())          # uses SLURM_PROCID
if "WORLD_SIZE" not in os.environ:
    os.environ["WORLD_SIZE"] = str(get_world_size_safe())  # uses SLURM_NTASKS
if "MASTER_ADDR" not in os.environ:
    os.environ["MASTER_ADDR"] = get_master_addr_safe()     # parses SLURM_NODELIST
if "MASTER_PORT" not in os.environ:
    os.environ["MASTER_PORT"] = str(get_master_port_safe()) # derives from SLURM_JOB_ID

Pièges clés

1. Deux phases srun pour uv sync : Exécuter d'abord un srun simple-processus pour 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 double-quotée, échapper les $ internes pour les variables Slurm qui doivent s'étendre à l'exécution (pas au moment sbatch) :

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

4. Bridge rm -rf nemo_experiments : Ajouter avant training pour éviter une auto-reprise de checkpoint stale.

5. MLM nécessite 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 : GPUs totaux = NNODES * 8. Doit satisfaire : TP * PP * EP * DP >= GPUs_totaux où DP = GPUs_totaux / (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 sequences packé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 packées soient partagées. Sans cela, les ranks sur d'autres nœuds échouent avec TypeError: 'NoneType' object is not an iterator.

Modèles complets et corps de commandes

Pour scaffolding sbatch copiable et corps TRAIN_CMD Bridge/MLM-spécifiques, lire references/templates.md.