如何利用PyTorch在GPU上高效并行强化学习实验？

问题 #

需运行1350个独立强化学习实验，串行单实验耗时约10分钟。实验采用相同NN架构但模型与参数独立，且模型规模极小，串行时GPU显存仅被少量占用。

使用Google Colab的T4 GPU及torch.multiprocessing包实现并行，代码示例如下：

import torch.multiprocessing as mp
from torch.multiprocessing import Pool

def run_experiment(param):
   # 创建环境
   # 基于param创建智能体（模型）
   # 让智能体与环境交互
   return

if __name__ == '__main__':
   params = ...
   mp.set_start_method('spawn') # CUDA runtime不支持fork方法
   num_p = 8                    # 并行进程数
   with Pool(num_p) as p:
      p.map(run_experiment, params)

设置num_p=8时虽有提速，但总耗时仍超28小时，GPU显存利用率仍极低。为探究原因开展小实验，结果显示：在未超出15GB可用显存的情况下，进程数翻倍并未使总执行速度翻倍，且首次run_experiment()完成时间随进程数增加而变长。现询问如何正确实现GPU并行以充分利用显存？

解决方案 #

1. 替换Pool，手动管理进程并绑定GPU #

spawn启动的进程会重复初始化CUDA上下文，Pool的调度逻辑也会带来额外开销。可以手动控制进程数量（因模型极小，可大幅提升至16-32），并确保所有进程共享单卡资源：

import torch.multiprocessing as mp
import torch

def run_experiment(param):
    device = torch.device('cuda:0')
    # 创建环境
    agent = create_agent(param).to(device)
    # 智能体与环境交互逻辑
    return

if __name__ == '__main__':
    params = [...]
    mp.set_start_method('spawn')
    num_p = 24  # 根据GPU显存调整，模型极小可尝试更高值
    processes = []
    
    for param in params:
        p = mp.Process(target=run_experiment, args=(param,))
        p.start()
        processes.append(p)
        
        # 控制并发数，避免进程过载
        if len(processes) >= num_p:
            for p in processes:
                p.join()
            processes = []
    
    # 等待剩余进程完成
    for p in processes:
        p.join()

2. 降低CUDA上下文初始化开销 #

spawn启动每个进程都会重新初始化CUDA上下文，这是首次任务耗时变长的核心原因。可尝试两种优化方式：

• 改用forkserver启动方法（需确保主进程未提前初始化CUDA）：

if __name__ == '__main__':
    mp.set_start_method('forkserver')
    # 后续进程启动逻辑

• 若模型极小，可尝试fork方法（CUDA官方不推荐，但部分场景下能避免重复初始化）：

if __name__ == '__main__':
    torch.cuda.init()  # 主进程提前初始化CUDA
    mp.set_start_method('fork')
    # 后续进程逻辑

3. 批量调度任务，减少进程创建销毁成本 #

将多个实验参数打包成批次，让单个进程处理多组实验，降低进程频繁创建销毁的开销：

def run_batch(params_batch):
    device = torch.device('cuda:0')
    for param in params_batch:
        agent = create_agent(param).to(device)
        # 运行实验
    return

if __name__ == '__main__':
    params = [...]
    batch_size = 5  # 每批处理的实验数量
    param_batches = [params[i:i+batch_size] for i in range(0, len(params), batch_size)]
    
    mp.set_start_method('spawn')
    num_p = 16
    with Pool(num_p) as p:
        p.map(run_batch, param_batches)

4. 排查环境交互的CPU瓶颈 #

强化学习的环境交互（如游戏环境）通常是CPU密集型，可能导致GPU闲置。可优化方向：

• 用多线程或numba加速环境交互逻辑
• 减少CPU与GPU间的数据传输，仅将模型计算相关数据移至GPU

5. 实时监控负载定位瓶颈 #

用nvidia-smi查看GPU显存和利用率，用top查看CPU负载：

• 若GPU利用率低但CPU负载高，瓶颈在环境交互
• 若两者均低，说明进程调度或CUDA上下文开销过大

内容的提问来源于stack exchange，提问作者q.uijote