_fused_adam: weight decay uses the AdamW rule, wrong results and wrong optimizer state for weight_decay != 0

[truong-v]

Summary

_fused_adam / _fused_adam_ produce wrong results whenever weight_decay != 0.

aten::_fused_adam is Adam, not AdamW. Weight decay must be an L2 penalty added to the
gradient before the moments are updated, so it also feeds exp_avg and exp_avg_sq:

g = g + weight_decay * p

The Triton kernel instead applies the decoupled AdamW rule straight to the parameter update,
so the moments never see the decay at all. The decoupled form belongs to the separate
aten::_fused_adamw op.

Both the updated parameter and the optimizer state come out wrong, and the error compounds over
a training run since exp_avg / exp_avg_sq are carried forward. With weight_decay = 0 the
two rules coincide, which is why this went unnoticed.

Reproduction

import torch
import flag_gems


def step(use_gems, weight_decay):
    torch.manual_seed(0)
    dev = flag_gems.device
    p = torch.randn(1024, device=dev)
    g = torch.randn(1024, device=dev)
    m = torch.randn(1024, device=dev)
    v = torch.randn(1024, device=dev).abs()
    mx = torch.zeros(1024, device=dev)
    t = torch.tensor([5.0], device=dev)  # fused adam expects float32 state steps

    args = ([p], [g], [m], [v], [mx], [t])
    kwargs = dict(lr=0.001, beta1=0.9, beta2=0.999, weight_decay=weight_decay,
                  eps=1e-8, amsgrad=False, maximize=False)
    if use_gems:
        with flag_gems.use_gems():
            torch.ops.aten._fused_adam_(*args, **kwargs)
    else:
        torch.ops.aten._fused_adam_(*args, **kwargs)
    return p, m, v


for wd in (0.0, 0.1):
    ref = step(False, wd)   # native PyTorch
    got = step(True, wd)    # FlagGems
    diffs = " ".join(f"{n}:{(a - b).abs().max().item():.2e}"
                     for n, a, b in zip(("param", "exp_avg", "exp_avg_sq"), ref, got))
    print(f"weight_decay={wd}: max abs err  {diffs}")

weight_decay=0.0: max abs err  param:0.00e+00 exp_avg:2.38e-07 exp_avg_sq:2.38e-07
weight_decay=0.1: max abs err  param:4.05e-04 exp_avg:3.32e-02 exp_avg_sq:1.07e-03

exp_avg is off by 3.3e-02 — five orders of magnitude above fp32 noise, not a precision issue.

Root cause

src/flag_gems/ops/_fused_adam.py#L94-L104:

# Compute the update step
# For weight decay, we use AdamW style (decay the parameters)
if weight_decay > 0:
    # AdamW: param = param - lr * (m / denom + weight_decay * param)
    update = corrected_exp_avg / denom + weight_decay * param_load
else:
    # Adam: param = param - lr * m / denom
    update = corrected_exp_avg / denom

For comparison, _single_tensor_adam in PyTorch
(torch/optim/adam.py#L397-L398)
folds the decay into the gradient before exp_avg.lerp_ / exp_avg_sq.mul_().addcmul_().
The ATen CUDA kernel does the same: ADAM_MODE::ORIGINAL adds param * weight_decay to the
gradient, and only ADAM_MODE::ADAMW decays the parameter.

The weight_decay > 0 guard is also too narrow — a negative weight decay is silently ignored.

Fix

Fold the decay into the gradient right after the maximize negation, and drop the extra term
from the parameter update:

if weight_decay != 0:
    grad_load = grad_load + weight_decay * param_load
...
param_new = param_load - lr * corrected_exp_avg / denom

With this the same script gives param:2.98e-08 exp_avg:2.38e-07 exp_avg_sq:2.38e-07 at
weight_decay=0.1, i.e. fp32 rounding.

tests/test_fused_adam.py does not catch this because it only ever runs weight_decay=0.0,
and its hand-written reference implements the AdamW rule too.

PR to follow.

Environment

• FlagGems master (bcf3f5b7e7c2); the kernel is unchanged since it was added
• torch 2.6.0+cu124, triton 3.2.0, NVIDIA H100, CUDA 12.4