十年匠心定制 · 商业建站与技术教学双线并行 咨询热线:400-886-1026 service@lmnt.cn
ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

CleanRL 多GPU分布式训练实战:用 torchrun 给 Atari PPO 提速

CleanRL 多GPU分布式训练实战:用 torchrun 给 Atari PPO 提速 CleanRL 多GPU分布式训练实战用 torchrun 给 Atari PPO 提速【免费下载链接】cleanrlHigh-quality single file implementation of Deep Reinforcement Learning algorithms with research-friendly features (PPO, DQN, C51, DDPG, TD3, SAC, PPG)项目地址: https://gitcode.com/GitHub_Trending/cl/cleanrl在 Atari 上单卡跑 PPO1000 万步要占住一张卡好几天排期被环境采样拖垮。CleanRL 的 ppo_atari_multigpu.py 用 torch.distributed 做多GPU分布式训练面向会用 PyTorch 和 gym、想压缩训练墙钟时间的工程师每张卡独立采样梯度跨卡同步。下文按启动命令、收益数据、调参排障展开。多GPU训练值不值得先看基准数据官方基准的结论同硬件下 ppo_atari_multigpu.py 比单卡 ppo_atari.py 快约 30%样本效率不损失。三个标准 Atari 任务上两者回报均值基本重叠Breakout 多GPU 414.94 ± 20.60单卡 414.66 ± 28.09BeamRider 多GPU 2414.65 ± 643.74单卡 1915.93 ± 484.58。也就是说多卡不是用回报换速度收益来自采样吞吐环境越吃并行多GPU训练越划算。多GPU训练启动命令与关键参数装好 Atari 依赖后用 torchrun 启动即可官方文档同时提示该脚本不支持 Windows 与 macOSpip install -r requirements/requirements-atari.txt torchrun --standalone --nnodes1 --nproc_per_node2 \ cleanrl/ppo_atari_multigpu.py \ --env-id BreakoutNoFrameskip-v4 --backend nccl参数含义可直接查 cleanrl/ppo_atari_multigpu.py 中的参数表。--nproc_per_node决定子进程数world size第 i 个进程默认用 cuda:{i}--local-num-envs是每进程环境数总环境数 local_num_envs × world sizebatch 大小与迭代次数都由它推导--num-steps默认 128是每个环境每轮 rollout 的步数--backend默认 gloo多卡应换 nccl--device-ids手动指派显卡--device-ids 0 0 1 1表示每张卡挂两个进程。数据并行与多进程环境并行怎么协同可以把它理解成多组数据采集工各组各跑各的环境采样本先做本地训练步每个更新步末对账——把梯度取平均再写回参数各卡因此始终持有同一份模型版本。代码里torchrun 启动 N 个进程每个进程持有一份模型用 SyncVectorEnv 跑自己的 local_num_envs 个环境采满 local_num_envs × num_steps 的本地 batch。一致性靠两处细节权重初始化用同一个 torch 种子各卡起点相同随机状态用 args.seed local_rank 错开各进程的环境随机性不相关。关键在梯度同步——backward 后把所有参数梯度拼成一个张量一次 all_reduceSUM再除以 world size 写回同步逻辑如下摘自源码# 所有参数梯度拼成一个张量一次 all_reduce替代逐参数通信 all_grads torch.cat([p.grad.view(-1) for p in agent.parameters() if p.grad is not None]) dist.all_reduce(all_grads, opdist.ReduceOp.SUM) offset 0 for p in agent.parameters(): if p.grad is not None: p.grad.data.copy_(all_grads[offset : offset p.numel()].view_as(p.grad.data) / args.world_size) offset p.numel()随后照常做梯度裁剪与 optimizer.step()。TensorBoard 只由 rank 0 写global_step 每步按总环境数递增曲线可直接与单卡运行对比。单卡与多卡基准结果怎么读环境单卡 ppo_atari.py多卡 ppo_atari_multigpu.py结论PongNoFrameskip-v420.36 ± 0.2020.34 ± 0.43均值持平波动略大BeamRiderNoFrameskip-v41915.93 ± 484.582414.65 ± 643.74多卡均值更高BreakoutNoFrameskip-v4414.66 ± 28.09414.94 ± 20.60均值持平波动更小读表先看方差是否重叠而不是均值谁大这些是多次运行的均值方差重叠即样本效率不损失这也是官方基准的校验标准。再看墙钟时间官方结论是同硬件快约 30%原始数据见 多GPU基准测试。调优与排障显存不足、通信瓶颈与负载不均显存不足。现象启动或首个 rollout 即 CUDA out of memory。原因rollout 存储obs、actions、rewards、values 等张量全量驻留显存体积正比于 num_steps × local_num_envs。处理先降--local-num-envs再降--num-steps卡数不动。通信瓶颈。现象SPS每秒步数不随卡数增长。原因默认 gloo 后端走 CPU 通信多卡梯度同步成为瓶颈。处理--backend nccl。训练不稳。现象回报曲线剧烈抖动、explained_variance 偏低。原因每卡环境数太少、本地 batch 偏小脚本已按 args.seed local_rank 为各进程错开种子一般无需手动处理。处理调大--local-num-envs必要时降--learning-rate。GPU 负载不均。现象nvidia-smi 里一张卡满载、另一张接近空闲。原因进程数超过卡数时脚本回退到不带索引的 cuda多进程挤到同一张卡。处理用--device-ids逐个进程指定卡。ppo_atari_multigpu.py 给 PPO 这类 on-policy 算法提供了直接的多GPU训练路径每卡环境并行采样加梯度 all_reducePPO 超参数一个不动。如果瓶颈是单卡跑 Atari 或 Procgen 的墙钟时间它是第一选择。本地/全局参数推导与批量梯度同步的完整实现见 PPO 官方文档。【免费下载链接】cleanrlHigh-quality single file implementation of Deep Reinforcement Learning algorithms with research-friendly features (PPO, DQN, C51, DDPG, TD3, SAC, PPG)项目地址: https://gitcode.com/GitHub_Trending/cl/cleanrl创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表