SLURM集群下PyTorch Lightning DDP的GPU配置错误问询

环境 #

• PyTorch Lightning
• 单节点4GPU的SLURM集群

Trainer YAML配置 #

trainer:
  _target_: lightning.pytorch.trainer.Trainer
  default_root_dir: ${paths.run_dir}
  # ... 其他设置 ...
  accelerator: gpu
  devices: 4
  strategy: ddp

SBATCH脚本 #

#!/bin/bash -l
#SBATCH --partition=gpu
#SBATCH --nodes=1
#SBATCH --gres=gpu:4
#SBATCH --ntasks-per-node=4
#SBATCH --output=/project/home/<project>/output/output.txt
#SBATCH --error=/project/home/<project>/output/error.txt

# 设置Python路径至home目录下的poetry环境

cd repos/myrepo
srun python -m myrepo.main

问题现象 #

1. 当Trainer配置devices: 4时，重复四次报错：

Error in call to target 'lightning.pytorch.trainer.trainer.Trainer':
MisconfigurationException('You requested gpu: [0, 1, 2, 3]\n But your machine only has: [0]')

2. 将devices改为-1时，报错：

ValueError: You set devices=1 in Lightning, but the number of tasks per node configured in SLURM --ntasks-per-node=4 does not match. HINT: Set devices=4.

3. 替换srun python为python后可运行，但收到警告：

The srun command is available on your system but is not used. HINT: If your intention is to run Lightning on SLURM, prepend your python command with srun like so: srun python .../repos/myrepo

疑问 #

1. 如何正确配置PyTorch Lightning的Trainer以适配SLURM的DDP环境？
2. SBATCH脚本中srun python与python的差异是什么？

1. 正确配置PyTorch Lightning Trainer适配SLURM DDP #

问题核心：用srun启动4个任务时，SLURM会为每个任务单独分配1个GPU（--ntasks-per-node=4对应--gres=gpu:4，任务与GPU一一绑定），但原Trainer配置devices:4会让每个任务尝试占用所有4个GPU，导致资源冲突。

正确配置方案 #

• 修改Trainer配置：将devices设为1或auto（PyTorch Lightning会自动识别SLURM分配的单GPU），保留strategy: ddp：

  trainer:
    _target_: lightning.pytorch.trainer.Trainer
    default_root_dir: ${paths.run_dir}
    # ... 其他设置 ...
    accelerator: gpu
    devices: 1  # 或用auto
    strategy: ddp
  

• SBATCH脚本保持不变：继续使用--ntasks-per-node=4和srun python启动，SLURM会生成4个进程，每个进程绑定1个GPU，PyTorch Lightning的DDP会自动完成进程间通信。

简化配置选项 #

PyTorch Lightning支持自动检测SLURM环境，可直接省略strategy和devices，只保留accelerator: gpu，框架会自动启用DDP并适配资源：

trainer:
  _target_: lightning.pytorch.trainer.Trainer
  default_root_dir: ${paths.run_dir}
  # ... 其他设置 ...
  accelerator: gpu

2. srun python与python的差异 #

• srun python：SLURM根据--ntasks-per-node设置启动对应数量的Python进程，每个进程绑定分配的专属资源（这里是1个GPU），是SLURM分布式任务的标准启动方式，能确保DDP正常利用多GPU并行训练。
• python：仅启动单个Python进程，该进程会占用所有4个GPU，此时PyTorch Lightning会自动切换为dp（DataParallel）模式，虽能运行，但无法发挥DDP的高效分布式性能，且会触发SLURM规范警告。

额外验证建议 #

在训练脚本中添加以下代码，可打印每个进程绑定的GPU信息，确认资源分配正确：

import torch
if torch.distributed.is_initialized():
    print(f"Process {torch.distributed.get_rank()} using GPU: {torch.cuda.current_device()}")

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