DeepSpeed

c456c059 - [BUG] Replace deprecated `npu_apply_adam(_w)` with `torch._fused_adam(w)_` on NPU (#8640)

[BUG] Replace deprecated `npu_apply_adam(_w)` with `torch._fused_adam(w)_` on NPU (#8640) ## Summary The NPU backend of `FusedAdam` (`op_builder/npu/fused_adam.py`) currently relies on `torch_npu.npu_apply_adam` / `torch_npu.npu_apply_adam_w`. Both APIs are marked **"planned for deprecation"** in the official TorchNPU documentation, which recommends `torch.optim.Adam` / `torch.optim.AdamW` as replacements. This PR switches the implementation to `torch._fused_adam_` / `torch._fused_adamw_`, the fused multi-tensor ops used by `torch.optim.Adam/AdamW(fused=True)` and supported by TorchNPU on Ascend. This follows the same migration already adopted in MindSpeed (gitee.com/ascend/mindspeed/pulls/2205). The change also fixes a correctness issue on Ascend910_93 (Atlas A3) with CANN 9.1 + torch_npu 2.10.0.post2. In a single-step test with shape `(1024,)` and three seeds: | API | fp32 | bf16 | fp16 | | --- | --- | --- | --- | | `npu_apply_adam` (Adam) | NaN | NaN | NaN | | `npu_apply_adam_w` (AdamW) | OK | NaN | NaN | | `torch._fused_adam_` / `torch._fused_adamw_` | OK | OK | OK | This means the current NPU `FusedAdam` path is broken for bf16/fp16 AdamW training and for Adam mode on this platform. ## Implementation - Replace `torch_npu.npu_apply_adam` / `torch_npu.npu_apply_adam_w` with `torch._fused_adam_` / `torch._fused_adamw_`. - Keep the existing `multi_tensor_adam` interface unchanged. - Pass the caller-managed step tensor to the fused kernels. ## Behavior changes - `bias_correction=False` is still not supported by this NPU path because the fused kernels always apply bias correction. The previous implementation also ignored this flag. - Adam mode (`adam_w_mode=False`) now applies `weight_decay` as L2 regularization, consistent with the CUDA implementation and `torch.optim.Adam`. The previous NPU implementation silently ignored `weight_decay` in Adam mode. ## Verification Tested on Ascend910_9382 (Atlas A3) with CANN 9.1, torch 2.10.0, and torch_npu 2.10.0.post2. ### Operator-level reference test Compared against the same fp32 reference used by `test_fused_adam_matches_reference` across Adam/AdamW × fp32/bf16/fp16 × `weight_decay` {0, 0.1}, with 5 optimizer steps and 3 tensors per configuration. All 12 configurations passed within tolerance (fp32 max error: `1.9e-6`). For fp32 AdamW, the only tested configuration where the old API remained finite, max `|Δparam|` was `7.2e-7` after 5 steps. ### Existing DeepSpeed tests `tests/unit/ops/adam/test_adamw.py`: 22/22 passed, including engine integration and reference-matching cases. ### End-to-end SFT A 30-step LLaMA-Factory SFT run with Qwen3-0.6B, bf16, and ZeRO-2 completed without NaNs; loss decreased from `2.32` to `1.82`. Compared with `torch_adam=true` (`torch.optim.AdamW`) using the same seed and data: ```text max |Δloss| = 0.011 mean |Δloss| = 0.0024 ``` <img width="1050" height="600" alt="image" src="https://github.com/user-attachments/assets/83a46371-d84c-4c72-a9dd-38120b0043b7" /> Signed-off-by: Yinlei Sun <ginray0215@gmail.com>