如何通过Slurm作业数组实现单GPU多任务共享运行？

你的Slurm脚本配置逻辑有误，导致每个作业数组任务都尝试占用独立GPU。核心问题是--ntasks=8与作业数组0-7的组合，让Slurm混淆了资源分配规则，错误地为每个任务分配GPU资源。以下是两种可行的解决方法：

方法一：移除冗余的--ntasks配置 #

作业数组0-7已经对应8个独立任务实例，无需额外指定--ntasks=8。修改后的脚本如下：

#!/bin/bash

#SBATCH --job-name parallel_finetune
#SBATCH --nodes=1
#SBATCH --nodelist=node3
#SBATCH -t 48:00:00
#SBATCH -o ./shell_output/output_%A_%a.output
#SBATCH -e ./shell_output/error_%A_%a.error
#SBATCH --array=0-7
#SBATCH --cpus-per-task=2
#SBATCH --mem-per-cpu=4GB
#SBATCH --gpus=1  # 为整个作业数组分配1个共享GPU

关键说明： #

• 去掉--ntasks=8后，Slurm会为每个数组任务分配指定的CPU和内存，所有任务共享这1个GPU。
• 执行Python脚本时，Slurm会自动设置CUDA_VISIBLE_DEVICES环境变量为分配的GPU编号（通常是0），你的Python代码会默认使用这个GPU。如果需要显式指定，可在代码中添加：

  import torch
  torch.cuda.set_device(0)
  

方法二：明确禁止单任务分配GPU #

如果需要保留--ntasks配置，可通过--gpus-per-task=0强制Slurm不为单个任务分配GPU，仅为整个作业分配1个GPU：

#!/bin/bash

#SBATCH --job-name parallel_finetune
#SBATCH --nodes=1
#SBATCH --nodelist=node3
#SBATCH -t 48:00:00
#SBATCH -o ./shell_output/output_%A_%a.output
#SBATCH -e ./shell_output/error_%A_%a.error
#SBATCH --ntasks=8
#SBATCH --cpus-per-task=2
#SBATCH --mem-per-cpu=4GB
#SBATCH --gpus=1
#SBATCH --gpus-per-task=0  # 禁止每个任务单独获取GPU
#SBATCH --array=0-7

额外注意事项 #

• 显存限制：为避免任务间显存竞争，需在Python脚本中限制单进程显存占用。例如使用PyTorch时：

  import torch
  # 限制单进程使用1/8的GPU显存（对应8个任务）
  torch.cuda.set_per_process_memory_fraction(0.125)
  

• 节点显存检查：确认目标节点GPU显存≥24GB（8个任务×3GB/任务），避免显存不足导致任务崩溃。
• 资源验证：提交作业后，可通过sacct -j <JOB_ID> --format=JobID,Node,Gres查看GPU资源分配情况，确认所有数组任务共享同一个GPU。

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