Slurm调度提交任务时PyTorch无法检测适配版本的CUDA

核心差异定位 #

Jupyter能正常检测GPU但Slurm提交不行，问题根源是Slurm任务运行环境与Jupyter环境不一致，优先从资源分配、环境变量、CUDA可见性这几个方向排查。

具体排查步骤 #

1. 检查Slurm脚本的GPU资源配置 #

确认脚本是否正确申请GPU资源，常见错误点：

• 未添加--gres=gpu参数（需指定数量，如--gres=gpu:1）
• 未指定GPU专属分区（部分集群需加--partition=gpu）
• 未配置对应QoS（部分集群GPU任务需指定QoS）

示例正确配置片段：

#!/bin/bash
#SBATCH --job-name=torch_gpu_test
#SBATCH --partition=gpu
#SBATCH --gres=gpu:1
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1

2. 对比Jupyter与Slurm的环境变量 #

在Jupyter中执行：

!printenv | grep -E "(CUDA|NVIDIA|PATH|LD_LIBRARY_PATH)"

在Slurm脚本中添加同样命令，对比输出差异：

• 重点看CUDA_VISIBLE_DEVICES：Slurm应显示分配的GPU编号（如0），为空则说明资源未分配
• 检查PATH/LD_LIBRARY_PATH：是否包含CUDA 11.8路径（如/usr/local/cuda-11.8/bin、/usr/local/cuda-11.8/lib64）
• 确认PYTHONPATH：是否和Jupyter使用的PyTorch虚拟环境一致

3. 验证Slurm任务中CUDA基础可用性 #

在Slurm脚本中添加以下命令，排除CUDA本身问题：

nvidia-smi
nvcc --version

• 若nvidia-smi无GPU输出：Slurm未正确分配GPU，联系集群管理员确认节点配置
• 若nvcc版本不是11.8：CUDA工具包路径未正确加载

4. 确保PyTorch环境一致性 #

在Slurm脚本中明确激活你的虚拟环境，避免使用系统默认Python：

# 替换为你的虚拟环境路径
source ~/torch_env/bin/activate
python -c "import torch; print(torch.__version__); print(torch.cuda.is_available()); print(torch.cuda.device_count())"

对比Jupyter中torch.__version__和torch.cuda.device_count()的输出，确保版本完全匹配。

5. 检查Slurm资源绑定状态 #

在Slurm任务中执行：

scontrol show job $SLURM_JOB_ID | grep Gres

查看输出是否包含你申请的GPU资源（如gres=gpu:1），无此内容则说明资源申请未生效。

补充信息定位 #

请提供以下内容以便进一步排查：

• torch collect_env完整输出
• Slurm提交脚本全文
• Slurm任务执行时的printenv和nvidia-smi输出

内容的提问来源于stack exchange，提问作者jgbaroja