Slurm 实战：基于 Docker 的分布式训练提交模板

本文介绍如何通过 Slurm 调度系统 提交一个包含 Docker 容器环境 的分布式训练任务。该方案融合了 Slurm 的资源调度能力与 Docker 的环境隔离能力，适用于大规模集群下的深度学习训练（如 Megatron-LM）。

1. 方案概览

在 HPC 集群中，通常要求用户使用非 Root 权限运行作业，且需要特定的 CUDA/Python 环境。本方案的核心逻辑如下：

1. 资源申请：通过 sbatch 向 Slurm 申请计算节点和 GPU 资源。
2. 容器启动：使用 srun 在申请到的每个节点上启动 Docker 容器。
3. 脚本生成：在 Shell 脚本中动态生成容器内部执行的训练脚本 (run_training.sh)，通过挂载卷传递给容器。
4. 分布式通信：利用 PyTorch Elastic (c10d) 后端，结合 Slurm 环境变量自动发现 Master 节点并建立通信。

---

2. 完整作业脚本 (Sbatch Template)

以下是一个可直接使用的 megatron_docker.slurm 脚本模板。该脚本以 Megatron-DeepSeek-R1-Distill-Qwen-1.5B 模型训练为例。

#!/bin/bash
#SBATCH --job-name=megatron              # 作业名称
#SBATCH --output=output_%j.log           # 标准输出日志 (%j 为作业ID)
#SBATCH --error=error_%j.log             # 错误输出日志
#SBATCH --partition=TEST1                # 指定计算分区
#SBATCH --account=test1267               # 账户名 (计费/权限)
#SBATCH --exclude=g[01-86,88-89,91]      # (可选) 排除特定节点
#SBATCH --gres=gpu:8                     # 每个节点申请 8 块 GPU
#SBATCH --ntasks=2                       # 总任务数 (此处等于节点数)
#SBATCH --nodes=2                        # 申请 2 个节点
#SBATCH --ntasks-per-node=1              # 每个节点启动 1 个任务 (即 1 个 Docker 容器)
#SBATCH --cpus-per-task=64               # 每个任务分配 64 个 CPU 核心
#SBATCH --mem=1000G                      # 每个节点分配 1000G 内存

# ==========================================
# 1. 环境变量与路径配置
# ==========================================
# Docker 镜像与路径映射
docker_image="megatron_pytorch:latest"
docker_workspace="/workspace/files"      # 容器内路径
host_workspace="/home/test/test06/qzk/"  # 宿主机路径 (将被挂载)

# 训练数据与模型路径 (相对于容器内路径)
load_path="${docker_workspace}/PLM/Megatron-DeepSeek-R1-Distill-Qwen-1.5B/"
save_path="data/Distill-Qwen-1.5B-OpenR1-Math-94k/checkpoints/"
tensorboard_path="data/Distill-Qwen-1.5B-OpenR1-Math-94k/tensorboard/"

# ==========================================
# 2. 动态生成训练脚本 (run_training.sh)
# ==========================================
# 注意：此脚本将在 Docker 容器内部执行
cat <<RUN_EOF > run_training.sh
#!/bin/bash
set -ex

# 2.1 安装特定依赖 (如需)
# pip install modelbest_sdk-0.2.5.7-py3-none-any.whl

# 2.2 获取分布式环境参数
# 自动检测当前节点 GPU 数量
GPUS_PER_NODE=\$(nvidia-smi --query-gpu=gpu_name --format=csv,noheader | wc -l)

# 从 Slurm 环境变量读取 World Info
WORLD_SIZE=\${SLURM_NTASKS}
RANK=\${SLURM_PROCID}
# 获取主节点 IP (取 hostlist 第一个)
MASTER_ADDR=\$(scontrol show hostname \$SLURM_NODELIST | head -n 1)
MASTER_PORT=6420

echo "Node Info: Rank \${RANK}/\${WORLD_SIZE} | Master: \${MASTER_ADDR}:\${MASTER_PORT} | GPUs: \${GPUS_PER_NODE}"

# 2.3 模型与数据配置
TOKENIZER_MODEL=${docker_workspace}/PLM/DeepSeek-R1-Distill-Qwen-1.5B
# 数据路径格式: "权重 路径"
DATA_PATH="1.0 ${docker_workspace}/Datasets/OpenR1-Math-prefix/OpenR1-Math-220k_deepseek_qwen-prefix"

# 2.4 构建 PyTorch Distributed 参数 (torchrun)
DISTRIBUTED_ARGS=(
    --nproc_per_node \$GPUS_PER_NODE
    --nnodes \$WORLD_SIZE
    --node_rank \$RANK
    --master_addr \$MASTER_ADDR
    --master_port \$MASTER_PORT
    --rdzv_backend c10d
    --rdzv_endpoint \$MASTER_ADDR:\$MASTER_PORT
)

# 2.5 Megatron 模型参数 (示例)
GPT_MODEL_ARGS=(
    --vocab-size 151936
    --make-vocab-size-divisible-by 1
    --num-layers 28
    --hidden-size 1536
    --num-attention-heads 16
    --seq-length 4096
    --max-position-embeddings 4096
    --untie-embeddings-and-output-weights
)

# 2.6 训练超参
TRAINING_ARGS=(
    --micro-batch-size 4
    --global-batch-size 512
    --lr 1e-5
    --lr-decay-style cosine
    --train-iters 500000
    --weight-decay 0.1
    --clip-grad 1.0
    --log-interval 10
    --save-interval 1000
    --eval-interval 1000
    --eval-iters 10
)

# 2.7 输出配置
OUTPUT_ARGS=(
    --save ${docker_workspace}/${save_path}
    --load ${load_path}
    --tensorboard-dir ${docker_workspace}/${tensorboard_path}
)

# 2.8 启动命令
torchrun \
    "\${DISTRIBUTED_ARGS[@]}" \
    pretrain_gpt.py \
    "\${GPT_MODEL_ARGS[@]}" \
    "\${TRAINING_ARGS[@]}" \
    "\${OUTPUT_ARGS[@]}" \
    --tokenizer-model \${TOKENIZER_MODEL} \
    --data-path \${DATA_PATH}
RUN_EOF

# 赋予脚本执行权限
chmod +x run_training.sh

# ==========================================
# 3. 提交任务到 Docker
# ==========================================
echo "Starting Docker container on all nodes..."

# srun 会在 Slurm 分配的每个节点上执行以下命令
srun bash -c "
docker run --rm \\
  --gpus all \\
  --network host \\
  --ipc=host \\
  --ulimit memlock=-1 \\
  --ulimit stack=67108864 \\
  --privileged=true \\
  -v ${host_workspace}:${docker_workspace} \\
  -w ${docker_workspace} \\
  ${docker_image} \\
  bash run_training.sh
"

---

3. 核心机制解析

3.1 路径映射策略

容器内外通过 Volume 挂载保持数据一致性，建议使用统一的变量管理：

• host_workspace: 物理机上的存储路径 (e.g., Lustre/NFS 挂载点)。
• docker_workspace: 容器内的标准工作目录。
• 技巧: 将 load_path 和 save_path 都设置为 docker_workspace 的子目录，确保挂载一次即可访问所有资源。

3.2 分布式环境注入

PyTorch Elastic (torchrun) 需要知晓集群拓扑。我们在 Slurm 脚本中动态获取这些信息并传递给容器：

参数	来源	说明
WORLD_SIZE	$SLURM_NTASKS	总进程数 (通常设为节点数)
RANK	$SLURM_PROCID	当前节点的全局序号 (0, 1, ...)
MASTER_ADDR	scontrol show hostname ...	获取第一个节点的主机名作为 Master
GPUS_PER_NODE	nvidia-smi	动态检测，避免硬编码

3.3 Docker 关键参数

为了获得裸机级别的性能，Docker 启动参数至关重要：

• --gpus all: 透传所有 GPU 设备。
• --network host: 使用宿主机网络栈，避免 NAT 损耗，这对 IB/RoCE 网络至关重要。
• --ipc=host: 共享宿主机共享内存，防止多进程 DataLoader 通信死锁。
• --ulimit memlock=-1: 解除内存锁定限制，支持 RDMA Pin Memory。

---

4. 最佳实践与注意事项

调试建议

首次提交时，建议将脚本中的 sbatch 参数改为 #SBATCH --time=00:10:00 并使用 srun --pty ... 进行交互式测试，确保 Docker 镜像可以正常拉取且挂载路径正确。

1. 镜像缓存: 在没有外网的大规模集群中，确保 docker_image 已提前通过 docker pull 或 docker load 分发到所有计算节点，或者配置了本地 Harbor 仓库。
2. 权限控制: 如果集群禁止普通用户使用 docker 命令，需联系管理员配置 User Namespace 或改用 Singularity/Apptainer (HPC 领域的容器替代方案)。
3. 日志管理: 脚本使用了 set -ex，所有的执行命令都会打印到 output_%j.log 中。如果训练卡死，请检查 error_%j.log 中的 NCCL 报错信息。
4. 死链检查: 若代码中引用了本地 .whl 包 (如 modelbest_sdk), 请确保该文件存在于 host_workspace 路径下，否则容器内 pip install 会失败。