Tiankuan-lanzhichuan

[JerryJonny]

Description

Add an NPU backend implementation for multi_tensor_adam in TransformerEngine-FL.

This implementation adapts the TE-FL multi_tensor_adam interface to the NPU Adam/AdamW execution path, including input validation, optimizer state handling, dtype handling, noop behavior, and backend registration.

With this change, multi_tensor_adam can dispatch to the NPU vendor backend instead of falling back to the reference.torch implementation.

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

• Add the NPU implementation of multi_tensor_adam.
• Register multi_tensor_adam for the NPU vendor backend.
• Support Adam/AdamW update semantics required by the TE-FL interface.
• Handle optimizer states, dtype conversion, noop behavior, and in-place parameter updates.
• Preserve the existing reference.torch implementation as the reference/fallback path.

Validation

The implementation was validated with dedicated local optimizer tests and end-to-end Qwen3.5 MoE training on 16 Ascend NPUs.

Training configuration:

• Sequence length: 2048
• Seed: 42
• Same random initialization for reference and NPU runs
• BF16 training
• FP32 main parameters, gradients, and optimizer states
• 100 training iterations

Correctness and stability:

• Relevant optimizer tests passed locally.
• Both reference and NPU runs completed 100 iterations without NaN.
• No skipped iterations were observed.
• Loss and parameter-norm trajectories remained close to the reference.torch implementation.
• Small numerical differences accumulated after the first few optimizer updates, as expected from implementation and operation-order differences.

Performance was measured over iterations 11-100 to exclude warm-up:

Backend	Average iteration latency
reference.torch	1.4844 s
vendor.npu	1.3199 s

The NPU backend reduced end-to-end iteration latency by approximately 11.1%, corresponding to about 1.125x speedup in this configuration.

Checklist

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[CLAassistant]

All committers have signed the CLA.