Honor adam_w_mode in the CPU multi_tensor_adam binding - #8307
Conversation
Sorry, something went wrong.
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)
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 GitHubYour 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".
Sorry, something went wrong.
|
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. |
Sorry, something went wrong.
| 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.
Sorry, something went wrong.
There was a problem hiding this comment.
They were placeholders, yes — update_state and IncrementStep overwrite every one of them on the first step (IncrementStep takes its full-reset branch whenever the incoming betas differ from the stored ones). Done in a4c407e: a single named constant kPlaceholderHyperparam = 0.0f with a comment saying exactly that, so the values no longer look meaningful.
Sorry, something went wrong.
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?
Sorry, something went wrong.
Let me fix it today. |
Sorry, something went wrong.
Resolve tests/unit/ops/adam/test_adamw.py in favor of master's reference-based FusedAdam test from #8300, which supersedes the bf16 trim this branch carried for the old torch-comparison test. Signed-off-by: PKUWZP <zhipeng.rainbowserie@gmail.com>
…ders The module-level fp16 skip also hid the dtype-parametrized FusedAdam reference test on CPU runners without fp16 support, which is exactly how the adam_w_mode bug stayed unnoticed. Only TestAdamConfigs needs fp16 (its config enables it), so the skip moves onto that class and the reference test now runs everywhere with its own per-dtype skips. The CPU binding's construction hyperparameters become a named placeholder constant: update_state and IncrementStep overwrite all of them on every step, so any value works and zero says so honestly. Signed-off-by: PKUWZP <zhipeng.rainbowserie@gmail.com>
…-mode-fix Signed-off-by: PKUWZP <zhipeng.rainbowserie@gmail.com>
|
@tohtana done in a4c407e — the module-level fp16 skip is removed and scoped onto TestAdamConfigs only (its config hard-codes fp16 enabled). After merging master, the dtype-parametrized reference test from #8300 (which replaced test_fused_adam_matches_torch) now runs on CPU runners regardless of fp16 support, with its own per-dtype skips — so this regression can't hide behind runner hardware again. Also merged master to resolve the branch conflicts (took master's reference-based test, kept the kernel fix), and re modal-torch-latest: @FU-max-boop's triage looks right — the failed run hit the 3600s sandbox lifetime at 97% of the full 1232-item suite with no Adam-related failure; the new push re-runs it. |
Sorry, something went wrong.
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