ApexW优化器面向大型GPU工作负载
一名开发者分享了ApexW，这是一款采用紧凑优化器状态和基于Triton的执行路径、面向大型GPU工作负载的PyTorch优化器。在开发者进行的仅优化器测试中，Triton紧凑状态路径处理少量大张量时，速度约为PyTorch融合AdamW的1.5倍，优化器状态内存用量则大致减半。开发者表示，这些结果并未反映完整的模型训练；自动模式无法稳定取得相同的加速效果，包含许多小张量的工作负载也比PyTorch融合优化器慢得多。开发者还称，在端到端MLP测试中，PyTorch整体速度略快。
ApexW展示了开发者测试所得的一种面向大张量PyTorch训练的记忆体与速度权衡，但其优化器级性能收益并未转化为开发者端到端MLP测试中的整体加速。
Hey, I've been working on a PyTorch optimizer called ApexW and figured I'd share it here to get some feedback from people who know more about this stuff than I do.
The main goal is to reduce optimizer-state memory usage and see how much faster I can make optimizer steps on larger GPU workloads. I've been experimenting with compact optimizer states and a Triton-based execution path.
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.
The results aren't all good, though. Those numbers are from optimizer-only benchmarks, not full model training. The automatic execution mode doesn't consistently get the same speedup yet, and workloads with lots of small tensors are much slower than PyTorch's fused optimizer. So I'm definitely not claiming this is a better AdamW replacement for every workload.
I want to test it on more realistic models and different GPUs, and I'd appreciate a second pair of eyes on both the implementation and the benchmarks. the Reality Check (The Trade-Offs)