# 我为一个面向大规模 GPU 工作负载的 PyTorch 优化器做了实现——希望获得反馈 \[P\] - 原标题：I built a PyTorch optimizer for large GPU workloads — looking for feedback\[P\] - 类型：帖子 - 来源：Reddit - 发布时间：2026-10-27 - 首次精选：2026-10-29 01:09 - 分级：值得看 - 学习层级：深入 - 主题：[inference-hardware-memory]()（第 17 位）、[model-serving-compute]()（第 30 位） - 原文链接：[https://reddit.com/r/MachineLearning/comments/1xgqtxp/i\_built\_a\_pytorch\_optimizer\_for\_large\_gpu]() - 本站页面：[https://oplearnx.com/resources/pk\_a8455e0a2003]() ## 摘要 这是 r/MachineLearning 的一篇社区帖子（\[P\] 表示项目分享），作者介绍自己开发的 PyTorch 优化器 ApexW，目标是降低优化器状态显存占用并加快较大 GPU 工作负载上的优化器更新步骤，做法是紧凑的优化器状态加一条基于 Triton 的执行路径。作者称在只有几个大张量时的优化器单项基准中，Triton 加紧凑状态路径比 PyTorch 的 fused AdamW 快约 1.5 倍，优化器状态显存约减半；同时说明自动执行模式尚未稳定取得同样加速，含大量小张量的负载明显慢于 PyTorch fused 优化器（例如 1.28 亿参数分散在众多小张量时，PyTorch 原生为 16.1 ms/步），并强调这些数字来自优化器单项测试而非完整模型训练，希望有人在实现与基准测试上把关。适合从事 PyTorch 训练优化、关注优化器状态显存与 Triton 内核实现的人阅读。 ## 原帖 > 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. — Ready\_Parking\_5435 · 2026-10-27 · [原帖]() ## 中文大意 发帖人介绍自己写的 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 或参数分组缓解。（[出处]()） - 如果是自动调度路径本身有问题，就需要重新考虑它选择融合路径的方式。（[出处]()） --- 来源：LearnMap · 学径（https://oplearnx.com）