nemo-mbridge-multi-node-slurm

What this skill does

# 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

| Approach | `ntasks-per-node` | Process spawning | Best for |
|---|---|---|---|
| **srun-native** (preferred) | 8 | Slurm spawns 8 tasks/node | Conversion, inference, Bridge scripts |
| **uv run torch.distributed** (legacy) | 1 | `uv run python -m torch.distributed.run` spawns 8 procs/node | MLM pretrain_gpt.py |

**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

```text
<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

```bash
#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

```bash
#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:

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

With multi-node (inside `TRAIN_CMD` string):

```bash
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

```bash
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

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

### Step 2: Launch container shell

```bash
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

```bash
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):

```bash
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:**

| Error | Cause | Fix |
|---|---|---|
| `Cannot find GPU specification` | Missing `--gpus-per-node` | Always include `--gpus-per-node=8` in both `salloc` and `srun` |
| `invalid partition specified: pool0` | Wrong partition name | Use `interactive` for interactive, `batch` for sbatch. Check: `sinfo --summarize` |
| `Invalid account or account/partition combination` | Partition not available for account | Check combos: `sacctmgr -nP show assoc where user=$USER format=account,partition` |
| `Unable to create step for job... Requested node configuration is not available` | `-w <node>` conflicts with allocation | Remove `-w` flag — HF cache is on shared filesystem, accessible from any node |
| `uv: command not found` inside container | Container doesn't have `uv` pre-installed | Use a container with `uv` pre-installed, or `pip