用Slurm高效提交深度学习任务：以KBQA实体消歧任务为例

用Slurm高效提交深度学习任务：以KBQA实体消歧任务为例

在深度学习研究中，我们常需要在集群上运行大模型任务（如KBQA中的实体消歧），而Slurm（Simple Linux Utility for Resource Management）作为主流的集群资源管理器，能帮我们高效分配CPU/GPU、管理任务队列。本文将以运行run_disamb.py实体消歧脚本为例，从环境配置、Slurm脚本编写、任务管理到问题排查，完整讲解如何用Slurm提交Python深度学习任务，解决你可能遇到的“模块缺失”“GPU跑满”“任务超时”等常见坑。

一、前置背景：我们要解决的任务

本次案例的核心任务是运行KBQA（基于知识库的问答）中的实体消歧脚本run_disamb.py，依赖Hugging Face transformers库加载预训练模型（如BERT），此前曾遇到过模块缺失（packaging） 、模型路径错误、CPU跑模型速度过慢等问题。接下来，我们将通过Slurm规范流程，避免这些问题并高效执行任务。

二、第一步：环境准备（避坑基础）

在提交Slurm任务前，需确保集群上的Python环境正确配置，避免“依赖缺失”“环境不匹配”等低级错误。

1. 激活Conda环境

我们之前创建了专门的PLM环境（用于预训练语言模型），需在Slurm脚本中显式激活（Slurm默认不加载本地环境）：

# 查看已有的Conda环境
conda info --envs
# 激活PLM环境（本地测试时先确认环境可用）
conda activate PLM

2. 安装依赖（解决ModuleNotFoundError）

若遇到No module named 'packaging'或transformers缺失，需在环境中补装依赖，且确保安装路径在PLM环境内：

# 安装packaging（之前踩过的坑）
pip install packaging -i https://pypi.tuna.tsinghua.edu.cn/simple
# 安装transformers及PyTorch（GPU版本）
pip install transformers torch torchvision -i https://pypi.tuna.tsinghua.edu.cn/simple

3. 确认模型文件（避免路径错误）

脚本run_disamb.py需加载预训练模型（如bert-base-uncased），需提前准备好模型文件（本地路径或在线访问权限）：

• 本地路径：将模型文件（含config.json、pytorch_model.bin、vocab.txt）放在../pretrain/bert-base-uncased；
• 在线访问：确保集群能连接Hugging Face Hub（需解决代理问题，或提前下载到本地）。

三、核心：编写Slurm提交脚本

Slurm任务通过 .sh脚本 定义资源需求和执行步骤，以下是针对run_disamb.py的完整脚本（命名为submit_disamb.sh），关键参数已标注注释：

#!/bin/bash
# -------------------------- 1. Slurm资源配置（根据集群调整） --------------------------# 任务提交的分区（GPU分区，用sinfo查看集群分区）可选！若集群有默认分区，可注释掉这行，自动分配）
#SBATCH --partition=gpu_partition 
# 任务名称（自定义，方便后续查看）
#SBATCH --job-name=entity_disamb
# 分配1个任务实例，每个任务用4个CPU核心（配合GPU数据加载）
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=4
# 分配1块GPU（核心！没有这行默认用CPU，速度极慢）
#SBATCH --gres=gpu:1
# 分配32G内存（避免数据缓存不足，根据模型调整）
#SBATCH --mem=32G
# 任务最大运行时间（24小时，避免超时被杀死，可后续优化）
#SBATCH --time=24:00:00
# 输出日志文件（%j自动替换为任务ID，方便区分）
#SBATCH --output=disamb_log_%j.out
# 错误日志文件（单独存放，排查错误更清晰）
#SBATCH --error=disamb_err_%j.err# -------------------------- 2. 环境初始化（避坑关键） --------------------------
# 加载Conda（Slurm默认不加载，需指定你的miniconda路径）
source /home/liang/miniconda3/etc/profile.d/conda.sh
# 激活PLM环境（确保依赖都在这个环境里）
conda activate PLM
# 设置离线模式（避免连接Hugging Face超时，本地有模型时启用）
export HF_HUB_OFFLINE=1# -------------------------- 3. 执行Python脚本 --------------------------
# 切换到脚本所在目录（确保能找到components/model.py等依赖）
cd /fs1/private/user/liang/PLMs-in-Practical-KBQA/main/entity_diamb/
# 打印当前环境信息（调试用，确认GPU和Python路径）
echo "当前Python路径：$(which python)"
echo "GPU是否可用：$(python -c 'import torch; print(torch.cuda.is_available())')"
# 执行脚本（若需传参数，可在后面添加，如--epoch 100）
python run_disamb.py

脚本关键参数解读（必看！）

参数	作用
--partition	可选参数！指定集群分区（若集群有默认分区，未指定时会自动分配）。但需注意：
--gres=gpu:1	必选参数（若需 GPU）！申请 1 块 GPU（核心中的核心！之前 CPU 跑 19 秒 / 迭代，GPU 可缩至 0.1 秒 / 迭代），无论是否指定分区，要 GPU 必须加这行
--cpus-per-task	分配CPU核心数（建议4-8，配合DataLoader多线程加载数据，避免GPU空闲）
source conda.sh	加载Conda环境（否则Slurm会报“conda: command not found”）
HF_HUB_OFFLINE=1	强制离线模式（本地有模型时启用，避免网络问题导致任务失败）

如果你没指定 --partition 却能成功提交任务，想知道自动分配到了哪个分区，可通过以下命令查看：
1、查看集群所有分区及状态：

sinfo  # 输出中“PARTITION”列是分区名，“DEFAULT”列标“*”的是默认分区

2、查看已提交任务的分区信息：

squeue -j 81507 -o "%j %P %T"  # 81507是任务ID，%P会显示任务所在分区

四、任务提交与状态管理

脚本编写完成后，通过简单命令即可提交任务，并实时查看任务状态。

1. 提交任务

在终端进入脚本所在目录，执行sbatch命令提交：

sbatch submit_disamb.sh

若提交成功，会返回类似信息：Submitted batch job 81507（81507是任务ID，后续用它管理任务）。

2. 查看任务状态

用ssqueue命令查看任务队列，重点关注任务ID和状态：

# 查看当前用户的所有任务
squeue -u liang
# 查看指定任务的详细状态（如81507）
squeue -j 81507

任务状态含义：

• PD（Pending）：排队中（等待集群资源，耐心等即可）；
• R（Running）：运行中（正常状态，可查看日志了解进度）；
• FAILED（Failed）：任务失败（需查看错误日志disamb_err_81507.err）。

3. 取消任务

若任务跑错或无需继续，用scancel命令取消：

# 取消任务ID为81507的任务
scancel 81507
# 取消当前用户的所有任务（谨慎使用！）
scancel -u liang

五、常见问题排查（避坑指南）

在提交任务时，你可能会遇到和我之前一样的问题，以下是解决方案汇总：

1. 问题1：ModuleNotFoundError: No module named ‘packaging’

• 症状：任务日志显示缺失packaging或其他模块；
• 原因：Conda环境未激活，或依赖安装在其他环境；
• 解决：
  1. 确认脚本中source conda.sh和conda activate PLM路径正确；
  2. 登录集群后，手动激活PLM环境，重新安装依赖：pip install packaging。

2. 问题2：OSError: 找不到config.json文件

• 症状：日志提示../pretrain/bert-base-uncased无config.json；
• 原因：模型路径错误，或本地未下载完整模型文件；
• 解决：
  1. 用ls ../pretrain/bert-base-uncased查看路径是否存在；
  2. 若路径错误，修改脚本中cd命令后的路径，或调整模型加载路径；
  3. 若文件缺失，从Hugging Face Hub手动下载模型，上传到集群对应路径。

3. 问题3：速度极慢（19秒/迭代）

• 症状：进度条显示19.91s/it，且GPU利用率为0%；
• 原因：未申请GPU，或代码未将模型移到GPU；
• 解决：
  1. 确认Slurm脚本有--gres=gpu:1，且分区是GPU分区；
  2. 在代码中添加GPU适配逻辑（关键！）：

     import torch
     # 自动检测GPU
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     # 模型和数据移到GPU
     model = model.to(device)
     batch = {k: v.to(device) for k, v in batch.items()}
     

  3. 用nvidia-smi查看GPU利用率：srun --jobid=81507 nvidia-smi，若GPU-Util>30%则正常。

4. 问题4：ProxyError: 无法连接PyPI或Hugging Face

• 症状：安装依赖或加载在线模型时，提示“Connection refused”；
• 原因：代理设置拦截网络请求；
• 解决：
  1. 在Slurm脚本中清除代理环境变量：

     unset HTTP_PROXY
     unset HTTPS_PROXY
     

  2. 安装依赖时指定国内镜像：pip install packaging -i https://pypi.tuna.tsinghua.edu.cn/simple。

六、效率优化技巧

解决了“能跑通”的问题后，我们可以进一步优化任务效率，缩短运行时间：

1. 优化DataLoader（减少CPU瓶颈）

在代码中调整DataLoader参数，让CPU数据加载速度跟上GPU计算：

dataloader = DataLoader(dataset=your_dataset,batch_size=16,  # 增大batch size（GPU内存足够时，如16、32）num_workers=4,  # 多线程加载（等于Slurm分配的--cpus-per-task）pin_memory=True,  # 锁定内存，加速数据传输到GPUdrop_last=True  # 丢弃最后一批不足batch size的数据
)

2. 调整Epoch数（避免过拟合+节省时间）

若原脚本设置1000个Epoch（如日志显示0/1000），可先改为10-20个Epoch测试效果，避免不必要的计算：

# 在run_disamb.py中修改Epoch数
args.epochs = 20  # 原1000，根据任务收敛情况调整

3. 查看GPU利用率（避免资源浪费）

任务运行时，用watch实时监控GPU状态，确保GPU没有空闲：

srun --jobid=81507 watch -n 1 nvidia-smi

• 若GPU-Util长期<10%：可能是batch size太小，或num_workers不足；
• 若GPU-Util长期100%：资源利用充分，无需调整。

七、总结

用Slurm提交深度学习任务的核心是“资源配置正确+环境适配+问题快速排查”，针对run_disamb.py这类KBQA任务，记住以下3个关键点：

1. GPU是效率核心：必须在Slurm脚本中申请GPU（--gres=gpu:1），并在代码中把模型/数据移到GPU；
2. 路径和环境别踩坑：确认模型路径存在config.json，Conda环境激活命令正确；
3. 日志是调试神器：任务失败时，优先查看错误日志（disamb_err_%j.err），定位问题比盲目重试更高效。

按照本文的流程，你不仅能跑通run_disamb.py，还能将Slurm脚本适配到其他深度学习任务（如文本分类、问答模型），让集群资源为你的研究高效服务！