我为大规模 GPU 工作负载开发了一个 PyTorch 优化器——希望获得反馈[P]

嘿，我最近一直在开发一个叫 ApexW 的 PyTorch 优化器，想着把它分享到这里，听听比我更懂这方面的人的反馈。主要目标是减少优化器状态的内存占用，并看看在较大的 GPU 工作负载上能把优化器更新步骤加速多少。我一直在尝试紧凑的优化器状态和基于 Triton 的执行路径。在目前的测试中，对于包含少量大型张量的工作负载，Triton + 紧凑状态路径比 PyTorch 的融合版 AdamW 快约 1.5 倍，优化器状态的内存占用也大致减半。不过，结果并非全都理想。这些数据来自仅针对优化器的基准测试，而不是完整的模型训练。自动执行模式目前还无法稳定获得同样的加速，而且包含大量小张量的工作负载比 PyTorch 的融合优化器慢得多。所以，我绝不是在声称它能在所有工作负载中更好地替代 AdamW。我想在更贴近实际的模型和不同的 GPU 上进行测试，也希望有人能帮忙审视一下实现和基准测试。现实检验（取舍）我想在这里实事求是地说明——这是工程上的取舍，不是什么魔法：小张量上的 Python 循环开销：在测试包含大量微小张量的布局（many-small 测试）时，ApexW 的 Triton 内核受到 Python 内核启动延迟的严重影响。例如，当 1.28 亿个参数分散在大量小张量中时，PyTorch 原生实现每步耗时 16.1 毫秒，而 ApexW 紧凑版每步耗时 76.0 毫秒（加速比约为 0.21 倍）。如果你的模型由大量微小层组成，标准 PyTorch 明显更好。端到端影响：虽然优化器更新步骤本身的速度最高可达原来的 1.5 倍，但它只占完整训练步骤的一小部分（主要耗时在前向和反向传播上）。在我的端到端 MLP 测试中，由于 C++ 的深度集成，PyTorch 的整体运行速度仍然略快。你为什么应该关注它，又该在什么时候关注？如果你正在训练较大的深度学习模型（例如 Transformer 模型或进行大语言模型微调），并且经常遇到内存不足（OOM）的瓶颈，**，**将优化器状态的内存占用严格减半，就能让你在不升级硬件的情况下增大批量大小或扩大模型规模。GitHub：https://github.com/omen42515-prog/apexw/tree/main 如果你开发过 PyTorch 优化器、Triton 内核，或训练过较大的模型，我很想听听你的看法。尤其想知道，我的基准测试设置是否存在明显问题，或者是否遗漏了什么重要测试。