DeepSpeed
53e336cd - Reject amsgrad in DeepSpeedCPUAdam instead of ignoring it (#8694)

Commit
19 days ago
Reject amsgrad in DeepSpeedCPUAdam instead of ignoring it (#8694) `DeepSpeedCPUAdam(amsgrad=True)` silently runs plain Adam. The CPU kernel has no AMSGrad path, so the flag is stored in the param group and never read. With `"amsgrad": true` in the optimizer config, `FusedAdam` raises, but turning on `offload_optimizer` swaps in CPU Adam and the run trains with a different algorithm and no warning. Fix: reject `amsgrad` in `add_param_group`, the same way `FusedAdam` and `FusedLamb` do. That covers the constructor, the engine's config path, ZenFlow's subclass and groups added later. `__del__` now skips the C++ destroy when construction failed, so the error isn't followed by an `AttributeError`. Test: `test_cpu_adam_rejects_amsgrad` fails on master (DID NOT RAISE) and passes with the fix. `tests/unit/ops/adam/` and the Muon partial-training tests pass: 158 passed, 14 skipped. Signed-off-by: Vineeth Sai <vineethsai4444@gmail.com>
Author
Parents
Loading