nemo-mbridge-multi-node-slurm

Multi-Node Slurm

Convert single-node uv run python -m torch.distributed.run commands into multi-node Slurm sbatch scripts with Enroot container support, and debug common multi-node failures.

First Answer Checklist

When converting or debugging Bridge multi-node jobs, answer in this order:

1. Prefer the srun-native launch shape for Bridge scripts that reach initialize.py: #SBATCH --ntasks-per-node=8 and a direct srun ... uv run python <script> ... launch. Do not wrap these jobs in python -m torch.distributed.run.
2. State that Bridge derives RANK, WORLD_SIZE, LOCAL_RANK, MASTER_ADDR, and MASTER_PORT from SLURM variables during initialize.py distributed init.
3. Require shared paths and matching container mounts for the repo, data, logs, HF_HOME, UV_CACHE_DIR, and NEMO_HOME.
4. For NCCL timeout reports, do these first-log checks before speculating:
   • grep for real errors while filtering warning/frame noise
   • inspect Failures: to find the first failed rank and node
   • grep for ncclUniqueId, timeout, or crash on rank 0

Two Approaches: srun-native vs uv run torch.distributed

Prefer srun-native — simpler, avoids shell escaping issues with TRAIN_CMD. Megatron Bridge auto-derives RANK, WORLD_SIZE, LOCAL_RANK, MASTER_ADDR, MASTER_PORT from SLURM env vars (SLURM_PROCID, SLURM_NTASKS, SLURM_LOCALID, SLURM_NODELIST) via common_utils.py helpers called during initialize.py distributed init, so you never need to set them manually.

Cluster Environment

Use a shared filesystem for the repository, data, logs, HF_HOME, UV_CACHE_DIR, and NEMO_HOME. NEMO_HOME must not use the container-local default (/root/.cache/nemo) for multi-node SFT/PEFT jobs, because packed-sequence data prepared on node 0 must be visible to the other nodes.

Keep credentials out of sbatch templates and logs. Provide HF_TOKEN, GH_TOKEN, and WANDB_API_KEY through the scheduler environment or a restricted secrets file, and never hardcode token values in the script body. For copy-paste environment and sbatch templates, read references/templates.md.

Log Directory

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

srun-native Approach (Preferred)

Slurm spawns all processes directly. No torch.distributed.run, no TRAIN_CMD escaping.

SBATCH Headers

#SBATCH --job-name=<model>-<task>
#SBATCH --nodes=<NNODES>
#SBATCH --ntasks-per-node=8          # Slurm spawns 8 tasks per node
#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

Build and Launch

Use a two-phase srun pattern: first run a single-process uv sync to populate the shared cache, then launch the full multi-node job. The full copy-paste version lives in references/templates.md.

srun-native Key Points

• Phase 1 runs uv sync once on a single node/process, building all wheels into the shared cache on Lustre
• Phase 2's uv sync is a fast no-op (everything is cached) — safe to run on all ranks without sleep guards
• initialize.py + common_utils.py auto-set RANK, WORLD_SIZE, LOCAL_RANK, MASTER_ADDR, MASTER_PORT from SLURM env vars
• Env vars like HF_TOKEN, HF_HOME, UV_CACHE_DIR exported at sbatch level are inherited by srun tasks
• Reference: examples/models/glm/glm_45v/slurm_sft.sh, examples/models/minimax/minimax_m2/slurm_conversion.sh

uv run torch.distributed Approach (Legacy)

Use when the script requires torch.distributed.run (e.g., MLM pretrain_gpt.py) or when Bridge's initialize.py is not in the call path.

1. Add SBATCH Headers

Critical: --ntasks-per-node=1, NOT 8. uv run python -m torch.distributed.run --nproc_per_node=8 spawns 8 processes per node. Using ntasks-per-node=8 causes EADDRINUSE port collisions (8 tasks x 8 procs = 64 per node).

2. Convert to Multi-Node

Replace single-node:

With multi-node (inside TRAIN_CMD string):

MASTER_ADDR and MASTER_PORT are auto-derived from SLURM env vars by initialize.py / common_utils.py — no need to set them.

3. Wrap in TRAIN_CMD + two-phase srun

Use the same two-phase pattern: first a single-process srun to warm the uv cache, then the full run.

Set runtime variables inside the container, but do not inject token values into a long bash -c string. Export credentials through the scheduler or source a restricted secrets file before the job starts. Keep HF_HOME, UV_CACHE_DIR, and NEMO_HOME on shared storage.

4. Launch (two-phase)

Use the two-phase launch template in references/templates.md, keeping #SBATCH --ntasks-per-node=1 for this legacy approach.

5. (Optional) Add Loss Extraction Footer

Interactive GPU Allocation (salloc + srun)

For ad-hoc testing (inference, conversion debugging), always follow these 3 steps:

Step 1: Allocate the node

Step 2: Launch container shell

Step 3: Set up environment inside container

Then run commands with uv run (uses the synced virtualenv):

Pitfalls with interactive allocation:

Debugging Multi-Node Failures

Quick Diagnosis

Check the log for these patterns (in order):

Debugging Checklist

When a multi-node job fails:

1. Check exit code: 1 = Python error, 9 = OOM killed, 143 = SIGTERM (timeout or cascade)
2. Find first failure: Which task/node crashed first? Others get SIGTERM (143) as cascade
3. grep the actual error: Filter out UserWarnings, NCCL frame dumps
4. Check rank 0 specifically: Most save/export errors happen on rank 0
5. Verify EP sizing: For MoE models, ensure num_experts / EP fits in GPU memory with headroom
6. Try interactive first: Use salloc -N 2 -p interactive to iterate faster than sbatch queue

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

Symptom: All ranks on node 2+ show:

Root causes (check in order):

The save_artifacts problem: When trust_remote_code=True, rank 0 runs save_artifacts() (downloads tokenizer, config, custom modeling code) while all other ranks skip directly to dist.barrier(). If save_artifacts is slow or crashes, other ranks timeout.

Fix for ImportError in save_artifacts (hf_pretrained/base.py):

OOM for MoE Models

Symptom: torch.OutOfMemoryError: CUDA out of memory during model loading or forward pass.

Key insight: TP does NOT reduce expert memory. Only EP splits experts across GPUs.

Sizing formula:

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

Rules of thumb:

• Roundtrip (weight-only): can use more experts per GPU (~60GB model params OK)
• Inference (forward pass + KV cache): needs headroom (~40GB model params max)
• Training (activations + optimizer): needs even more headroom (~30GB model params max)

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

Cause: Container's pre-installed megatron-core conflicts with local 3rdparty/Megatron-LM.

Fix: Add uv sync before running:

FP8 Weight Mismatch in Roundtrip

Symptom: Roundtrip completes but shows ❌ for all expert weights and raises ValueError: Weight mismatch detected.

Cause: Original HF weights are FP8, Megatron stores in BF16. Exported weights are BF16. Comparison against original FP8 exceeds atol=1e-1.

This is expected for FP8 models. The conversion is correct; the comparison tolerance is insufficient for the FP8→BF16 precision gap.

WORLD_SIZE Not Set with srun

Symptom: Script exits with "must be launched with torchrun".

Cause: Scripts check os.environ.get("WORLD_SIZE") which torchrun sets but srun doesn't.

Fix: Also check SLURM_NTASKS:

Bridge's common_utils.py helpers (called by initialize.py) populate env vars from SLURM:

Key Gotchas

1. Two-phase srun for uv sync: Run a single-process srun first to warm the cache, then the full multi-node srun. The second uv sync is a fast no-op since everything is already cached on the shared filesystem.

2. --no-container-mount-home is an srun flag, NOT an #SBATCH directive.

3. Escaping inside TRAIN_CMD: Since TRAIN_CMD is a double-quoted string, escape inner $ for Slurm variables that must expand at runtime (not sbatch time):

   • \${SLURM_PROCID}, \${SLURM_JOB_NUM_NODES}, \${SLURM_NODEID}
   • Host-side variables like $GH_TOKEN, $LOGDIR, $WORKDIR expand at sbatch time — no escaping needed.

4. Bridge rm -rf nemo_experiments: Add before training to avoid stale checkpoint auto-resume.

5. MLM needs PYTHONPATH: For pretrain_gpt.py scripts, add inside TRAIN_CMD:

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

6. Node count heuristic: Total GPUs = NNODES * 8. Must satisfy: TP * PP * EP * DP >= total_GPUs where DP = total_GPUs / (TP * PP * EP).

7. NEMO_HOME on shared filesystem for multi-node SFT: The default nemo cache (/root/.cache/nemo) is container-local. Multi-node SFT with packed sequences prepares .npy files on one node that are invisible to others. Set export NEMO_HOME=<SHARED_FS>/cache/nemo so packed data is shared. Without this, ranks on other nodes fail with TypeError: 'NoneType' object is not an iterator.

Full Templates and Command Bodies

For copyable sbatch scaffolding and Bridge/MLM-specific TRAIN_CMD bodies, read references/templates.md.