ApexW PyTorch optimizer targets large GPU workloads
A developer shared ApexW, a PyTorch optimizer that uses compact optimizer states and a Triton-based execution path for large GPU workloads. In the developer's optimizer-only tests, the Triton compact-state path was about 1.5 times faster than PyTorch's fused AdamW on workloads with a few large tensors while using roughly half the optimizer-state memory. The developer said these results did not represent full model training, that automatic mode did not consistently achieve the same speedup, and that workloads with many small tensors were much slower than PyTorch's fused optimizer. The developer also reported that PyTorch was slightly faster overall in an end-to-end MLP test.
ApexW demonstrates a developer-tested memory and speed trade-off for PyTorch training with large tensors, but its optimizer-level gains did not translate into an overall speedup in the developer's end-to-end MLP test.
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)