feat: validate AMP support on MUSA (#145)

Run the shared PrivateUse1 autocast and GradScaler contract on native MUSA builds, and document measured FP16/BF16 behavior on MTT S5000. Keep GradScaler unscale on the existing correctness-oriented CPU fallback because mudnn has no fused AMP foreach primitive.

Co-authored-by: lvyufeng <lvyufeng@flagopen.baai.ac.cn>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

Author: 3 people

README.md

@@ -81,7 +81,7 @@ Capability availability and validation status vary by platform. A feature existi
 | PPU | CUDA boxing against PPU's CUDA-13-compatible SDK | Eager, autograd | **Experimental** | [PPU](docs/vendors/ppu/installation.md) |
 | Hygon DCU | CUDA boxing over hipified DTK torch | Eager, autograd, FP16/BF16 AMP, profiler | **Beta** | [DCU](docs/vendors/dcu/installation.md) |
 | Enflame GCU | Native topsaten backend, CPU fallback for unrouted/int64 ops | Eager | **Experimental** | [GCU](docs/vendors/gcu/installation.md) |
-| Moore Threads MUSA | Native mudnn backend, CPU fallback for unrouted ops | Eager | **Experimental** | [MUSA](docs/vendors/musa/installation.md) |
+| Moore Threads MUSA | Native mudnn backend, CPU fallback for unrouted ops | Eager, FP16/BF16 AMP | **Experimental** | [MUSA](docs/vendors/musa/installation.md) |
 | D-Robotics BPU | No eager kernels; `torch.compile` graph path via hbdk4 | Graph compilation only | **Runtime only** | [BPU](docs/vendors/bpu/installation.md) |
 | TsingMicro | Runtime build selector exists | Setup not documented | **Runtime only** | — |
 

docs/reference/compatibility.md

@@ -28,7 +28,7 @@
 | PPU | `ACCELERATOR=cuda` + `PPU_SDK`/`PPU_HOME` detection | Same CUDA-boxing path as NVIDIA CUDA, against the PPU's CUDA-13-compatible SDK | Experimental (FP16/BF16 autocast and GradScaler measured on PPU hardware, not in CI) | Not validated | Experimental (NCCL fallback via vendor-adapted `libnccl.so.2`; not CI-covered) | Not validated on this vendor's tracer | Experimental (vendor-index Triton required) | Experimental |
 | Hygon DCU | `ACCELERATOR=dcu` | CUDA boxing over the hipified DTK torch build (HIP kernels under the CUDA dispatch key) | Beta (including FP16/BF16 autocast and GradScaler) | Experimental (FlagTree HCU validated on `gfx936`; not in CI) | Experimental (RCCL via DTK; all_reduce/DDP measured on 2 cards, not in CI) | Beta (parity suite runs in CI) | Beta (Python dispatch only) | Beta |
 | Enflame GCU | `ACCELERATOR=gcu` | Native `libtopsaten.so` operator backend, with CPU fallback for unrouted/int64 ops | Experimental | Not validated | Not validated | Not validated | Experimental (Python dispatch, requires vendor Triton) | Experimental |
-| Moore Threads MUSA | `ACCELERATOR=musa` | Native `mudnn` operator backend, with CPU fallback for unrouted ops | Experimental | Not validated | Not validated | Experimental (MUPTI device timeline measured on MTT S5000; CPU-Kineto linkage is environment-dependent) | Experimental (Python dispatch, requires vendor Triton) | Experimental |
+| Moore Threads MUSA | `ACCELERATOR=musa` | Native `mudnn` operator backend, with CPU fallback for unrouted ops | Experimental (including FP16/BF16 autocast and GradScaler measured on MTT S5000) | Not validated | Not validated | Experimental (MUPTI device timeline measured on MTT S5000; CPU-Kineto linkage is environment-dependent) | Experimental (Python dispatch, requires vendor Triton) | Experimental |
 | D-Robotics BPU | `ACCELERATOR=bpu` | No eager kernel sets are built; eager ops run on CPU | Runtime only (CPU fallback for eager) | Experimental (`torch.compile(backend="bpu")` graph path via hbdk4) | Not applicable | Not validated | Not applicable (no per-op kernel build) | Runtime only |
 | TsingMicro | `ACCELERATOR=tsingmicro` | Runtime/build selector present; no per-op kernel set documented | Runtime only | Not validated | Not validated | Not validated | Not applicable | Runtime only |
 
@@ -155,11 +155,17 @@ with a documented CPU fallback for ops the vendor kernel table does not cover (s
 [`setup.py`](../../setup.py), lines 440-457, and
 `tests/integration/ops/test_musa_dispatch.py`, lines 15-24). No CUDA boxing kernels or FlagGems
 C++ kernels are compiled for this platform; FlagGems Python dispatch is optional and needs the
-vendor Triton backend. Distributed support remains unvalidated on hardware. The optional MUPTI
-tracer has been measured on the available MTT S5000 host and emits real positive-duration kernel,
-runtime, and memcpy activities with device, stream, name, and Chrome-trace metadata. CPU-to-device
-linkage depends on whether the installed PyTorch/Kineto build supplies the PrivateUse1 resolver;
-the CPU-only PyTorch 2.10 wheel used for this measurement does not justify a general parity claim.
+vendor Triton backend. The shared `AutocastPrivateUse1` policies and native mudnn routes were
+measured on an MTT S5000 with both FP16 and BF16. The 25-case AMP suite covers lower-precision
+matmul, linear, and convolution; float32 and promote policies; nested autocast state; finite scale
+growth; overflow backoff and optimizer-step skipping; and end-to-end GradScaler training. The
+GradScaler unscale operation currently follows the correctness-oriented CPU fallback and copies
+its mutated tensors and `found_inf` flag back to MUSA rather than using a native mudnn foreach
+kernel. Distributed support remains unvalidated on hardware. The optional MUPTI tracer has been
+measured on the available MTT S5000 host and emits real positive-duration kernel, runtime, and
+memcpy activities with device, stream, name, and Chrome-trace metadata. CPU-to-device linkage
+depends on whether the installed PyTorch/Kineto build supplies the PrivateUse1 resolver; the
+CPU-only PyTorch 2.10 wheel used for this measurement does not justify a general parity claim.
 
 ### D-Robotics BPU
 

docs/reference/dtype-support.md

@@ -25,13 +25,29 @@ is not an AMP target and is not accepted by its native matmul API.
 | Ascend | float16, bfloat16, float32, float64, integer, uint8, bool | Vendor coverage; unsupported ACLNN combinations use the CPU fallback | float16, bfloat16, float32 natively; float64 and unsupported types use the CPU fallback | float16, bfloat16 |
 | MetaX boxing | MACA libtorch CUDA-compatible coverage | MACA libtorch CUDA-compatible coverage | MACA libtorch CUDA-compatible coverage | float16, bfloat16 |
 | DCU | Vendor library coverage | Vendor library coverage | Vendor library coverage | Backend-specific |
-| MUSA | Vendor library coverage | Vendor library coverage | Vendor library coverage | Backend-specific |
+| MUSA | float16, bfloat16, float32, float64, integer, bool | mudnn coverage; unrouted operations use the CPU fallback | float16, bfloat16, and float32 measured for AMP paths; broader support is vendor-dependent | float16, bfloat16 |
 
 The MetaX AMP entry is measured for the CUDA-boxing path on C550 with MACA
-3.8.0; it does not cover the legacy handwritten MetaX kernel mode. The DCU and
-MUSA entries intentionally do not claim AMP target parity without hardware
-measurements for the specific library release. Their native boxing or
-code-generated kernels determine the available operator/dtype combinations.
+3.8.0; it does not cover the legacy handwritten MetaX kernel mode. The DCU entry
+remains vendor-dependent. MUSA AMP coverage below is measured for the specific
+S5000 and mudnn release rather than inferred from its route configuration.
+
+## MUSA Boundaries
+
+MUSA AMP uses the shared `AutocastPrivateUse1` policy lists with the native
+mudnn operator routes. On the measured MTT S5000 and mudnn v3300 setup, both
+float16 and bfloat16 execute the lower-precision matmul, linear, and convolution
+policies. Logarithm, layer normalization, MSE loss, and the default softmax
+policy return float32; an explicit softmax dtype is preserved; mixed inputs use
+the promote policy.
+
+`torch.amp.GradScaler("flagos")` uses the existing PrivateUse1 fallback for
+`_amp_foreach_non_finite_check_and_unscale_`. That fallback moves the list and
+scalar operands to CPU, runs the reference kernel, and copies the mutated values
+and `found_inf` flag back to MUSA. This is correctness-oriented rather than a
+claim of native mudnn AMP-foreach acceleration. Finite scale growth, overflow
+backoff, optimizer-step skipping, and FP16/BF16 autocast training were measured
+on the same S5000 host.
 
 ## Ascend Boundaries
 

docs/vendors/musa/installation.md

@@ -99,6 +99,14 @@ pytest tests/integration/ops/test_musa_dispatch.py -v
 
 This test file checks that ops route to the `musa` backend, exercises per-op environment overrides, and validates results against CPU references. It replaces the generic per-op tests in `tests/integration/ops/`, which assert `-> cuda` routing that MUSA builds cannot produce (no CUDA boxing kernels are compiled).
 
+Run the autocast and GradScaler suite:
+
+```bash
+TORCH_DEVICE_BACKEND_AUTOLOAD=0 ACCELERATOR=musa \
+  LD_LIBRARY_PATH=/usr/local/musa/lib:$LD_LIBRARY_PATH \
+  pytest tests/integration/test_amp.py -v
+```
+
 Run common operator smoke tests:
 
 ```bash
@@ -135,6 +143,44 @@ Category-driven codegen via `scripts/codegen_mudnn.py` covers:
 - Handles both in a single device pass
 - Without this, `copy_`/`clone`/`contiguous` would reach the CUDA `DispatchStub` and fail
 
+### AMP and GradScaler
+
+MUSA uses the shared `AutocastPrivateUse1` policies with its native mudnn routes.
+Both `torch.float16` and `torch.bfloat16` are supported autocast targets:
+
+```python
+import torch
+import torch.nn.functional as F
+import torch_fl  # noqa: F401
+
+model = torch.nn.Linear(8, 4, device="flagos")
+optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
+scaler = torch.amp.GradScaler("flagos")
+inputs = torch.randn(2, 8, device="flagos")
+targets = torch.randn(2, 4, device="flagos")
+
+with torch.autocast("flagos", dtype=torch.bfloat16):
+    output = model(inputs)
+    loss = F.mse_loss(output, targets)
+
+scaler.scale(loss).backward()
+scaler.step(optimizer)
+scaler.update()
+```
+
+The standard policy groups cast matmul, linear, and convolution to the selected
+lower-precision dtype. Logarithm, layer normalization, MSE loss, and the default
+softmax policy run in float32; operations in the promote group use the widest
+input dtype. Nested autocast state and explicit softmax dtypes follow PyTorch's
+PrivateUse1 behavior.
+
+GradScaler's `_amp_foreach_non_finite_check_and_unscale_` currently reaches the
+existing CPU fallback because mudnn has no fused AMP-foreach primitive. The
+fallback copies tensor-list operands and scalar state to CPU, invokes PyTorch's
+reference kernel, and copies the mutated gradients and `found_inf` flag back to
+MUSA. This preserves finite-step growth, overflow backoff, and optimizer-step
+skipping, but it is not a native device-side performance path.
+
 ### Native RNG and CPU fallback
 
 Native RNG uses a per-device Philox generator from `libmurand.so`. Each operation consumes one seed from the selected PyTorch `flagos` generator under its mutex, then configures a cached muRAND generator on the shared MUSA stream. This preserves `torch.manual_seed`, `torch.flagos.manual_seed`, `manual_seed_all`, `get_rng_state`/`set_rng_state`, explicit `generator=` isolation, and independent per-device sequences without maintaining a second Python-side RNG state. The optional FlagGems bridge reserves operation seeds through the same C++ API and starts each FlagGems Philox invocation at offset zero, so mixed native/FlagGems call order has one authoritative state.
@@ -183,6 +229,7 @@ MUSA has no CI runner (no `.github/configs/musa.yml` config exists), so all test
 
 The native RNG and hybrid FlagGems implementation were measured on 2026-08-17 on an eight-device Moore Threads MTT S5000 host. Device 0 reported capability 3.1, 60 multiprocessors, and 85,813,358,592 bytes of memory. The environment used CPU PyTorch 2.10.0, the installed `/usr/local/musa` toolkit (`mudnn` v3300), FlagGems 5.0.2, and the vendor `flagtree-0.5.0+mthreads3.1` wheel (Triton 3.1.0, backend `mthreads`). Results:
 
+- `pytest tests/integration/test_amp.py -v`: **25 passed** in 9.00 seconds. Both FP16 and BF16 passed lower-precision matmul/linear/convolution, float32 and promote policies, nested state, mutable non-finite unscale, finite scale growth, overflow backoff, optimizer-step skipping, and autocast/GradScaler training.
 - `pytest tests/integration/ops/test_rng_dispatch.py -m "main_ops" -v`: the unified RNG suite passed, including explicit generators, state round trips, `torch.manual_seed`, full-width int64 ranges, native dropout, shared native/FlagGems reservation ordering, and multi-device sequence isolation. MUSA-only generator and reservation cases are selected by the `musa` mark.
 - `pytest tests/integration/ops/test_musa_dispatch.py -v`: **89 passed** in 36.31 seconds.
 - `pytest tests/unit/test_vendor_routing.py tests/unit/test_musa_rng_bridge.py -v`: **24 passed** in 2.19 seconds.

tests/integration/test_amp.py

@@ -15,23 +15,23 @@
 import os
 
 import pytest
-import torch_fl
 import torch
 import torch.nn.functional as F
+import torch_fl
 
 
 DEVICE = "flagos:0"
 AMP_DTYPES = (torch.float16, torch.bfloat16)
 
-# PPU uses the CUDA-compatible build path, so ACCELERATOR alone cannot
-# distinguish it from an NVIDIA build. MetaX has an explicit build selector.
-# Run this hardware contract only on those two validated vendor runtimes.
+# PPU shares the CUDA build selector, so its SDK marker distinguishes it from
+# NVIDIA. Native AMP backends have explicit selectors. Device availability is
+# enforced by integration conftest.
 BUILD_ACCELERATOR = torch_fl._build_accelerator()
 PPU_RUNTIME = bool(os.environ.get("PPU_SDK") or os.environ.get("PPU_HOME"))
-AMP_RUNTIME = BUILD_ACCELERATOR == "metax" or PPU_RUNTIME
+AMP_RUNTIME = BUILD_ACCELERATOR in {"ascend", "dcu", "metax", "musa"} or PPU_RUNTIME
 pytestmark = pytest.mark.skipif(
     not AMP_RUNTIME,
-    reason="AMP tests require a MetaX build or PPU_SDK/PPU_HOME runtime",
+    reason="AMP tests require an Ascend, DCU, MetaX, MUSA, or PPU runtime",
 )