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
#SBATCH --exclusive

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

#SBATCH --job-name=<model>-<framework>
#SBATCH --nodes=<NNODES>
#SBATCH --ntasks-per-node=1          # ALWAYS 1 — torchrun handles per-node spawning
#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

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:

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

With multi-node (inside TRAIN_CMD string):

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

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

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

Interactive GPU Allocation (salloc + srun)

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

Step 1: Allocate the node

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

Step 2: Launch container shell

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

Step 3: Set up environment inside container

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

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

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

Pitfalls with interactive allocation:

Debugging Multi-Node Failures

Quick Diagnosis

Check the log for these patterns (in order):

# 1. Find the actual error (filter noise)
grep -a 'Error\|OOM\|CUDA out of memory\|FAILED\|Killed' job.log \
  | grep -v 'UserWarning\|AllocatorConfig\|transformer_engine\|frame\|srun: error'

# 2. Check which rank crashed first
grep -a 'Failures:' -A 20 job.log | head -25

# 3. Check for NCCL timeout
grep -a 'ncclUniqueId\|timeout\|crash on rank 0' job.log | head -5

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