Fused AdamW Optimizer

• Fused AdamW optimizer is an adaptive method that fuses parameter updates into the backward pass, reducing memory traffic and eliminating sequential bottlenecks.
• It interleaves gradient computation with parameter updates to expose fine-grained parallelism and improve data locality, enhancing GPU performance.
• Empirical studies show a 1.15–1.20× speedup in image classification tasks without altering convergence behavior, making it a drop-in replacement for standard AdamW.

A fused AdamW optimizer is an implementation of the AdamW update rule in which optimizer logic is “fused” directly into the backward pass of neural network training. Instead of executing a separate optimizer step after the forward and backward computations, the fused strategy updates parameters and optimizer states the moment each gradient element is produced. This backward-fusion methodology enhances data locality, halves memory traffic, removes sequential bottlenecks inherent to standard eager-mode frameworks, and exposes additional fine-grained parallelism, all without altering the mathematical optimizer algorithm or affecting convergence properties. Empirical results demonstrate that this strategy yields 1.15–1.20× end-to-end speedups in typical image classification workloads on modern GPUs (Jiang et al., 2021).

1. Standard AdamW Update Formulation

AdamW is a widely used adaptive moment-based optimizer with decoupled weight decay. For parameters $\theta_t \in \mathbb{R}^D$ at iteration $t$, mini-batch gradients $g_t = \nabla_\theta \ell(\theta_t)$, first and second moment decay rates $0 < \beta_1, \beta_2 < 1$, learning rate $\eta$, weight-decay coefficient $\lambda$, and $\epsilon$ for numerical stability, the canonical update rules are: $m_t = \beta_1 m_{t-1} + (1 - \beta_1) g_t$

$$v_t = \beta_2 v_{t-1} + (1 - \beta_2) g_t \odot g_t$$

$$\hat{m}_t = m_t / (1 - \beta_1^t), \quad \hat{v}_t = v_t / (1 - \beta_2^t)$$

$t$0

This decouples weight decay from the adaptive step and requires separate optimizer logic after gradient accumulation (Jiang et al., 2021).

2. Fused Backward-Mode AdamW: Reordering and Interleaving

A conventional eager-mode training workflow executes: (A) a forward pass to read $t$1, (B) a backward pass accumulating $t$2, and (C) a separate optimizer update pass reading $t$3, $t$4, $t$5, and $t$6 again and writing back updates. This results in redundant data movement and prevents overlapping optimizer updates with ongoing gradient computations.

The fused AdamW approach merges the (B) backward pass and (C) optimizer pass into a single sweep. As soon as the gradient $t$7 for parameter $t$8 is computed, the corresponding moment buffers $t$9, $g_t = \nabla_\theta \ell(\theta_t)$0 are updated, bias correction (optional) applied, the parameter $g_t = \nabla_\theta \ell(\theta_t)$1 incremented, and all new states written back—in one tight loop over elements. There is no dependency across indices, so updates can be interleaved and parallelized efficiently.

3. Fused-AdamW Kernel: Pseudocode

The following gives the canonical fused-AdamW routine. The backward automatic-differentiation engine invokes this per parameter block, passing pointers to gradient, weights, first- and second-moment buffers and scalar hyperparameters: $g_t = \nabla_\theta \ell(\theta_t)$9 Each optimizer buffer ($g_t = \nabla_\theta \ell(\theta_t)$2) is read and written exactly once per element, and the arithmetic is fused so all operands remain in L1/L2 cache or register (Jiang et al., 2021).

4. Data Locality and Parallelism Benefits

Locality

In the baseline (non-fused) scheme, each buffer for $g_t = \nabla_\theta \ell(\theta_t)$3, $g_t = \nabla_\theta \ell(\theta_t)$4, $g_t = \nabla_\theta \ell(\theta_t)$5, and $g_t = \nabla_\theta \ell(\theta_t)$6 is read and written at least twice per iteration. For 4-byte elements, this results in 16 bytes per element. The fused kernel reduces this to exactly one read and one write per buffer (8 bytes per element), cutting memory traffic by half.

Parallelism

Standard eager execution enforces a strict sequential order: forward $g_t = \nabla_\theta \ell(\theta_t)$7 backward $g_t = \nabla_\theta \ell(\theta_t)$8 optimizer. Fused-AdamW eliminates this bottleneck: as soon as the last use of a parameter in gradient accumulation completes, its update can be issued in parallel with gradient computation for other parameters, increasing GPU thread-block concurrency. This exposes more wavefronts for scheduling, improving hardware utilization.

Roofline Argument

On contemporary GPUs, separate AdamW kernels are frequently memory-bandwidth bound. By halving traffic, the fused kernel moves training closer to the compute-flop limit, boosting effective FLOP/s and reducing total iteration wall-clock time (Jiang et al., 2021).

5. Empirical Speedup, Scalability, and Memory Bandwidth Impact

On an ImageNet classification task (MobileNetV2, batch size 32, Pascal Titan Xp, float32 precision), baseline eager AdamW took 98.8 ms/iter. Fused-AdamW backward-mode achieved 83.0 ms/iter (1.19× speedup, saving 15.7 ms), while forward-fusion yielded 84.5 ms/iter (1.17× speedup) (Jiang et al., 2021).

Memory-Bandwidth Counters

• Baseline optimizer step: ~16.7 ms and 180 GB/s
• Fused step: ~0 ms in separate kernel; total backward bandwidth drops to ~95 GB/s

Scalability Across Batches and Models

Sweeping batch sizes from 1 to 200, the absolute time saved (~15 ms) is constant, since AdamW's elementwise cost is batch-independent and gradient computation cost dominates at large batch sizes. Across other models (ResNet-18/50/152, DenseNet-201, VGG19_BN, batch 32), backward-fusion consistently gave 1.05–1.18× speedups, with relatively larger benefits on architectures with small parameters per layer (e.g., MobileNetV2).

Mixed Precision

On mixed-precision (fp16) using NVIDIA Apex, fused backward-pass delivers an additional 5–10% gain over Apex's fused-optimizer kernels, since Apex does not overlap optimizer update with gradient computation; backward-fusion closes that efficiency gap.

6. Summary and Significance

By fusing AdamW's update logic into the backward pass, memory traffic is halved, kernel launches are reduced, and new opportunities for parallel execution are exposed. These gains translate into consistent 1.15–1.20× end-to-end training accelerations on modern GPUs without changing optimizer convergence behavior or accuracy. The fused-AdamW framework remains a drop-in replacement for standard AdamW in any training workflow, as it modifies the execution schedule but not the optimizer algorithm itself (Jiang et al., 2021).