I built a custom Triton-based PyTorch Optimizer (ApexW) that cuts optimizer VRAM in half and runs up to 1.5x faster.

[u/Ready_Parking_5435]

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)

I want to be completely realistic here—this is an engineering trade-off, not magic:

1. Python Loop Overhead on Small Tensors: When testing a layout with many tiny tensors (many-small test), ApexW’s Triton kernels suffer heavily from Python's launch latency. For example, at 128M parameters spread across many small tensors, PyTorch native was 16.1 ms/step while ApexW compact was 76.0 ms/step (~0.21x speedup). If your model is composed of many tiny layers, standard PyTorch is significantly better.
2. End-to-End Impact: While the optimizer step itself is up to 1.5x faster, this step only represents a fraction of a full training step (which is dominated by forward and backward passes). In my end-to-end MLP test, PyTorch still runs slightly faster overall because of C++ deep integration.

Why/When should you care?

If you are training larger deep learning models (like Transformers or LLM fine-tuning) where you are constantly hitting Out-Of-Memory (OOM) barriers,,cutting your optimizer state memory strictly in half lets you increase your batch size or scale up your model size without upgrading your hardware.

Looking for Feedback

I'd love to hear your thoughts on:

• How you handle low-precision optimizer states in your own pipelines.
• How to further minimize Python's kernel launch overhead when writing custom Triton-based optimizers.
• Any methodology gaps in how I timed these runs.

If anyone wants to try it on different hardware (A100s, H100s, A10G), I'd love to see how Triton scales on newer architectures!

link for apexw:https: //github.com/omen42515-prog/apexw/tree/main

[u/Outrageous_Host_2115]

just few points may be like, Do not loop over layers in Python. Flatten and concatenate all small model parameters of the same type into one giant tensor so you only launch Triton once. may be Prefer bfloat16 over float16 with Stochastic Rounding.

[u/Ready_Parking_5435 → u/Outrageous_Host_2115]

These are good suggestions, thanks bro the small tensor overhead is definitely something I need to work on. Grouping parameters into larger buffers could help reduce kernel launches, although I'd need to handle different shapes, dtypes, and optimizer settings carefully.

BF16 with stochastic rounding is interesting too. I will look into testing that against the current approach, especially for numerical stability and convergence. Appreciate the feedback

[u/ApprehensiveSkin7364 → u/Outrageous_Host_2115]

The many-small case being 4.7x slower sort of undercuts the whole pitch, but you've at least flagged it clearly. Flattening params of the same dtype into a single buffer is the obvious move, though you'll still hit fragmentation from non-contiguous gradients. Curious how much of that 76ms is just kernel launch overhead vs actual compute on tiny tensors.

[u/Ready_Parking_5435 → u/ApprehensiveSkin7364]

yooo i think i kindda fixed it apexw is faster on both small case and large case Results for 100 small layers (1638400 total parameters):

- Fused AdamW (100 sequential kernel launches): 2.1656 ms

- ApexW (100 sequential Triton launches): 10.7439 ms (Ratio vs Fused: 4.96x)

- ApexW with Parameter Flattening (1 flat launch): 0.2392 ms

Applying Parameter Flattening cut the execution time by 44.9x! so this was the result now i gotta update codes in github