企业官网建设流程全解析

在 SLURM 集群上运行 Megatron-LM 分布式训练：sbatch 脚本骨架、torch.distributed.run 环境配置与故障诊断实战指南

【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM

本篇技术指南以 Megatron-LM 仓库内置的mcore-run-on-slurm技能（见 skills/mcore-run-on-slurm/SKILL.md）为核心，讲解如何在 SLURM 集群上正确提交与运行 Megatron-LM 分布式训练任务。内容覆盖最小可用的 sbatch 脚本骨架、torch.distributed.run所需的环境变量（MASTER_ADDR/MASTER_PORT/WORLD_SIZE）、CUDA_DEVICE_MAX_CONNECTIONS在不同硬件与并行模式下的取值规则、容器（enroot/pyxis/singularity）运行约定、任务监控，以及按 rank 逐层定位多节点任务失败根因的排查方法。读者阅读完本文后，将能够独立编写、提交、监控并调试自己的多节点 Megatron-LM 训练作业。

速查常量：先记住这些结论

在深入完整脚本之前，先把面向纯文本提问的“答案优先”（Answer-First）常量记住，它们几乎覆盖了 90% 的 SLURM 配置问题：

• 共享工作区：从所有节点可见的共享 worktree 路径提交作业，并在脚本中cd到该目录后再启动训练。代码、数据、检查点与输出必须对所有节点可达。
• 每节点一个srun任务：使用uv run python -m torch.distributed.run拉起 worker，而不是裸调torchrun。
• 环境变量：MASTER_ADDR通过scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n1取首节点；MASTER_PORT默认 29500；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：pre-Blackwell（Hopper/Ampere）且 TP>1 或 CP>1 且非 FSDP 时取1；Blackwell/GB200 无需设置；Torch-FSDP2 或 Megatron-FSDP 场景严禁取1；开启overlap_moe_expert_parallel_comm时取32。

前提条件

在提交第一个作业之前，请确认以下三项就绪：

1. SLURM 集群账户：拥有登录权限，并能向 GPU 分区（partition）提交作业。
2. 共享文件系统上的 Megatron-LM：仓库必须签出在分配内所有节点都能访问的文件系统上（NFS、Lustre 等）。所有节点必须能访问到完全相同的代码、数据、检查点与输出路径——节点本地路径对其他 rank 不可见，这一点在多节点规则部分还会强调。
3. uv环境：在提交前，于 worktree 内先执行一次uv sync --extra training --extra dev（或--extra lts），把.venv物化出来并保证每个节点都能看到。这些 extra 定义于仓库根目录的 pyproject.toml（[project.optional-dependencies]一节）：training包含flask-restful、sentencepiece、tiktoken、wandb、transformers等训练依赖，dev包含nvidia-modelopt[torch]、nvidia-resiliency-ext等开发/调试依赖，而lts目前是保留的别名（注释明确说明 LTS 依赖已直接钉在 docker/Dockerfile.ci.lts 中，保留该别名仅用于兼容pip install megatron-core[lts]）。

从源码结构看，Megatron-LM 使用uv管理虚拟环境，uv sync生成的.venv是torch.distributed.run得以发现全部依赖包的前提；这也是为什么技能文档反复强调提交前先uv sync。

最小 sbatch 脚本

将以下内容保存为 worktree 中的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}"

脚本设计要点逐一说明：

• --ntasks-per-node=1+ 单任务启动 worker：每个节点只起一个srun任务，该任务内部再通过torch.distributed.run派生--nproc-per-node个 worker 进程。NODE_RANK直接取自SLURM_NODEID（每节点一个任务时二者一一对应），无需额外计算。
• set -euo pipefail：任何命令失败立即退出，避免半崩溃任务占着分配继续空转。
• 环境变量导出：MASTER_ADDR、MASTER_PORT、NNODES、GPUS_PER_NODE、WORLD_SIZE是torch.distributed.run与其 worker 建立通信所必需的；WORLD_SIZE = NNODES × GPUS_PER_NODE即全局进程总数。
• 训练入口：示例使用pretrain_gpt.py（仓库根目录的 pretrain_gpt.py），同样的骨架可直接换成 pretrain_hybrid.py、pretrain_mamba.py、pretrain_vlm.py 或其他入口，<MEGATRON_ARGS>处填入对应--tensor-model-parallel-size、--pipeline-model-parallel-size、--micro-batch-size、--train-iters等训练参数。

多节点运行规则

• 从预期运行的 worktree 提交，或在脚本中cd到它。所有节点必须通过共享文件系统（NFS、Lustre 等）访问同一路径——节点本地路径对其他 rank 不可见。
• 在所有节点上使用同一个torchrunworker 组，不要启动相互独立的单节点作业。整个作业必须是一个统一的torch.distributed.run调用跨全部节点。
• --nproc-per-node应等于每节点可见 GPU 数（通常即--gpus-per-node），确保 worker 数与物理资源一致。
• 检查点、TensorBoard 数据与结构化日志全部写入共享存储，保证任务结束后仍可读取、可续训。

CUDA_DEVICE_MAX_CONNECTIONS 的取值规则

不要无条件导出这个环境变量，正确取值取决于硬件与并行模式：

源码层面的依据

CUDA_DEVICE_MAX_CONNECTIONS=1之所以在 pre-Blackwell 的 TP/CP 场景下是硬性要求，是因为 Megatron-LM 的张量并行反向传播依赖它来强制通信 kernel 先于计算 kernel 被调度。在 megatron/core/tensor_parallel/layers.py 中可以看到多处这样的注释：异步 all-gather、异步 all-reduce、reduce-scatter 都“依赖CUDA_DEVICE_MAX_CONNECTIONS=1确保通信在权重梯度计算之前被调度”，以让通信与计算真正重叠——这对正确性不是必需的，但对性能至关重要，因此代码通过环境变量强制调度顺序。

反向依赖关系同样存在：DDP 的梯度桶（bucket）分组逻辑在 megatron/core/distributed/distributed_data_parallel.py 与 megatron/core/distributed/param_and_grad_buffer.py 中明确说明，当 fp8 与 bf16 桶并存且启用 vpp 时，“因为使用了CUDA_DEVICE_MAX_CONNECTIONS=1，背靠背的多次通信会阻碍通信与计算重叠”，因此需要把多个桶合并进桶组。这说明该变量与 DDP 桶策略存在联动，随意取值会破坏性能设计。

另外，在 megatron/core/parallel_state.py 中，UCC 后端初始化会对该变量做硬断言：CUDA_DEVICE_MAX_CONNECTIONS被设置为1时直接报错（assert os.environ["CUDA_DEVICE_MAX_CONNECTIONS"] != "1"），因为 UCC 要求该值大于 1 才能保证重叠通信。这进一步印证了“不要盲目设置”的告诫。

overlap_moe_expert_parallel_comm是模型并行配置项，定义于 megatron/core/model_parallel_config.py：其作用是把 MoE 专家并行（EP）的 all-to-all 通信与流水线 1F1B 阶段（或非流水线调度）中不同 micro-batch 的独立计算重叠。该配置与 FSDP 的交互约束也在 megatron/core/distributed/fsdp/mcore_fsdp_adapter.py 中有详细体现（例如与fsdp_double_buffer、细粒度参数 gather hook 的联动约束）。

结论：当你的配置需要时，在 sbatch 脚本中显式设置该变量；不需要时保持不设置，不要盲目写export CUDA_DEVICE_MAX_CONNECTIONS=1。

容器化运行约定

许多站点在容器内运行 Megatron-LM（部分集群使用 enroot/pyxis，另一些使用 singularity）。若采用容器方式，需满足两个前提：

1. uv 管理的.venv必须位于容器内可见的路径上，否则 worker 找不到依赖；
2. 容器镜像必须提供仓库所期望的 CUDA / NCCL / torch 版本。仓库通过 docker/.ngc_version.dev（当前为nvcr.io/nvidia/pytorch:26.08-py3）与 docker/.ngc_version.lts（当前为nvcr.io/nvidia/pytorch:25.09-py3）两个文件声明参考的 NGC PyTorch 容器版本，dev 与 lts 两套镜像分别对应不同依赖锁定策略，详见 docker/Dockerfile.ci.dev 与 docker/Dockerfile.ci.lts。

上述 sbatch 骨架保持不变，只需用调度器（SLURM）的容器标志包住srun调用，例如--container-image=…、--container-mounts=…（具体标志名取决于你的 SLURM 插件实现）。

监控与结果收集

任务提交后，用以下命令监控状态：

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

• squeue查看任务状态（State）、运行时长（TIME）、节点数（NODES）与所在节点（NODELIST(REASON)）；
• sacct在任务结束后查询最终 State、ExitCode 与 Elapsed，用于事后复盘；
• scancel用于取消任务、释放分配。

进阶建议：轮询产物而非仅看任务状态。如果训练脚本会写出结果产物（例如 rank 0 输出的 JSON 指标文件、最终检查点等），应轮询该产物，而不是只等squeue状态变化。有用的输出往往在 SLURM 标记任务完成之前就已出现——一旦检测到产物落地即可主动scancel，而不是白白占用分配直到超时。

失败诊断：逐 rank 定位根因

排查多节点训练失败时的第一原则：扫描每个 rank 的 stderr，而不是只看 rank 0。最早的、非 NCCL 的 Python traceback 通常是根因，其后其他 rank 出现的 NCCL 超时只是第一次崩溃的下游症状。

按以下分类快速定位：

• OOM（显存不足）：记录 rank、阶段（forward / backward / optimizer）、batch size、序列长度、并行配置（TP/DP/CP/PP）与峰值显存，再调整显存策略（如减小 micro-batch、开启重计算等，Megatron-LM 的激活重计算配置见 megatron/core/recompute.py）。
• 形状 / 整除性错误：校验WORLD_SIZE = TP × DP × CP × PP是否成立，以及 head 数是否能被 TP 整除（num_attention_heads % TP == 0）。后者在仓库中有硬性校验：见 megatron/core/transformer/transformer_config.py 中num_attention_heads % tensor_model_parallel_size != 0时直接抛出ValueError的实现。
• Import 错误：通常是错误的 worktree、漏掉uv sync，或PYTHONPATH过期。确认启动前确实cd <MEGATRON_WORKTREE>，并核对.venv是否已物化且对所有节点可见。
• 无 Python traceback 的 NCCL 失败：优先检查分配是否正常、端口是否可达、MASTER_ADDR解析是否一致，以及各 rank 的命令是否完全一致（尤其注意环境变量是否在每个节点都生效）。

常见陷阱清单

总结

在 SLURM 集群上跑通 Megatron-LM 多节点训练，核心就三件事：共享文件系统上的 worktree + 物化好的 uv venv、每节点一个srun任务拉起统一torch.distributed.runworker 组、按硬件与并行模式精确设置CUDA_DEVICE_MAX_CONNECTIONS。在此基础上，配合squeue/sacct/scancel的监控组合、对结果产物的轮询策略，以及“从最早的非 NCCL traceback 入手、逐 rank 扫 stderr”的诊断原则，即可快速定位 OOM、整除性、Import 与 NCCL 四类高频故障。本指南对应的完整技能定义与评估元数据分别见 skills/mcore-run-on-slurm/SKILL.md 与 skills/mcore-run-on-slurm/skill-card.md。

【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM