--- title: Fused AdamW Optimizer url: https://www.emergentmind.com/topics/fused-adamw-optimizer type: topic --- # Fused AdamW Optimizer 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 [2104.00237]. ## 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) \] \[ \theta_{t+1} = \theta_t - \eta \cdot \hat{m}_t / (\sqrt{\hat{v}_t} + \epsilon) - \eta \cdot \lambda \cdot \theta_t \] This decouples weight decay from the adaptive step and requires separate optimizer logic after gradient accumulation [2104.00237]. ## 2. Fused Backward-Mode AdamW: Reordering and Interleaving A conventional eager-mode training workflow executes: (A) a forward pass to read $\theta$, (B) a backward pass accumulating $\partial \ell/\partial \theta$, and (C) a separate optimizer update pass reading $\theta$, $m$, $v$, and $\partial \ell/\partial \theta$ 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 $g_t[i]$ for parameter $\theta[i]$ is computed, the corresponding moment buffers $m_{t}[i]$, $v_{t}[i]$ are updated, bias correction (optional) applied, the parameter $\theta[i]$ 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: ```python def BACKWARD_FUSED_ADAMW(θ_ptr, g_ptr, m_ptr, v_ptr, β1, β2, ε, η, λ, t): one_minus_β1 = 1 - β1 one_minus_β2 = 1 - β2 inv_bias1 = 1 / (1 - β1**t) # optional inv_bias2 = 1 / (1 - β2**t) # optional N = len(θ_ptr) # Elementwise parallel execution: parallel_for i in range(N): g = g_ptr[i] θ = θ_ptr[i] m = m_ptr[i] v = v_ptr[i] m_new = β1 * m + one_minus_β1 * g v_new = β2 * v + one_minus_β2 * (g * g) m_hat = m_new * inv_bias1 v_hat = v_new * inv_bias2 update = η * (m_hat / (v_hat**0.5 + ε) + λ*θ) θ_new = θ - update θ_ptr[i] = θ_new m_ptr[i] = m_new v_ptr[i] = v_new g_ptr[i] = 0 # zero-out gradient ``` Each optimizer buffer ($\theta, m, v, g$) is read and written exactly once per element, and the arithmetic is fused so all operands remain in L1/L2 cache or register [2104.00237]. ## 4. Data Locality and Parallelism Benefits ### Locality In the baseline (non-fused) scheme, each buffer for $\theta[i]$, $m[i]$, $v[i]$, and $g[i]$ 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 $\to$ backward $\to$ 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 [2104.00237]. ## 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) [2104.00237]. ### 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 [2104.00237]. Source: https://www.emergentmind.com/topics/fused-adamw-optimizer