Honor adam_w_mode in the CPU multi_tensor_adam binding - #8307
Conversation
The CPU fused_adam extension created its Adam_Optimizer once with default arguments and ignored the mode parameter, so FusedAdam on the CPU backend always applied decoupled (AdamW) weight decay even when constructed with adam_w_mode=False. Keep one optimizer instance per mode instead; everything else is already passed per call. This surfaced as cpu-torch-latest failures in the fp32-adam case of test_fused_adam_matches_torch. The test's bf16 cases are dropped: torch.optim in bf16 does its math in bf16 while the fused kernels compute in fp32, so it was never a valid bf16 reference. Low-precision dtypes get an explicit fp32-math reference in the FusedAdam rework. Signed-off-by: PKUWZP <zhipeng.rainbowserie@gmail.com> (cherry picked from commit a746b27)
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 1e17e84f2b
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
|
CI triage note: modal-torch-latest tested the exact head 1e17e84. Because csrc/** is a run-all trigger, it ran the full tests/unit/v1 suite (1232 items), not the changed CPU Adam test. The run reached 97% without an Adam-related failure, then exited 137 almost exactly at the controller's 3600-second Modal Sandbox lifetime; the unrelated test_offload_activation case was marked failed only after its xdist worker was forcibly terminated. The same 1232-item full suite passed on the base SHA in 2479.20 seconds (run 32692693729), and cpu-torch-latest passed on this exact head (job 97354322902). The failed Modal job therefore looks like a full-suite lifetime timeout rather than evidence of an Adam regression. A rerun should distinguish runtime/order variance; if it recurs, the shared 3600-second Sandbox lifetime may need more headroom. |
| static bool initialized[2] = {false, false}; | ||
| const int optimizer_id = mode; | ||
| if (!initialized[mode]) { | ||
| create_adam_optimizer(optimizer_id, 1e-3f, 0.9f, 0.999f, 1e-8f, 0.0f, mode == 1); |
There was a problem hiding this comment.
Are these '1e-3f, 0.9f, 0.999f, 1e-8f, 0.0f' value place holder? If they are better define a macro as place holder and put it here.
tohtana
left a comment
tohtana
left a comment
There was a problem hiding this comment.
The fix looks good to me, thank you @PKUWZP! This is important to unblock the CI.
I noticed that the new regression test is not run in CPU CI on hosts without fp16 support. This means its coverage can vary depending on the underlying GitHub Actions runner hardware. As I understand it, this bug had existed for a long time but remained hidden because the runners being used did not support fp16, causing the entire test module to be skipped.
Could we scope the fp16 skip to TestAdamConfigs only, while leaving test_fused_adam_matches_torch eligible to run on CPU regardless of fp16 support?
4 participants
Summary
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'sADAM_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 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 againsttorch.optim.Adam/AdamWon the active accelerator.