Developer Builds Custom PyTorch Optimizer ApexW: Half the Memory, 1.5x Faster

A developer shared a PyTorch optimizer called ApexW that he built, targeting optimizer-state memory usage and update speed under large-GPU workloads. According to the author, in optimizer-only benchmarks, ApexW—using a Triton execution path plus a compact state representation—is about 1.5x faster than PyTorch fused AdamW in "few large tensors" scenarios, with optimizer-state memory usage roughly halved. But he also admits that the automatic execution mode cannot yet reliably reproduce this speedup, and that "many small tensors" workloads are noticeably slower than PyTorch's fused optimizer.

Compact State and Triton Kernels: The Two Technical Pillars of ApexW

ApexW's design centers on two things: compressing optimizer state, and rewriting the execution path with Triton. Adaptive optimizers like AdamW need to maintain first- and second-order momentum for each parameter; the larger the parameter count, the more significant the memory this state occupies. The compact state representation attempts to reduce this overhead, while the Triton kernel fuses element-wise update operations into fewer GPU launches, thereby reducing scheduling and memory-access overhead.

Neither direction is a new concept, but combining them into a drop-in optimizer implementation and publishing benchmark results is what makes this project worth attention. The author doesn't stop at the slogan "I wrote a faster optimizer" but proactively lists the failure scenarios—something not often seen in similar shares.

AdamW and optimizer state: AdamW is the decoupled weight decay variant of the Adam optimizer, widely used for training deep networks such as Transformers. It needs to maintain two sets of sliding statistics with the same shape as the parameters—the first moment (momentum, the exponential moving average of gradients) and the second moment (the exponential moving average of squared gradients)—usually stored in FP32. For a 7-billion-parameter model, these two states alone require about 56 GB of memory, far exceeding the model parameters themselves (about 14 GB when stored in BF16). The so-called "compact state representation" essentially reduces the precision of these two momentum tensors or compresses their storage using blocking/quantization, thereby cutting this fixed overhead; the cost is a possible loss of update precision, so it needs to be verified whether convergence is affected in real training.

Triton kernels and kernel launch overhead: Triton is an open-source GPU programming language and compiler from OpenAI; you write CUDA-like kernels in Python, and they can be compiled and executed across NVIDIA and AMD GPUs. Optimizer updates are element-wise operations; if each tensor calls a native PyTorch operator once, it produces a large number of independent kernel launches, each with scheduling and synchronization overhead. The idea of a fused optimizer is to merge multiple element-wise operations into a single kernel, reducing launch count and memory round-trips. This also explains the phenomenon of being faster on "few large tensors" and slower on "many small tensors": the fixed overhead brought by fusion can only be amortized on large tensors.

The Premise of the 1.5x Speedup: Only the Optimizer Was Tested, Not Full Training

A key caveat that must be emphasized is that these numbers come from optimizer-only benchmarks, not full model training. The optimizer step accounts for only part of the time in end-to-end training; how much benefit actual training gains depends on the ratio of time among forward pass, backward pass, and optimizer update. If the optimizer itself accounts for a small share, even a 1.5x faster single step will be diluted in overall training speedup.

On the other hand, "few large tensors" is a relatively narrow sweet spot. The parameter matrices in large language models or large-scale vision models do fit this characteristic, but the parameter distribution of real networks is often mixed with a large number of small tensors such as normalization layers and bias terms—precisely ApexW's current weak spot.

What optimizer-only benchmarks mean: Such benchmarks only loop over the optimizer step and do not include forward or backward propagation, so what is measured is the pure time cost of the optimizer step. In typical training, the optimizer update usually accounts for only 5%–15% of total single-step time (depending on model and batch), with forward and backward taking the lion's share. If the optimizer itself is 1.5x faster, then when the optimizer accounts for 10%, the end-to-end speedup is about 1/(0.9+0.1/1.5)≈1.03, i.e. only about 3%. Therefore, when evaluating such results, you must convert them based on the actual proportion of the optimizer step in the profile; you cannot equate 1.5x directly with training speedup.

PyTorch fused AdamW baseline: Since 2.x, PyTorch has provided an AdamW implementation with fused=True, completing parameter updates in a single CUDA kernel, and it is already the official optimized path. Any new optimizer comparing against it should do so under the same precision state and the same compilation options (such as torch.compile); otherwise the speedup may come from an improperly configured baseline rather than the algorithm itself.

Small-Tensor Scenarios and Automatic Mode: Two Unresolved Weaknesses

The author explicitly states that the automatic execution mode cannot yet reliably achieve the same speedup ratio. This means the current version may require manually specifying the execution path to trigger the best performance, adding a layer of tuning cost for ordinary users.

Let's also mention a detail: performance degradation in small-tensor scenarios. Under workloads with many small tensors, ApexW is noticeably slower than PyTorch's fused optimizer. One possible explanation is that each tensor requires an independent kernel launch or state handling, and the fixed overhead of the compact state cannot be amortized on small tensors. The author himself admits that this is not a solution that can unconditionally replace AdamW.

Automatic execution mode and manual path selection: The so-called automatic execution mode means the optimizer automatically decides, based on heuristic rules such as parameter shape and count, whether to use the Triton fused kernel or fall back to the regular implementation. If the heuristic judgment is inaccurate, it may incorrectly enable the fusion path on small tensors, or fall back when fusion should be used, leading to unstable performance. Manually specifying the execution path can bypass this uncertainty, but it requires users to understand their own workload characteristics, increasing the tuning threshold. This also suggests that for a drop-in optimizer to be truly usable, the robustness of the automatic scheduling strategy is as important as performance.

What to Watch in Community Feedback: Benchmark Setup and Missing Tests

In the post, the author proactively solicits two types of feedback: implementation-level issues, and benchmark setup-level issues. He particularly wants to know whether the benchmark setup has obvious flaws and whether important tests are missing.

From an external perspective, several directions worth questioning include: whether the baseline uses the same precision state representation, whether equivalent optimizations such as torch.compile are enabled, whether timing includes kernel launch and synchronization overhead, and whether memory statistics measure peak or steady state. These details directly determine whether the two numbers—1.5x and "halved"—can be reproduced. In addition, validation across multiple GPU architectures has not yet been completed; the author says he hopes to continue testing on different GPUs and more realistic models.

Metrics to check when reproducing benchmarks: Memory statistics should distinguish peak (peak allocated/reserved) from steady-state usage; PyTorch can measure them with torch.cuda.max_memory_allocated and memory_reserved respectively. Timing needs synchronization after each iteration (torch.cuda.synchronize) and several warm-up steps to exclude the effects of asynchronous kernel launches and first-time compilation. The baseline state precision must be aligned with ApexW's compact representation; otherwise "halved" may simply be a precision difference rather than an algorithmic gain. In addition, different GPU architectures (such as Hopper and Ampere) have different sensitivity to kernel fusion and memory bandwidth, so single-GPU results are not enough to represent general performance.

What It Means for Practitioners

For now, ApexW is better seen as an experimental project worth tracking rather than a production replacement for AdamW. If your training workload is dominated by large tensors and the optimizer step accounts for a relatively high proportion in the profile, it offers a possible optimization direction; if your model parameters are highly fragmented, its current performance may be worse than the built-in solution. The author's attitude of openly soliciting feedback also gives this project a foundation for continued iteration.

The project code has been released on GitHub, with the repository address github.com/omen42515-prog/apexw.