发布

  • Honor adam_w_mode in the CPU multi_tensor_adam binding (#8307)

    frostbyte_neo 发布于 2026-08-27 17:12:18 +00:00 | 59 次提交 在此版本后已推送到 master

    Summary

    The CPU fused_adam extension created its Adam_Optimizer once with
    default arguments and ignored the mode parameter entirely
    (csrc/cpu/adam/fused_adam.cpp), so FusedAdam(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 via ds_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_torch to fp32: its bf16 cases
    compared against torch.optim running 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 on test_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 against torch.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

    下载附件