Slurm+Optuna+PyTorch Lightning调参：多GPU运行差异排查

问题背景 #

我编写了一个示例脚本，使用Optuna结合5折交叉验证，针对感知器的learning rate超参数进行调优。集群节点配备多GPU，采用Slurm作为任务调度器，使用DDP进行分布式计算。但脚本通过Slurm启动与直接在计算节点本地启动时表现不同：

• 通过Slurm命令srun --gres=gpu:2 python ./ttt.py启动时，仅打印一次"=Script Start="，且仅使用1个GPU（经nvidia-smi验证）；
• 直接在计算节点本地用python ./ttt.py启动时，会打印两次"=Script Start="，且正常使用2个GPU。

差异原因 #

核心是Slurm的进程管理逻辑和本地启动的多进程逻辑冲突：

• 本地启动时，脚本里应该用了PyTorch自带的多进程启动机制（比如torch.distributed.run或spawn），会自动启动与GPU数量匹配的2个进程，所以打印两次启动信息，占满2个GPU；
• 用srun启动时，未显式指定任务数，Slurm默认只启动1个进程，而脚本可能检测到Slurm环境后，不再自行启动多进程，因此只跑1个进程，仅占用1个GPU。

解决方法 #

1. 调整Slurm启动命令 #

显式指定任务数等于GPU数量，让Slurm启动对应数量的进程：

srun --gres=gpu:2 --ntasks=2 python ./ttt.py

如果集群对CPU核心有要求，可添加--cpus-per-task=8（数字根据实际需求调整），为每个进程分配足够的CPU资源。

2. 修改脚本的DDP初始化逻辑 #

移除脚本中手动启动多进程的代码（比如torch.multiprocessing.spawn），改用Slurm环境变量初始化DDP：

import os
import torch
import torch.distributed as dist

# 从Slurm环境变量获取分布式信息
rank = int(os.environ['SLURM_PROCID'])
world_size = int(os.environ['SLURM_NTASKS'])
local_rank = int(os.environ['SLURM_LOCALID'])

# 绑定当前进程到对应GPU
torch.cuda.set_device(local_rank)

# 初始化分布式进程组
dist.init_process_group(
    backend='nccl',
    rank=rank,
    world_size=world_size
)

# 后续训练逻辑正常编写即可

3. 验证效果 #

在脚本开头添加打印语句，确认每个进程能正确获取Slurm环境变量：

print(f"===Script Start===")
print(f"当前进程rank: {os.environ.get('SLURM_PROCID')}")
print(f"总进程数: {os.environ.get('SLURM_NTASKS')}")

用修改后的srun命令启动后，应该会打印两次"=Script Start="，且nvidia-smi能看到2个GPU被正常占用。

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