DeepSpeed
d47e4f4d - Enable optimized Adam backend for MuonWithAuxAdam optimizer (#8278)

Commit
4 days ago
Enable optimized Adam backend for MuonWithAuxAdam optimizer (#8278) ### Motivation This pull request aim to extend the Muon optimizer auxiliary Adam(w) optimizer to support more built-in optimized implementation such as `FusedAdam`. ### Changes * Extracted Adam/AdamW backend selection logic into a new `_select_adam_optimizer` method, improving modularity and allowing consistent backend selection for both Adam and Muon optimizers. * Keep original naive implementation as fallback path and extend to support weight-decay (AdamW) * Support `CPUAdam` and `Zenflow` ### Tests ``` pytest tests/unit/ops/muon/test_muon_partial_training.py ``` <img width="1894" height="379" alt="图片" src="https://github.com/user-attachments/assets/125b2d0c-21d5-44a0-93af-7433aac5f5ba" /> tested on RTX 5090: 128 Tensors 16,384 FP32 elements, warmup 30, steps 200. --------- Signed-off-by: iLeGend <824040212@qq.com>
Author
Parents
Loading