nemo-mbridge-multi-node-slurm

多节点 Slurm

将单节点 uv run python -m torch.distributed.run 命令转换为支持 Enroot 容器的多节点 Slurm sbatch 脚本，并排查常见的多节点故障。

首次回答检查清单

转换或调试 Bridge 多节点作业时，请按以下顺序回答：

1. 对于会执行到 initialize.py 的 Bridge 脚本，优先采用 srun-native 启动方式：使用 #SBATCH --ntasks-per-node=8，并直接通过 srun ... uv run python <script> ... 启动。不要用 python -m torch.distributed.run 包装这些作业。
2. 说明 Bridge 会在 initialize.py 分布式初始化期间，根据 SLURM 变量推导出 RANK、WORLD_SIZE、LOCAL_RANK、MASTER_ADDR 和 MASTER_PORT。
3. 要求仓库、数据、日志、HF_HOME、UV_CACHE_DIR 和 NEMO_HOME 使用共享路径，并确保容器挂载配置一致。
4. 对于 NCCL 超时报告，在推测原因之前，先检查首份日志：
   • 过滤掉警告和堆栈帧噪声后，搜索真正的错误
   • 检查 Failures:，找出第一个失败的 rank 和节点
   • 搜索 ncclUniqueId、timeout 或 crash on rank 0

两种方式：srun-native 与 uv run torch.distributed

方式	ntasks-per-node	进程启动方式	适用场景
srun-native（首选）	8	Slurm 为每个节点启动 8 个任务	转换、推理、Bridge 脚本
uv run torch.distributed（旧版）	1	uv run python -m torch.distributed.run 为每个节点启动 8 个进程	MLM pretrain_gpt.py

优先使用 srun-native ——更简单，也能避免 TRAIN_CMD 的 shell 转义问题。Megatron Bridge 会通过 initialize.py 分布式初始化期间调用的 common_utils.py 辅助函数，根据 SLURM 环境变量（SLURM_PROCID、SLURM_NTASKS、SLURM_LOCALID、SLURM_NODELIST）自动推导 RANK、WORLD_SIZE、LOCAL_RANK、MASTER_ADDR 和 MASTER_PORT，因此无需手动设置。

集群环境

仓库、数据、日志、HF_HOME、UV_CACHE_DIR 和 NEMO_HOME 应使用共享文件系统。对于多节点 SFT/PEFT 作业，NEMO_HOME 不能使用容器内的默认路径（/root/.cache/nemo），因为在节点 0 上准备的打包序列数据必须对其他节点可见。

不要在 sbatch 模板和日志中泄露凭据。通过调度器环境或受限的机密文件提供 HF_TOKEN、GH_TOKEN 和 WANDB_API_KEY，切勿在脚本正文中硬编码令牌值。可复制粘贴的环境变量和 sbatch 模板请参阅 references/templates.md。

日志目录

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

srun-native 方式（首选）

Slurm 直接启动所有进程。无需使用 torch.distributed.run，也无需转义 TRAIN_CMD。

SBATCH 标头

#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

构建和启动

使用两阶段 srun 模式：先运行单进程 uv sync 以填充共享缓存，然后启动完整的多节点作业。完整的可复制粘贴版本位于 references/templates.md。

srun-native 要点

• 阶段 1 在单节点上运行一次 uv sync，将所有 wheel 构建到 Lustre 上的共享缓存中
• 阶段 2 的 uv sync 会快速空操作（所有内容均已缓存），可安全地在所有 rank 上运行，无需 sleep 防护
• initialize.py + common_utils.py 会根据 SLURM 环境变量自动设置 RANK、WORLD_SIZE、LOCAL_RANK、MASTER_ADDR 和 MASTER_PORT
• 在 sbatch 级别导出的 HF_TOKEN、HF_HOME、UV_CACHE_DIR 等环境变量会由 srun 任务继承
• 参考：examples/models/glm/glm_45v/slurm_sft.sh、examples/models/minimax/minimax_m2/slurm_conversion.sh

---

uv run torch.distributed 方式（旧版）

当脚本需要 torch.distributed.run（例如 MLM pretrain_gpt.py），或 Bridge 的调用路径中没有 initialize.py 时使用此方式。

1. 添加 SBATCH 标头

#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

重要：使用 --ntasks-per-node=1，而不是 8。uv run python -m torch.distributed.run --nproc_per_node=8 会为每个节点启动 8 个进程。使用 ntasks-per-node=8 会导致 EADDRINUSE 端口冲突（每个节点 8 个任务 × 8 个进程 = 64 个进程）。

2. 转换为多节点

将单节点配置：

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

替换为多节点配置（放在 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 和 MASTER_PORT 会根据 SLURM 环境变量由 initialize.py / common_utils.py 自动推导，无需设置。

3. 使用 TRAIN_CMD 和两阶段 srun 封装

采用相同的两阶段模式：先使用单进程 srun 预热 uv 缓存，然后执行完整任务。

在容器内设置运行时变量，但不要将令牌值注入长 bash -c 字符串。通过调度器导出凭据，或在作业启动前读取权限受限的机密文件。将 HF_HOME、UV_CACHE_DIR 和 NEMO_HOME 放在共享存储上。

4. 启动（两阶段）

使用 references/templates.md 中的两阶段启动模板，并在这种旧版方案中保留 #SBATCH --ntasks-per-node=1。

5.（可选）添加损失提取页脚

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

---

交互式 GPU 分配（salloc + srun）

对于临时测试（推理、转换调试），始终遵循以下 3 个步骤：

步骤 1：分配节点

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

步骤 2：启动容器 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

步骤 3：在容器内设置环境

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

然后使用 uv run 运行命令（使用已同步的虚拟环境）：

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

交互式分配中的常见问题：

错误	原因	修复方法
Cannot find GPU specification	缺少 --gpus-per-node	在 salloc 和 srun 中始终包含 --gpus-per-node=8
invalid partition specified: pool0	分区名称错误	交互式任务使用 interactive，sbatch 使用 batch。检查：sinfo --summarize
Invalid account or account/partition combination	该账户无法使用此分区	检查组合：sacctmgr -nP show assoc where user=$USER format=account,partition
Unable to create step for job... Requested node configuration is not available	-w <node> 与分配冲突	移除 -w 标志——HF 缓存位于共享文件系统上，可从任何节点访问
容器内出现 uv: command not found	容器未预装 uv	使用预装 uv 的容器，或运行 pip install uv
uv 或 pip 执行期间出现 No space left on device	容器的 /root/.cache/ 已满	重定向：export UV_CACHE_DIR=<SHARED_FS>/uv_cache
ModuleNotFoundError: No module named 'megatron.core.activations'	容器预装的 megatron-core 与本地 3rdparty/Megatron-LM 冲突	安装本地版本：pip install -e 3rdparty/Megatron-LM --no-deps --no-build-isolation

---

调试多节点故障

快速诊断

按以下顺序检查日志中的这些模式：

# 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

调试清单

多节点作业失败时：

1. 检查退出代码：1 = Python 错误，9 = 因 OOM 被终止，143 = SIGTERM（超时或级联终止）
2. 找出首次故障：哪个任务/节点最先崩溃？其他任务会因级联终止而收到 SIGTERM（143）
3. grep 实际错误：过滤掉 UserWarnings 和 NCCL 帧转储
4. 重点检查 rank 0：大多数保存/导出错误都发生在 rank 0
5. 验证 EP 配置：对于 MoE 模型，确保 num_experts / EP 在 GPU 显存范围内，并留有余量
6. 先尝试交互式运行：使用 salloc -N 2 -p interactive，比 sbatch 队列更快地迭代

dist.barrier() 处 NCCL 超时——“rank 0 崩溃”

症状：节点 2 及之后的所有 rank 都会显示：

[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

根本原因（按顺序检查）：

原因	验证方法	修复方法
rank 0 上的 save_artifacts 卡住	错误出现在 save_hf_weights → dist.barrier() 中	增加超时时间：init_process_group("nccl", timeout=timedelta(minutes=60))
自定义模型代码中出现 ImportError	grep ImportError job.log	在 save_artifacts 中捕获 ImportError（见下文）
rank 0 在导出期间 OOM	grep 'OutOfMemory' job.log	增加 EP 或节点数
节点间网络问题	仅跨节点 rank 出错	检查 sinfo，尝试其他节点

save_artifacts 问题：当 trust_remote_code=True 时，rank 0 会运行 save_artifacts()（下载分词器、配置和自定义建模代码），而所有其他 rank 会直接跳到 dist.barrier()。如果 save_artifacts 运行缓慢或崩溃，其他 rank 就会超时。

修复 save_artifacts 中的 ImportError（hf_pretrained/base.py）：

# Change:
except OSError:
    pass
# To:
except (OSError, ImportError):
    pass

MoE 模型 OOM

症状：模型加载或前向传播期间出现 torch.OutOfMemoryError: CUDA out of memory。

关键要点：TP 不会减少专家的内存占用。只有 EP 会将专家拆分到多个 GPU 上。

容量估算公式：

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

MiniMax-M2 示例（256 个专家，约 230GB fp8 → 约 460GB bf16）：

配置	节点数	GPU 数	每个 GPU 的专家数	结果
TP=2，EP=4	1	8	64	OOM（专家过多）
TP=2，EP=8	2	16	32	往返转换可用（仅权重），推理时 OOM
TP=1，EP=16	2	16	16	推理可用
TP=2，EP=32	8	64	8	适合训练

经验法则：

• 往返转换（仅权重）：每个 GPU 可以容纳更多专家（模型参数约 60GB 可行）
• 推理（前向传播 + KV 缓存）：需要预留空间（模型参数最多约 40GB）
• 训练（激活值 + 优化器）：需要预留更多空间（模型参数最多约 30GB）

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

原因：容器中预装的 megatron-core 与本地 3rdparty/Megatron-LM 冲突。

修复方法：运行前添加 uv sync：

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

往返转换中的 FP8 权重不匹配

症状：往返转换完成，但所有专家权重均显示 ❌，并报错 ValueError: Weight mismatch detected。

原因：原始 HF 权重为 FP8，Megatron 以 BF16 存储。导出的权重为 BF16。与原始 FP8 权重比较时，差异超出了 atol=1e-1。

这是 FP8 模型的预期行为。 转换是正确的；对于 FP8→BF16 的精度差异，比较容差不足。

使用 srun 时未设置 WORLD_SIZE

症状：脚本退出，并显示“must be launched with torchrun”。

原因：脚本检查 os.environ.get("WORLD_SIZE")，该变量由 torchrun 设置，但 srun 不会设置。

修复方法：同时检查 SLURM_NTASKS：

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

Bridge 的 common_utils.py 辅助函数（由 initialize.py 调用）会根据 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

---

关键注意事项

1. uv sync 分两阶段执行 srun：先运行单进程 srun 预热缓存，然后再运行完整的多节点 srun。第二次 uv sync 会快速无操作完成，因为所有内容都已缓存在共享文件系统中。

2. --no-container-mount-home 是 srun 标志，而不是 #SBATCH 指令。

3. 在 TRAIN_CMD 中转义：由于 TRAIN_CMD 是双引号字符串，因此必须转义需要在运行时（而非 sbatch 时）展开的 Slurm 变量中的 $：

   • \${SLURM_PROCID}、\${SLURM_JOB_NUM_NODES}、\${SLURM_NODEID}
   • 主机端变量（如 $GH_TOKEN、$LOGDIR、$WORKDIR）会在 sbatch 时展开，无需转义。

4. Bridge 中的 rm -rf nemo_experiments：在训练前添加此命令，以避免自动恢复过期的检查点。

5. MLM 需要 PYTHONPATH：对于 pretrain_gpt.py 脚本，在 TRAIN_CMD 中添加：

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

6. 节点数量经验规则：GPU 总数 = NNODES * 8。必须满足：TP * PP * EP * DP >= total_GPUs，其中 DP = total_GPUs / (TP * PP * EP)。

7. 多节点 SFT 的 NEMO_HOME 必须位于共享文件系统上：默认的 nemo 缓存（/root/.cache/nemo） 位于容器本地。对于使用打包序列的多节点 SFT，.npy 文件会在一个节点上生成， 其他节点无法访问。设置 export NEMO_HOME=<SHARED_FS>/cache/nemo，以共享打包数据。 否则，其他节点上的 rank 会因 TypeError: 'NoneType' object is not an iterator 而失败。

完整模板和命令内容

如需可复制的 sbatch 脚手架以及 Bridge/MLM 专用的 TRAIN_CMD 内容，请阅读 references/templates.md。