我为一个面向大规模 GPU 工作负载的 PyTorch 优化器做了实现——希望获得反馈 [P]
I built a PyTorch optimizer for large GPU workloads — looking for feedback[P]
这是 r/MachineLearning 的一篇社区帖子（[P] 表示项目分享），作者介绍自己开发的 PyTorch 优化器 ApexW，目标是降低优化器状态显存占用并加快较大 GPU 工作负载上的优化器更新步骤，做法是紧凑的优化器状态加一条基于 Triton 的执行路径。作者称在只有几个大张量时的优化器单项基准中，Triton 加紧凑状态路径比 PyTorch 的 fused AdamW 快约 1.5 倍，优化器状态显存约减半；同时说明自动执行模式尚未稳定取得同样加速，含大量小张量的负载明显慢于 PyTorch fused 优化器（例如 1.28 亿参数分散在众多小张量时，PyTorch 原生为 16.1 ms/步），并强调这些数字来自优化器单项测试而非完整模型训练，希望有人在实现与基准测试上把关。适合从事 PyTorch 训练优化、关注优化器状态显存与 Triton 内核实现的人阅读。
- 分级
- 值得看
- 收录时间
- 2026-10-29
原帖
In my current tests, the Triton + compact-state path was around 1.5x faster than PyTorch's fused AdamW on workloads with a few large tensors, and the optimizer-state memory usage was roughly cut in half.
中文大意
发帖人介绍自己写的 PyTorch 优化器 ApexW，目标是降低优化器状态显存占用、加快大 GPU 负载下的优化器步进，做法是紧凑优化器状态加 Triton 执行路径。在只有少量大张量的场景下，比 PyTorch fused AdamW 约快 1.5 倍，优化器状态显存约减半。但他明确说明这些只是纯优化器基准而非完整训练；自动执行模式尚不能稳定复现加速；小张量特别多的场景反而明显更慢（128M 参数分散在许多小张量时，PyTorch 原生 16.1 ms/step，ApexW compact 76.0 ms/step，约 0.21x），端到端 MLP 测试中 PyTorch 整体仍略快。他认为价值在于频繁撞上 OOM 的大模型训练（如 Transformer、LLM 微调），显存减半可换来更大 batch 或模型规模。他附上 GitHub 链接，希望有人帮忙看实现与基准设置，并问有没有…
讨论区核心观点
依据 2 条评论
- 发帖人承认还没分清小张量场景的减速有多少来自 Python 端循环、多少来自 kernel 启动/调度开销，暂不下结论。 出处（新标签页打开）
- 下一步打算剖析混合尺寸场景，并比较 eager 执行与 CUDA Graph 重放；若主要是启动开销，可用 CUDA Graph 或参数分组缓解。 出处（新标签页打开）
- 如果是自动调度路径本身有问题，就需要重新考虑它选择融合路径的方式。 出处（新标签页打开）