⚡

SLURM上运行Megatron-LM训练Skill mcore-run-on-slurm

该技能用于指导用户在SLURM集群上启动和运行分布式Megatron-LM训练任务。涵盖sbatch脚本编写、分布式环境变量配置、CUDA_DEVICE_MAX_CONNECTIONS规则、容器使用、监控与故障诊断。关键词：SLURM、Megatron-LM、分布式训练、sbatch、多节点、torch.distributed.run、CUDA_DEVICE_MAX_CONNECTIONS、NCCL、作业调度。

Megatron-Core训练 0 次安装 0 次浏览 更新于 9/7/2026
名称 mcore-run-on-slurm
描述 如何在SLURM集群上启动分布式Megatron-LM训练任务。涵盖最小sbatch骨架、torch.distributed.run的环境变量设置、不同硬件和并行模式下的CUDA_DEVICE_MAX_CONNECTIONS规则、容器约定、监控和每秩故障诊断。
开源协议 Apache-2.0 when_to_use: 提交SLURM作业；编写或调试sbatch脚本；配置多节点分布式训练；设置MASTER_ADDR / MASTER_PORT / WORLD_SIZE；诊断SLURM作业失败；‘如何在集群上运行’、‘sbatch’、‘多节点训练’。 metadata:
作者 Philip Petrakian ppetrakian@nvidia.com

在SLURM上运行Megatron-LM

首要常量

对于纯文本型SLURM设置问题，在完整脚本之前用以下常量回答：

  • 从每个节点可见的共享工作树路径提交；在启动训练前在脚本中cd到该路径。
  • 每个节点使用一个srun任务，并使用uv run python -m torch.distributed.run启动worker，而不是裸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：Blackwell之前的Hopper/Ampere，若TP>1或CP>1且非FSDP，使用1；Blackwell/GB200不需要；Torch-FSDP2或Megatron-FSDP必须不使用1；overlap_moe_expert_parallel_comm使用32。

先决条件

  • 具有GPU分区提交权限的SLURM集群登录。
  • 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))

# 仅当配置需要时设置CUDA_DEVICE_MAX_CONNECTIONS
# （参见下文）。例如，Blackwell之前且TP>1或CP>1（非FSDP）的情况：
#   export CUDA_DEVICE_MAX_CONNECTIONS=1

srun --ntasks=${NNODES} --ntasks-per-node=1 bash -c '
  # NODE_RANK来自每节点一个任务时的SLURM_NODEID。
  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等）访问相同路径——节点本地路径对同级rank不可见。
  • 跨所有节点使用一个torchrunworker组；不要启动独立的单节点作业。
  • --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的stderr，不仅仅是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>。
  • NCCL失败且没有Python回溯：验证分配、端口可达性、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。裸torchrun可能通过一个看不到venv包的python解释器调度，取决于venv的设置方式。