mcore-run-on-slurm

在 SLURM 上运行 Megatron-LM

优先给出答案的常量

对于仅涉及文本的 SLURM 设置问题，请先给出以下常量，再提供完整脚本：

• 从所有节点均可访问的共享工作树路径提交作业；在脚本中启动训练前先 cd 到该路径。
• 每个节点使用一个 srun 任务，并通过 uv run python -m torch.distributed.run 启动工作进程，而不是直接使用 torchrun。
• 使用 scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n1 设置 MASTER_ADDR，设置 MASTER_PORT、NNODES=${SLURM_NNODES}、GPUS_PER_NODE=<GPUS_PER_NODE>，以及 WORLD_SIZE=$((NNODES * GPUS_PER_NODE))。
• 向 torch.distributed.run 传入 --nnodes、--nproc-per-node、--node-rank、--master-addr 和 --master-port。
• CUDA_DEVICE_MAX_CONNECTIONS：对于 TP>1 或 CP>1 且非 FSDP 的 Blackwell 之前架构 Hopper/Ampere，使用 1；Blackwell/GB200 无需设置；Torch-FSDP2 或 Megatron-FSDP 不得使用 1；使用 overlap_moe_expert_parallel_comm 时使用 32。

前提条件

• 可登录 SLURM 集群，并有权向 GPU 分区提交作业。
• Megatron-LM 已检出到分配中的所有节点均可访问的文件系统（NFS、Lustre 或类似文件系统）。所有节点必须能访问代码、数据、检查点和输出所用的相同路径。
• 已安装 uv；提交前在工作树上运行一次 uv sync --extra training --extra dev（或 --extra lts），以生成 .venv 并确保所有节点均可访问。

最小 sbatch 脚本

在工作树中保存为 run_megatron.slurm：

#!/bin/bash
#SBATCH --job-name=megatron
#SBATCH --account=<SLURM_ACCOUNT>
#SBATCH --partition=<SLURM_PARTITION>
#SBATCH --nodes=<NODES>
#SBATCH --ntasks-per-node=1
#SBATCH --gpus-per-node=<GPUS_PER_NODE>
#SBATCH --time=<HH:MM:SS>
#SBATCH --output=logs/%x-%j.out
#SBATCH --error=logs/%x-%j.err

set -euo pipefail
cd <MEGATRON_WORKTREE>

export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n1)
export MASTER_PORT=${MASTER_PORT:-29500}
export NNODES=${SLURM_NNODES}
export GPUS_PER_NODE=<GPUS_PER_NODE>
export WORLD_SIZE=$((NNODES * GPUS_PER_NODE))

# Set CUDA_DEVICE_MAX_CONNECTIONS only when your configuration requires it
# (see the section below). Example for pre-Blackwell with TP>1 or CP>1
# (non-FSDP):
#   export CUDA_DEVICE_MAX_CONNECTIONS=1

srun --ntasks=${NNODES} --ntasks-per-node=1 bash -c '
  # NODE_RANK comes from SLURM_NODEID with one task per node.
  NODE_RANK=${SLURM_NODEID}
  uv run python -m torch.distributed.run \
    --nnodes='"${NNODES}"' \
    --nproc-per-node='"${GPUS_PER_NODE}"' \
    --node-rank=${NODE_RANK} \
    --master-addr='"${MASTER_ADDR}"' \
    --master-port='"${MASTER_PORT}"' \
    pretrain_gpt.py \
      <MEGATRON_ARGS>
'

提交：

mkdir -p logs && JOB_ID=$(sbatch --parsable run_megatron.slurm)
echo "Submitted ${JOB_ID}"

多节点规则

• 从计划运行作业的工作树提交，或在脚本中 cd 到该目录。所有节点必须能访问共享文件系统（NFS、Lustre 或类似文件系统）上的相同路径——节点本地路径对其他进程无法访问。
• 在所有节点上使用同一个 torchrun 工作进程组；不要启动彼此独立的单节点作业。
• --nproc-per-node 应等于每个节点可见的 GPU 数量。
• 将检查点、tensorboard 数据和结构化日志写入共享存储。

CUDA_DEVICE_MAX_CONNECTIONS

正确的值取决于硬件和并行模式。不要无条件导出此变量：

• **Blackwell 之前的架构（Hopper、Ampere），TP>1 或 CP>1，且非 FSDP：**设为 1。相关代码路径会对此进行断言——如果值不是 1，会产生断言错误，而不是静默死锁。
• **Blackwell：**无需设置；设置也不会产生效果。
• **Torch-FSDP2 或 Megatron-FSDP：**不得设为 1。保持环境变量未设置，或设为大于 1 的值。
• **启用 overlap_moe_expert_parallel_comm：**设为 32。

在配置要求时，请在 sbatch 脚本中显式设置。

容器

许多站点会在容器中运行 Megatron-LM（部分集群使用 enroot/pyxis，其他集群使用 singularity）。如果使用容器，uv 管理的 .venv 必须位于容器内可访问的路径上，并且容器镜像必须提供仓库所需的 CUDA / NCCL / torch 版本（参见 docker/.ngc_version.dev 和 .ngc_version.lts）。上方的基本脚本保持不变；使用调度器的容器标志（--container-image=…、--container-mounts=… 等）包装 srun 调用。

监控和收集

squeue -j "$JOB_ID" -o "%.10i %.8T %.10M %.6D %R"
sacct -j "$JOB_ID" --format=JobID,State,ExitCode,Elapsed
scancel "$JOB_ID"

如果训练脚本会写入结果产物（例如 rank 0 生成的 JSON 指标文件、最终检查点等），请轮询该产物，而不要只等待 squeue 状态。SLURM 将作业标记为完成前，通常就会出现有用的输出；轮询产物可让你在产物生成后立即取消作业，而不必让资源分配一直占用到超时。

故障诊断

扫描所有 rank 的标准错误输出，而不只是 rank 0。最早出现的非 NCCL Python 回溯通常就是根本原因；其他 rank 后续发生的 NCCL 超时，则是首次崩溃导致的连带症状。

快速分类：

• OOM：调整前，记录 rank、阶段（前向传播 / 反向传播 / 优化器）、批次大小、序列长度、并行方式（TP/DP/CP/PP）和峰值内存。
• 形状 / 可整除性错误：检查 WORLD_SIZE = TP × DP × CP × PP 和注意力头数的可整除性（num_attention_heads % TP == 0）。
• 导入错误：工作树错误、未运行 uv sync，或 PYTHONPATH 过期。启动前确认已执行 cd <MEGATRON_WORKTREE>。
• 没有 Python 回溯的 NCCL 故障：检查资源分配、端口连通性、MASTER_ADDR 解析情况，以及所有 rank 上的命令是否一致。

常见问题

• 首次提交前忘记运行 uv sync。如果 venv 不存在，每个作业都会在 srun 内重建它，导致每个作业耗费数分钟。
• 将日志写入作业结束后会消失的节点本地路径。始终写入共享文件系统。
• 盲目设置 CUDA_DEVICE_MAX_CONNECTIONS=1。合适的值取决于硬件和并行模式（参见上方专门章节）。在 FSDP 下将其设为 1 会导致另一种问题；在 Blackwell 上则不起作用；在 Blackwell 之前的硬件上使用 TP>1 或 CP>1（非 FSDP）时，代码会触发断言，而不会死锁。
• 直接运行 torchrun，而不是 uv run python -m torch.distributed.run。根据 venv 的配置方式，直接运行 torchrun 时可能会通过一个无法访问 venv 包的 Python 解释器来调度。