Slurm 作业提交与管理指南

摘要: 本文档为集群用户提供标准的 Slurm 作业提交模板，重点解析了从单机调试到多机多卡分布式训练 (PyTorch/DeepSpeed) 的脚本编写规范及常用管理命令。

1. Slurm 简介

Slurm (Simple Linux Utility for Resource Management) 是业界最主流的开源集群调度系统。它负责统一管理计算资源（CPU/GPU/内存），用户只需提交脚本，Slurm 会自动排队并在资源满足时执行。

2. 基础作业模板 (Single Node)

适用于单机调试、数据预处理或推理任务。

2.1 脚本示例 (submit_single.sh)

#!/bin/bash
#SBATCH --job-name=demo_job           # 作业名称
#SBATCH --partition=debug             # 分区名称 (请咨询管理员)
#SBATCH --nodes=1                     # 申请节点数
#SBATCH --ntasks-per-node=1           # 每节点任务数
#SBATCH --cpus-per-task=8             # 每个任务分配 CPU 核数
#SBATCH --gres=gpu:1                  # 申请 1 块 GPU
#SBATCH --output=%j.out               # 标准输出 (%j = JobID)
#SBATCH --error=%j.err                # 错误日志

echo "Job Start Time: $(date)"
echo "Running on node: $(hostname)"

# 加载环境
module load cuda/12.1
source activate my_env

# 执行命令
nvidia-smi
python train.py --epochs 10

echo "Job End Time: $(date)"

2.2 提交命令

sbatch submit_single.sh

3. 分布式训练模板 (Multi-Node)

适用于跨节点的大模型训练 (LLM Pretraining/SFT)。以下示例展示 4 节点 x 8 卡 (共 32 卡) 的配置。

3.1 提交脚本 (submit_dist.sh)

#!/bin/bash
#SBATCH --job-name=llm_train
#SBATCH --partition=gpu_p1
#SBATCH --nodes=4                     # 申请 4 个节点
#SBATCH --ntasks-per-node=1           # 每节点 1 个主进程
#SBATCH --cpus-per-task=64            # 独占节点所有 CPU
#SBATCH --gres=gpu:8                  # 每节点 8 卡
#SBATCH --exclusive                   # 独占节点 (推荐)
#SBATCH --output=logs/%x-%j.out

# 1. 获取 Master 节点信息 (用于 PyTorch DDP/Deepspeed 初始化)
nodes=( $( scontrol show hostnames $SLURM_JOB_NODELIST ) )
nodes_array=($nodes)
head_node=${nodes_array[0]}
head_node_ip=$(srun --nodes=1 --ntasks=1 -w "$head_node" hostname --ip-address)

export MASTER_ADDR=$head_node_ip
export MASTER_PORT=29500
export WORLD_SIZE=$(( SLURM_NNODES * SLURM_GPUS_ON_NODE ))  # 总卡数

echo "Node List: $SLURM_JOB_NODELIST"
echo "Master: $MASTER_ADDR:$MASTER_PORT"

# 2. 生成运行脚本 (run.sh)
cat <<EOF > run_task.sh
#!/bin/bash
# 激活环境
source /home/user/anaconda3/bin/activate llm

# 启动训练
torchrun \
    --nnodes=$SLURM_NNODES \
    --nproc_per_node=8 \
    --rdzv_id=$SLURM_JOB_ID \
    --rdzv_backend=c10d \
    --rdzv_endpoint=$MASTER_ADDR:$MASTER_PORT \
    train_llm.py \
    --deepspeed ds_config.json \
    --model_name_or_path /data/models/llama2-7b