多节点Slurm

将单节点的 uv run python -m torch.distributed.run 命令转换为带Enroot容器支持的多节点Slurm sbatch作业，并调试常见的多节点故障。

首要回答检查清单

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

对于会执行到 initialize.py 的Bridge脚本，优先采用 srun-native 启动形态：#SBATCH --ntasks-per-node=8 和直接 srun ... uv run python <script> ... 启动。不要将这些作业包装在 python -m torch.distributed.run 中。
说明Bridge在 initialize.py 分布式初始化期间会从SLURM变量中推导 RANK、WORLD_SIZE、LOCAL_RANK、MASTER_ADDR 和 MASTER_PORT。
要求仓库、数据、日志、HF_HOME、UV_CACHE_DIR 和 NEMO_HOME 使用共享路径，并确保容器挂载一致。
对于NCCL超时报告，在推测前先进行以下日志首查：
- 过滤警告/帧噪声，grep真实错误
- 检查 Failures: 找到第一个失败的rank和节点
- grep 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上准备的packed-sequence数据必须被其他节点看到。

不要在sbatch模板和日志中包含凭据。通过调度器环境或受限的secrets文件提供 HF_TOKEN、GH_TOKEN 和 WANDB_API_KEY，并切勿在脚本体中硬编码token值。有关可直接复制粘贴的环境变量和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 每个节点派生 8 个任务
#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 是快速no-op（所有内容已缓存）——可安全地在所有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          # 始终为1 — torchrun 处理每节点派生
#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 由 initialize.py / common_utils.py 从SLURM环境变量自动推导——无需设置它们。

3. 包装到TRAIN_CMD并采用两阶段srun

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

在容器内设置运行时变量，但不要将token值注入到长 bash -c 字符串中。通过调度器导出凭据或在作业开始前加载受限的secrets文件。将 HF_HOME、UV_CACHE_DIR、NEMO_HOME 放在共享存储上。

4. 启动（两阶段）

使用 references/templates.md 中的两阶段启动模板，保持此旧版方法的 #SBATCH --ntasks-per-node=1。

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

echo "======================================"
echo "完成。损失值："
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`
`No space left on device` 在 `uv` 或 `pip` 时	容器的 `/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. 查找实际错误（过滤噪声）
grep -a 'Error\|OOM\|CUDA out of memory\|FAILED\|Killed' job.log \
  | grep -v 'UserWarning\|AllocatorConfig\|transformer_engine\|frame\|srun: error'

# 2. 检查哪个rank最先崩溃
grep -a 'Failures:' -A 20 job.log | head -25

# 3. 检查是否存在NCCL超时
grep -a 'ncclUniqueId\|timeout\|crash on rank 0' job.log | head -5

调试检查清单

多节点作业失败时：

检查退出码：1 = Python错误，9 = OOM被杀，143 = SIGTERM（超时或级联效应）
找到第一个故障：哪个任务/节点最先崩溃？其他任务收到SIGTERM(143)作为级联
grep实际错误：过滤掉UserWarnings、NCCL帧转储
特别检查rank 0：大多数保存/导出错误发生在rank 0
验证EP大小：对于MoE模型，确保 num_experts / EP 在留有余量的情况下能装入GPU内存
先尝试交互式：使用 salloc -N 2 -p interactive 比sbatch队列迭代更快

NCCL在 `dist.barrier()` 超时 — “crash on 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

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

原因	如何验证	修复
`save_artifacts` 在rank 0上挂起	错误在 `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()（下载tokenizer、配置、自定义modeling代码），而其他所有rank直接跳到 dist.barrier()。如果 save_artifacts 很慢或崩溃，其他rank会超时。

修复 save_artifacts 中的ImportError (hf_pretrained/base.py)：

# 修改前：
except OSError:
    pass
# 修改后：
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模型参数OK）
推理（前向传播+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())          # 使用 SLURM_PROCID
if "WORLD_SIZE" not in os.environ:
    os.environ["WORLD_SIZE"] = str(get_world_size_safe())  # 使用 SLURM_NTASKS
if "MASTER_ADDR" not in os.environ:
    os.environ["MASTER_ADDR"] = get_master_addr_safe()     # 解析 SLURM_NODELIST
if "MASTER_PORT" not in os.environ:
    os.environ["MASTER_PORT"] = str(get_master_port_safe()) # 从 SLURM_JOB_ID 推导

关键注意事项

两阶段srun用于 uv sync：先运行单进程srun预热缓存，然后运行完整的多节点srun。第二个 uv sync 是快速no-op，因为所有内容已经缓存在共享文件系统上。
--no-container-mount-home 是 srun 标志，不是 #SBATCH 指令。
TRAIN_CMD内的转义：由于 TRAIN_CMD 是一个双引号字符串，内部的 $ 必须转义以在运行时（而不是sbatch时）展开Slurm变量：
- \\${SLURM_PROCID}、\\${SLURM_JOB_NUM_NODES}、\\${SLURM_NODEID}
- 主机端变量如 $GH_TOKEN、$LOGDIR、$WORKDIR 在sbatch时展开——无需转义。
Bridge rm -rf nemo_experiments：在训练前添加，以避免过时的checkpoint自动恢复。
MLM需要PYTHONPATH：对于pretrain_gpt.py脚本，在TRAIN_CMD内添加：

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

节点数估算：总GPU数 = NNODES * 8。必须满足：TP * PP * EP * DP >= total_GPUs，其中 DP = total_GPUs / (TP * PP * EP)。
多节点SFT使用共享文件系统上的 NEMO_HOME：默认nemo缓存（/root/.cache/nemo）是容器本地的。使用packed sequences的多节点SFT会在一个节点上准备 .npy 文件，而其他节点不可见。设置 export NEMO_HOME=<SHARED_FS>/cache/nemo 使packed数据共享。如果不设置，其他节点上的rank会因 TypeError: \'NoneType\' object is not an iterator 而失败。

完整模板和命令体

有关可复制的sbatch脚手架和Bridge/MLM特定的 TRAIN_CMD 命令体，请阅读 references/templates.md。