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.
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.
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?
Let me fix it today. |
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. |
…erator (deepspeedai#8335) ## Summary Closes the remaining gap in the Apple Silicon support series (deepspeedai#8293, deepspeedai#8300, deepspeedai#8303, deepspeedai#8307): none of the MPS paths were exercised by CI — every MPS-gated test skips on Linux runners, so regressions could only be caught on a developer's Mac. ### macOS CI workflow (`mps-torch-latest.yml`) Runs the MPS-green unit test subset on GitHub's arm64 macOS runners (`macos-15`), which expose a working MPS device: - `unit/ops/adam/test_adamw.py` — Metal/foreach FusedAdam vs fp32-math reference, CPU Adam configs incl. ZeRO-Offload - `unit/comm/test_dist.py` — gloo CPU-staging for collectives and P2P (`TestMpsStagedP2P`) - `unit/runtime/test_ds_config_dict.py` — config-driven `deepspeed.initialize` + training steps **Designed not to interfere with existing CI:** - PR triggers are scoped via `paths:` to MPS-relevant files (`accelerator/**`, `op_builder/mps/**`, `csrc/mps/**`, `deepspeed/comm/**`, the two test dirs, and the workflow itself) — the check does not even appear on unrelated PRs. - Separate workflow, own concurrency group with cancel-in-progress, hard `timeout-minutes: 45`. - Not a required check (that's a branch-protection setting; nothing here changes it), so even a red run cannot block merges of non-macOS work. - Nightly `schedule` + `workflow_dispatch` for coverage between touching PRs. ### torch floor check `MPS_Accelerator.__init__` now fails with a clear message on torch older than 2.3, where the `torch.mps` memory queries ZeRO depends on (`recommended_max_memory`) do not exist — previously this surfaced as a bare `AttributeError` deep inside ZeRO's flatten logic. Feature-detected rather than version-parsed. (The Metal FusedAdam kernel already degrades gracefully on torch without `compile_shader`.) ## Validation - The workflow's exact pytest command passes locally on an M5 Max (macOS 26.3, torch 2.13): 65 passed, 23 skipped (multi-device), 1m54s — comfortably inside the runner budget. - Guard verified both ways: normal construction unaffected; with `recommended_max_memory` hidden, construction raises the explicit `ValueError`. --------- Signed-off-by: PKUWZP <zhipeng.rainbowserie@gmail.com>
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.