用 Slurm 跑 GPU 集群 —— 比提交更重要的是知道为什么不跑

# 我能用的分区及其限制
sinfo -o "%20P %5a %10l %6D %10T %N"

# 我的账户与 QoS 的关联关系
sacctmgr show assoc user="$USER" format=Account,Partition,QOS,GrpTRES,MaxJobs

# 按 QoS 查看策略
sacctmgr show qos format=Name,Priority,MaxTRESPU%30,MaxJobsPU,Flags%30

# 某个分区里有多少张什么型号的 GPU
sinfo -p gpu-h100 -o "%20N %10c %10m %30G"

#!/bin/bash
#SBATCH --job-name=llm-pretrain
#SBATCH --account=research
#SBATCH --partition=gpu-h100
#SBATCH --qos=normal
#SBATCH --nodes=4
#SBATCH --ntasks-per-node=1
#SBATCH --gpus-per-node=8
#SBATCH --cpus-per-task=64
#SBATCH --mem=0
#SBATCH --time=24:00:00
#SBATCH --exclusive
#SBATCH --output=/scratch/%u/logs/%x-%j.out
#SBATCH --error=/scratch/%u/logs/%x-%j.err

set -euo pipefail
echo "job=$SLURM_JOB_ID nodes=$SLURM_JOB_NODELIST start=$(date -Is)"
srun --cpus-per-task="$SLURM_CPUS_PER_TASK" python -c "import torch; print(torch.cuda.device_count())"

# 节点内 GPU 互联拓扑与 NUMA 布局
nvidia-smi topo -m

# 实际绑定到了哪些核心
srun --cpu-bind=verbose,cores --gpu-bind=verbose,closest hostname

#!/bin/bash
#SBATCH --job-name=pretrain-7b
#SBATCH --partition=gpu-h100
#SBATCH --account=research
#SBATCH --nodes=8
#SBATCH --ntasks-per-node=1
#SBATCH --gpus-per-node=8
#SBATCH --cpus-per-task=96
#SBATCH --mem=0
#SBATCH --exclusive
#SBATCH --time=48:00:00
#SBATCH --output=/scratch/%u/logs/%x-%j.out

set -euo pipefail

# 把第一个节点当作 rendezvous 点
MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n1)
# 端口号从作业编号推导,避免与同时运行的其他作业冲突
MASTER_PORT=$(( 20000 + SLURM_JOB_ID % 20000 ))
export MASTER_ADDR MASTER_PORT
export OMP_NUM_THREADS=$(( SLURM_CPUS_PER_TASK / 8 ))
export NCCL_DEBUG=WARN
export TORCH_NCCL_ASYNC_ERROR_HANDLING=1

srun --cpus-per-task="$SLURM_CPUS_PER_TASK" --kill-on-bad-exit=1 \
  torchrun \
    --nnodes="$SLURM_NNODES" \
    --nproc-per-node=8 \
    --rdzv-backend=c10d \
    --rdzv-endpoint="$MASTER_ADDR:$MASTER_PORT" \
    --rdzv-id="$SLURM_JOB_ID" \
    train.py --config configs/7b.yaml

#SBATCH --nodes=8
#SBATCH --ntasks-per-node=8
#SBATCH --gpus-per-task=1
#SBATCH --cpus-per-task=12

srun python train.py --config configs/7b.yaml

# train.py 开头 —— 把 Slurm 的环境变量搬到分布式初始化所需的变量上
import os
import torch
import torch.distributed as dist

if "SLURM_PROCID" in os.environ and "RANK" not in os.environ:
    os.environ["RANK"] = os.environ["SLURM_PROCID"]
    os.environ["WORLD_SIZE"] = os.environ["SLURM_NTASKS"]
    os.environ["LOCAL_RANK"] = os.environ["SLURM_LOCALID"]

torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
dist.init_process_group(backend="nccl")
print(f"rank {dist.get_rank()}/{dist.get_world_size()} "
      f"on {os.uname().nodename} gpu {torch.cuda.current_device()}")

srun --container-image=/scratch/images/train-2026-08.sqsh \
     --container-mounts=/scratch:/scratch,/data:/data:ro \
     --container-workdir=/workspace \
     python train.py

#!/bin/bash
#SBATCH --job-name=lr-sweep
#SBATCH --array=0-11%4          # 12 个作业,同时最多跑 4 个
#SBATCH --gpus-per-node=1
#SBATCH --cpus-per-task=8
#SBATCH --time=02:00:00
#SBATCH --output=/scratch/%u/logs/%x-%A_%a.out

LRS=(1e-5 2e-5 5e-5 1e-4)
RANKS=(8 16 32)
lr=${LRS[$(( SLURM_ARRAY_TASK_ID % 4 ))]}
rank=${RANKS[$(( SLURM_ARRAY_TASK_ID / 4 ))]}

echo "task=$SLURM_ARRAY_TASK_ID lr=$lr rank=$rank"
srun python finetune.py --lr "$lr" --lora-rank "$rank" \
     --run-name "sweep-$SLURM_ARRAY_JOB_ID-$SLURM_ARRAY_TASK_ID"

# 把五段 24 小时的训练串联起来。每一段都要等前一段成功后才开始
prev=""
for i in $(seq 1 5); do
  if [ -z "$prev" ]; then
    prev=$(sbatch --parsable pretrain.sbatch)
  else
    prev=$(sbatch --parsable --dependency="afterok:$prev" pretrain.sbatch)
  fi
  echo "segment $i -> job $prev"
done

squeue -u "$USER" -o "%.10i %.12P %.20j %.8T %.10M %.6D %R"
squeue -j 123456 --start          # 预计开始时间
sprio -j 123456 -l                # 按优先级分量拆解出的分数
scontrol show job 123456          # 完整的请求资源与原因字符串
sinfo -R                          # down/drain 状态的节点及原因

sacct -j 123456 --format=JobID,JobName%20,State%20,ExitCode,Elapsed,MaxRSS,ReqTRES%40,NodeList%20
sacct -j 123456 --format=JobID,State,DerivedExitCode,Comment%40
seff 123456        # 如果装了 slurm-contribs,会给出一份摘要

srun -w gpu-node-042 nvidia-smi --query-gpu=index,name,ecc.errors.uncorrected.volatile.total --format=csv
srun -w gpu-node-042 nvidia-smi --query-remapped-rows=gpu_bus_id,remapped_rows.pending,remapped_rows.failure --format=csv
srun -w gpu-node-042 bash -c 'dmesg -T | grep -i xid | tail -20'

#SBATCH --signal=B:USR1@180     # 终止前 180 秒向批处理 shell 发送 USR1
#SBATCH --requeue               # 允许重新入队

CKPT_DIR=/scratch/$USER/ckpt/$SLURM_JOB_NAME
mkdir -p "$CKPT_DIR"

on_preempt() {
  echo "[$(date -Is)] 收到 USR1,请求保存检查点。"
  touch "$CKPT_DIR/SAVE_AND_EXIT"     # 训练循环每一步都会检查这个文件
  wait "$SRUN_PID"                    # 等待保存完成
  echo "[$(date -Is)] 正在重新入队。"
  scontrol requeue "$SLURM_JOB_ID"
  exit 0
}
trap on_preempt USR1

srun --cpus-per-task="$SLURM_CPUS_PER_TASK" python train.py --ckpt-dir "$CKPT_DIR" &
SRUN_PID=$!
wait "$SRUN_PID"