-
Honor adam_w_mode in the CPU multi_tensor_adam binding (#8307)
发布于
2026-08-27 17:12:18 +00:00 | 59 次提交 在此版本后已推送到 masterSummary
The CPU
fused_adamextension created itsAdam_Optimizeronce with
default arguments and ignored themodeparameter entirely
(csrc/cpu/adam/fused_adam.cpp), soFusedAdam(adam_w_mode=False)on
the CPU backend always applied decoupled (AdamW) weight decay instead of
L2. Everything else (lr, betas,eps,weight_decay, bias
correction) is already passed per call viads_adam_step; only the
AdamW-vs-L2 flag is fixed at construction. The fix keeps one optimizer
instance per mode (mode 1 == AdamW, matching the CUDA kernel's
ADAM_MODE_1).Also trims
test_fused_adam_matches_torchto fp32: its bf16 cases
compared againsttorch.optimrunning bf16 math, while the fused
kernels compute in fp32 — never a valid reference. Low-precision dtypes
get an explicit fp32-math reference test in the FusedAdam rework
(#8300).How this surfaced
Split out of #8303 at @delock's request: after a master merge,
cpu-torch-latest failed ontest_fused_adam_matches_torch[fp32-adam]
(98.7% of elements mismatched — systematic, not tolerance noise), and
the investigation traced it to this binding. The fix was verified green
on cpu-torch-latest in #8303's CI (run 32695...) before being extracted
here.Validation
test_fused_adam_matches_torch[fp32-adam]/[fp32-adamw]now
genuinely exercise both decay modes againsttorch.optim.Adam/AdamW
on the active accelerator.- cpu-torch-latest passed with this exact change as part of #8303's
branch; this PR carries it alone.
Signed-off-by: PKUWZP zhipeng.rainbowserie@gmail.com
Co-authored-by: Ma, Guokai guokai.ma@gmail.com下载附件