PufferLib 5.0 工程复盘：单 GPU 每秒 6000 万步的强化学习训练是怎么炼成的

一次「反向扩展」的极致工程

这是 PufferLib 5.0 发布配套的深度工程文章：这个强化学习库把单 GPU 有用训练吞吐做到每秒 6000 万步。在 4.0 版本，我们用一套自研 CUDA C 技术栈、历经 8 个月艰苦工程打破了旧格局；5.0 则是打磨掉所有毛边的精修版本。

线性扩展：小模型的性能反而没人管

PufferLib 的性能可以线性扩展到 10 万参数以下。峰值端到端训练吞吐大致是：10 万参数约每秒 5000 万步，100 万参数约 1000 万步，1000 万参数约 180 万步。

能意识到这有多疯狂的读者恐怕不多，所以让我把这件事放进语境里。硅谷所有拿着高薪的模型工程师都在做「扩展向上」的训练。其中一些挑战是直觉性的：GPU 越多，通信开销越大；权重越多，就越要精打细算每个计算放在哪里。问题不断叠加：先是模型装不进单卡，然后装不进单节点，接着单个层都开始放不下。规模越大，问题越多。

但反方向呢？当你尝试「扩展向下」会发生什么？你可以在任何使用 1000 万参数以下模型的机器学习项目里自己试一试：在标准框架下把模型规模减半，训练速度并不会翻倍。低于某个点之后，训练几乎不再变快。标准框架没有为小模型优化，每个操作都带一点开销——模型足够大时这些开销可以忽略，但小模型的 GPU 会有 95% 的时间在空转，因为开销占了主导。

在 1000 万参数量级，如果你的架构不需要自定义算子，Torch/Jax 还能有不错的性能；到 100 万参数，所有环节都需要重度融合的自定义算子；到 10 万参数——你需要 PufferLib。5.0 发布还有另外三篇文章讲关键改进：异步训练、GPU 环境支持、与算法简化的协同优化。本文则带你走一遍我的思考过程，以及做这个版本时遇到的各种权衡。

用打砖块当基准测试

我们在传播一些最大工程成就时有个难题：一个任务不到 1 秒就被解决，反而让怀疑者觉得整件事肯定很 trivial。打砖块（Breakout）一开始并不是玩具任务：是我们的进步把它变成了玩具任务。以下是历史上解决（甚至还没完全解决）这个任务所需的时间：

未解决阶段（300-400 分）：Nature DQN 需要好几天；早期 PPO 实现需要好几个小时；CleanRL + EnvPool PPO 需要 30 分钟；PufferLib 1.0 也是 30 分钟。

解决（850+ 分，Ocean 版环境）：PufferLib 2.0 用 5 分钟——100 倍吞吐（每秒 100 万步）；3.0 用 30 秒——训练器与算法改进（400 万步/秒）；4.0 用 4 秒——2000 万步/秒的 CUDA C 加 PufferNet 架构；5.0 不到 1 秒——6000 万步/秒的异步 GPU 实现。

在 2.0 里，我们把打砖块移植成基于状态的独立 C 环境，用 3 亿步（约 120 万个不同状态）在 5 分钟内解决。此后每个版本都在没有专门追求样本效率的情况下，同时改进了求解时间和样本效率。

另一种看待方式是：打砖块是我们仍然关心的最简单环境。不像 cartpole 和 pong 那种更简单的任务，打砖块的复杂度刚好够当一个有用的预测性基准：一个半残的算法通常解决不了它，而在打砖块上改善求解时间的算法/模型改动，通常能迁移到更复杂的环境。

研究中保持快速迭代周期的重要性怎么强调都不为过。30 秒和 10 分钟的区别，就是「能对想到的每个边缘情况都极度严谨地测试」和「基本只能凭感觉上」的区别。有了不到 1 秒的基准，你甚至能以合理的速度迭代外层循环。所以从 4 秒砍到 1 秒以内，实际上让我们在超参数扫描算法 Protein 上的开发变得可控得多：现在可以在单 GPU 上一小时内跑 1000 个实验。

基础设施优化

这个版本最大的性能变化是引入异步训练，但那其实是个很小的代码改动，不需要多少重新设计。真正的重活见我们的并行化与算法文章。这里我讨论的是为了精简库、改善环境支持所做的主要代码改动。

我喜欢用「屏障和边界」的方式思考代码：任何本质上把项目两个组件隔开的设计。大厂的思维是关注点分离是好事、能带来模块化可复用的代码。但我讨厌模块化。问题在于代码组件不是乐高积木。几乎不存在一个明显的边界，让你把功能切开后两块还能无缝拼回去而不带来额外臃肿。它更像把两个形状怪异的物体塞进两个过大的箱子里再摞起来。

抽象臃肿潜入我们的一个地方是向量化。环境此前由 gcc/clang 以 C99 编译，训练器由 NVCC 以 CUDA 编译。值得注意的是，我们的训练器完全用 .cu 文件写成，包括大量 CPU 侧 C99 代码——这是刻意设计，消除了通常需要为每个内核导出启动包装器的 CUDA 到 C 语言屏障。所以很自然地我们会问：能不能直接用 nvcc 编译环境？答案是：可以，而且相当轻松。有一两个常见语法怪癖，但大部分环境代码直接就能编译。它依然可以作为 C99 编译用于 CPU 演示，不需要维护两套版本。这个认知带来两大简化：第一，PufferLib 可以编译成单一单元、对着一个环境构建。不再有混乱的 C 到 CUDA 绑定，不再需要把常量定义从一边走私到另一边。对用户来说，现在写环境只需写一个文件，动作空间大小之类的直接在顶部定义成常量。第二，它顺手带来了 GPU 环境支持，我们在几个演示里用到了。

我喜欢这个例子，因为它展示了为一个小小的可用性改进要付出多少工作。2.0 你要写 Cython 绑定文件；3.0 写 C 绑定文件（4.0 又精简了一点）；5.0 你只写一个环境文件、导出几个常量。基础设施侧，我们借这次改动清理了向量化代码和构建脚本，还基本消除了 PufferLib 的启动时间——之前相当于约 10 个 epoch 的模型前向，是 Torch 留下的遗留问题（Torch 因为某种可怕的原因需要多次编译预热）。我们的做法是直接以捕获模式运行第一个 epoch 的有效工作，不需要快照一堆状态变量再在预热后恢复。

另一个大改动来自一封随机邮件。一位叫 Simon Zehnder 的读者发来了一些比 Puffer 更快的 Triton GAE 内核。这个优化对我影响不大（不到总时间的 1%，还增加不少路径专属内核代码），但它并行化加载的方式很有意思，所以我在 MinGRU 上试了试——结果长时序反向传播的配置端到端吞吐提升 10% 到 30%。

我们还做了一大堆削减几个百分点的小内核改动。其中较大的两个在 MinGRU 里。这个线性递归在 log 空间计算一些不稳定中间量，因为数值敏感，我们一直用 fp32 存储。我意识到可以找一个更稳定、能用 bf16 存储的中间量……结果发现那个量就是状态本身。

配合这次更新的算法简化（见算法文章），我们在更小的模型上又砍掉了 5% 到 10% 的开销。

彻底消灭 Python

所有性能关键基础设施在 4.0 就已移植到 CUDA C，当时剩下约 2000 行 Python 胶水代码。你可能会觉得留一些 Python 高层工具是有用的，把剩下的移植完纯属面子工程。但在紧凑项目里开发不是这样运作的。这么想：每加一段代码，你都要决定它住在 Python 还是 CUDA。在这个边界上来回搬东西不轻松——至少不像同语言里直接调函数那样轻松。导出绑定上总有摩擦，类型转换上也是。这个屏障让清理代码变得困难得多。同语言里一个函数调用另一个，你可以把被调函数体粘进调用方，直接看有没有冗余。听着笨，但这正是我每次重构的起点，而且我总能把提交的 PR 缩短——哪怕是资深工程师的。重构时我找的就是屏障：任何把一段逻辑与另一段隔开的东西。多语言项目天生更多这种屏障，所以在你愿意投入的固定清理精力下，膨胀一定比单语言项目多。

你大概能看出来我真的很在乎代码简洁。但单语言的理由不止于此。Python 依赖 frankly 就是一场彻头彻尾的噩梦。3.0 我们有 17 个依赖，4.0 砍到 9 个。但你知道比 9 个更好的是什么吗？零个。尤其考虑到最近的依赖腐烂速度。大部分替换都不难：各种性能分析库只是系统调用的包装，C/CUDA 里本来就能直接用。所以这里只讲几块大的。

gpytorch：这块最痛苦。我们的超参扫描算法 Protein 用到高斯过程。幸运的是，一位在该领域背景深厚的贡献者接手了移植任务，做了个最小化的 CUDA 版本。最初的 PR 约 3500 行，他后来缩到 2500 行，我花大约一天压到 1300 行。这是相对 4.0 最大的代码增量，但想想：1300 行换一个更快、确定性、自包含的扫描工具，甩掉一个重型 Torch 依赖——完全不亏。

Python 内置 .ini 解析器：PufferLib 生命周期里我大概在配置系统上花了上百小时。它们非常重要：做错了就会同时存在 3 种配置和文件格式，日志上堆积大量数据膨胀。4.0 我们用 .ini 做配置、json 做日志、还有个给 constellation 的奇怪聚合 json。5.0 我们把一切统一成单一格式——扩展了列表支持的 .ini。解析器约 500 行，但比 Protein 还难收紧。训练、扫描、CPU 评估演示、constellation 各自需要用略微不同的方式解析 .ini 数据。最初这些半重复/臃肿的实现加起来约 1500 行。

rich：所有 Python 依赖里，我最不想替换的就是它，因为 Rich 确实很好。我们用它画终端里的训练仪表盘。我不愿意接受移植到 C 就意味着忍受一个更简陋、更丑的仪表盘，而移植 Rich 这种量级的东西是个大工程。所以我们转而硬编码自己的仪表盘布局，然后尽力收紧代码。Python 版用 rich 是 110 行，我压到 250 行以内都做不到——这本身就说明 rich 有多好。如果有快速单文件 C 实现，我直接 vendored 了事。

大概就是这些。这个版本还有一大块代码用于支持自我对弈，但在 C/CUDA 里并不比 Python 里长多少。大部分移植工作从 AI 加验证开始，但（这点怎么强调都不过分）随后是数周由粗到细、越来越手动的清理。最初的纯 CUDA/C 版本超过 1.5 万行，基本不可读。

原文信息

作者：Joseph Suarez（@jsuarez），机器学习工程师、PufferLib 作者

原文地址：