feat(backend): add native Ascend 950DT support

[zhaosiying12138]

Related to #18.

DMuon currently assumes CUDA for device operations and optional kernels. This PR adds native Ascend 950DT execution through torch_npu and HCCL while retaining the existing public interfaces and shared optimizer implementation.

1. Native runtime: Add a private adapter for device selection, streams, events, synchronization, and allocator lifetime. Resolve the runtime from the actual tensors and device mesh without a CUDA-to-NPU transfer shim or emulated CUDA capability.
2. Shared optimizer: Reuse the existing ownership, DDP/FSDP2/HSDP scheduling, Muon/AdamW mathematics, parameter policies, and PyTorch Newton–Schulz paths. Guard CUDA-only imports, extensions, SYRK kernels, and graph replay without duplicating optimizer algorithms.
3. Accumulation and buffer lifetime: Preserve every microbatch's contribution to sharded AdamW gradients, honor explicit accumulation dtypes and no_sync, drain the complete reduction pipeline, and retain asynchronous inputs and outputs until their consumers finish.
4. Checkpoint placement: Honor CPU offload for ordinary AdamW tensor states and restore those states to the live parameter's device and dtype. Preserve non-tensor counters and metadata, and retain upstream's Muon momentum dtype policy.
5. Empty NPU shards: Initialize moments and advance step counters for zero-length FSDP shards while excluding those shards from unsupported NPU foreach launches. Nonempty updates and the CUDA path retain their existing mathematics.
6. Mandatory correctness gate: Extend test_hsdp_correctness.py with an explicit native NPU entry while preserving the existing CUDA default. Check exact loss, full model state, and optimizer state on every rank across 1D and 2x2 meshes, then verify continuation from one canonical rank-0 disk checkpoint. Add isolated tests that reject one-ULP differences, signed-zero differences, NaNs, and schema mismatches.
7. Documentation and regression coverage: Document CUDA and native Ascend installation and test commands. Add API/source guards, dtype and nested-cast tests, clipping, accumulation, asynchronous lifetime, and checkpoint-continuation coverage. CUDA-only scripts skip cleanly during pytest collection when their hardware or dependencies are unavailable.

Validation:

• Unit suite: 294 passed, 31 skipped; run without collection exclusions.
• Native mandatory gate: four Ascend NPUs, ten updates per topology plus two HSDP replay updates. Every rank passed exact tensor-bit comparisons. The original model dimensions and optimizer configuration are retained.
• Additional native checks: Newton–Schulz and clipping reference comparisons, mixed-precision updates, GAS4 sync/no_sync, asynchronous publication, and DP/HSDP restart cases.
• Ruff lint passed for dmuon/ and examples/; formatting passed for modified Python files. Twelve unchanged upstream files still fail the repository-wide format check and are intentionally left outside this PR's cleanup scope.
• Existing CUDA model and runner ASTs and all nine guarded CUDA implementation files remain unchanged. CUDA hardware execution was not available for this validation.

[StarrickLiu]

Thanks for the thorough work. The device abstraction in dmuon/_core/device.py is small and clean, and the sharded AdamW accumulation change fixes a real bug on main (every reduce reassigned _sharded_adamw_grad, so gradient accumulation without no_sync kept only the last micro-batch). I reproduced the unit results locally (294 passed / 31 skipped). A few requests before merge:

1. Please validate the CUDA path as far as you can

Several changes alter CUDA behavior, not only NPU: the no_sync merge / _grad_dtype accumulation and the new record_stream calls in fsdp2/group.py, the sharded AdamW accumulator, and checkpoint offload/restore. Since CUDA hardware wasn't available for your validation, please run as much of the following on GPUs as you can and paste the results:

pytest tests/unit/
torchrun --nproc_per_node=4 tests/distributed/test_hsdp_correctness.py
torchrun --nproc_per_node=4 tests/distributed/test_hsdp_async_correctness.py
torchrun --nproc_per_node=4 tests/distributed/test_grad_accum.py
torchrun --nproc_per_node=4 tests/distributed/test_muon_step.py
torchrun --nproc_per_node=4 tests/distributed/test_checkpoint.py
torchrun --nproc_per_node=4 tests/distributed/test_ddp_correctness.py

If you have no GPU access at all, let us know which ones you couldn't run.

2. Please remove tests/unit/test_public_api_contract.py and tests/unit/fixtures/public_api.json

test_cuda_specific_kernel_sources_are_unchanged pins the sha256 of 9 CUDA source files, and test_public_exports_and_signatures_are_unchanged pins AST hashes of the public signatures. Any future legitimate kernel change or new keyword argument would fail the unit suite; it already fails against one of our in-flight branches that adds a dedicate_params kwarg. These checks are good evidence for this PR, but that evidence belongs in the PR description, not the permanent test suite.

3. README: please slim down the Ascend section

The PR inserts 56 lines (including a full training snippet) ahead of "3-Line Integration", so the top-level README now reads Ascend-first. Please keep the README install section short and move the full NPU example and test commands to docs/getting-started/installation.md / tests/README.md. Something like:

## Install

```bash
git clone git@github.com:X-Square-Robot/dmuon.git && cd dmuon
pip install -e .
```

Optional CUDA kernels: `pip install -e ".[syrk]"` (SM80) or `".[quack]"` (SM90+).

**Ascend NPU (experimental):** activate CANN, then
`DMUON_BUILD_FAST_CLIP=0 pip install -e . --no-build-isolation`. Import `torch_npu` and use an
`"npu"` device mesh with the `hccl` backend; the API is unchanged. CUDA-only kernels (SYRK, fast
clip, CUDA-graph NS) are skipped automatically. See the
[installation guide](docs/getting-started/installation.md) for details.

Please also drop "The existing CUDA APIs and training examples below are unchanged." That's context for reviewing this PR, not user documentation.

Minor (non-blocking)

• dmuon/optim/muon.py:1856: param.device.type == "npu" and param.numel() == 0 adds a device special case to the shared AdamW path. Skipping empty shards on every device (after state init and step += 1) is equivalent on CUDA and keeps the shared path device-agnostic.
• tests/unit/test_fast_clip.py: loosening found_inf is True to bool(...) is True hides that the extension path returns a Python bool while the fallback returns a 0-d tensor. That inconsistency predates this PR (the strict test already fails on main in a CPU-only env), but it's better to normalize the type of GradClipBucketStats.found_inf in source and keep the strict assertion.
• dmuon/_core/device.py:63: cuda_is_available() requires torch.version.cuda, so ROCm builds fall through to NPU/CPU in default_device(). Fine if ROCm isn't a target; consider saying so in the docstring.
• dmuon/utils.py:174: _group_device_module(group) uses the loop variable leaked from the earlier for group in groups. It's correct today because groups must be non-empty to reach that line, but it's fragile; derive the module from groups[0] or from the reduce stream's device instead.
• dmuon/fast_clip.py:440: the try_clip_optimizer_grad_norm_buckets_ docstring still says None is returned only when the extension is unavailable. It now also returns None on non-CUDA devices.
• dmuon/_backends/fsdp2/group.py:1004: p._accumulated_grad = grad.clone() adds an extra copy when _get_grad_for_reduce has already produced a fresh tensor (dtype cast). Likewise, the merge path clones _accumulated_grad when there is no raw grad, even though the accumulator is dropped right afterwards.
• With grad_dtype=torch.float32, the no_sync accumulator is now a full unsharded fp32 tensor, which doubles accumulator memory for bf16 params compared to before. That's intended, but worth a line in the training guide.
• Docs: docs/zh/contributing.md and docs/zh/getting-started/installation.md were not updated along with the English pages.
• tests/native/ (~1.7k lines) only runs on Ascend and isn't in CI. Please add a short note to tests/README.md on who runs it and when, so CUDA-side contributors know how changes there are validated.

Related fix

While reviewing, we found that the HSDP replicate_grads_enabled=False path on main is broken: the partial-stash branch is unreachable, so gated micro-batches hand a shard-only partial to the optimizer. That's fixed separately in #20 (validated on 4×H100). It touches the same region of fsdp2/group.py, so whichever PR lands second needs a rebase. Note that after this PR, sharded AdamW entries live in _pending_reduce, so the gated branch there must keep them on every shard rank, not only on owner_shard.

[zhaosiying12138]

Thanks for the careful review. I have rebased onto main with #20 included, preserved the original 13 commits in order, and added separate follow-up commits for your requests. The current PR head is 18e06bd, including the three test corrections found during hardware validation. All hardware results below were rerun on this exact head on 14 October 2026; the two full log attachments include commands, versions, timestamps and exit codes.

1. CUDA validation: We now have results from 4x A800 with PyTorch 2.10 / CUDA 12.8. The existing SM80 fast-clip extension was loaded and its fastpath exercised, alongside the forced Torch fallback. The import-isolation tests now distinguish native CUDA availability and shipped adapters from missing external dependencies (initial fix, adapter correction).

   Requested test	Recorded result
   pytest tests/unit/	393 passed, 27 skipped, 0 failed
   test_hsdp_correctness.py	Passed, 4 ranks
   test_hsdp_async_correctness.py	Passed, 4 ranks
   test_grad_accum.py	Passed, 4 ranks
   test_muon_step.py	Passed, 4 ranks
   test_checkpoint.py	Passed, 4 ranks
   test_ddp_correctness.py	Passed, 4 ranks

   All six distributed scripts used the requested torchrun --nproc_per_node=4 entrypoints and unchanged acceptance criteria. The additional four-rank replicate-gate and storage-ownership checks passed with NCCL; the two-rank clipping appendix passed with CPU/Gloo. The real SM80 clipping extension and forced CUDA fallback each passed all eight strict contract cases. Of the 27 unit skips, eight NPU cases passed separately on Ascend; 19 require SM90+/Quack and remain unverified on A800.

   Log audit: all requested commands exited 0 and their success markers/assertions were checked. The checkpoint next-update loss difference was 0.147705, within its existing <0.5 threshold; this is not an exact state-equality claim. The full output retains the 14 TorchScript deprecations, pytest entrypoint deprecation, default OMP notices and barrier-device warnings. No unexpected CUDA errors, timeouts or training non-finite values were found.

2. Implementation snapshot guards: Removed test_public_api_contract.py and its JSON fixture. The one-time comparison of public signatures and nine CUDA implementation files remains supporting PR evidence; behavioral tests remain in the suite (fix).

3. README: Kept installation brief, marked Ascend experimental, moved the full NPU example and test commands to their guides, and removed the review-context sentence and obsolete backlink (fix).

4. Empty AdamW shards: Skip zero-element foreach inputs on every device after state initialization and step advancement, with bookkeeping and mixed-empty/nonempty regressions (fix).

5. Clipping result types: I agree that converting with bool(...) alone is insufficient to check the contract. On the current implementation, extension and fallback statistics both pass through _build_stats, which produces scalar device torch.bool tensors. I retained that device-resident contract and added strict type, shape, dtype, device, and value checks across CPU, actual CUDA extension, forced CUDA fallback, and NPU (contract checks). The subsequent CUDA assertion fix uses explicit scalar readback only in numerical tests, after checking the Tensor contract; the original numerical tolerances are unchanged.

6. ROCm scope: Clarified that the helper detects NVIDIA CUDA and does not provide ROCm selection, and added the corresponding mock regression (fix).

7. Aggregate stream waits: Select the runtime explicitly from groups[0], retaining the empty-group and different-stream paths (fix).

8. Fast-clip fallback documentation: Documented the non-CUDA, unavailable-extension, and unsupported-optimizer None returns, with regressions for early-return side effects (fix).

9. Redundant accumulation copies: Transfer freshly cast storage directly into the first accumulator and transfer the private accumulator when the final raw gradient is absent. Borrowed gradients and TP views still receive a protective copy. Alias and asynchronous allocator/reshard-pressure checks passed on CUDA and NPU (fix).

10. FP32 accumulation memory: Added English and Chinese guidance explaining the 4-versus-2-byte accumulator payload for BF16/FP16 parameters. The 2x comparison applies to the DP-unsharded, TP-local accumulator; temporary and communication buffers are separate (fix).

11. Chinese documentation: Synchronized the installation and contribution guides with the English NPU instructions and mandatory gate (fix).

12. Native test ownership: Documented who runs the manual Ascend suite, which changes trigger it, and the revision/environment/multi-rank/restart evidence to retain. The docs CI is not presented as hardware validation (fix).

13. fix(fsdp2): make the HSDP replicate-reduce gate accumulate partials correctly #20 integration and gradient verification: Adapted the gated path to collect _pending_reduce once, retain sharded AdamW partials on every shard rank, and preserve stream ordering, lifetime, flush, and cleanup (production fix). The original same-weight consumed-gradient comparisons keep rtol=1e-5, atol=1e-7. They passed on CUDA/NCCL and NPU/HCCL, with worst relative differences of 8.060e-7 and 9.141e-7.

The initial extra continuous-trajectory comparison was too strong: different accumulation orders produced tiny gradient differences that FP16 normalization and NS amplified. The test correction checks the actual gradients before any replacement and retains failures in the final rank-wide result, then separately tests three continuing optimizer updates with identical gradient inputs. Those model/state comparisons passed byte-exactly on both backends. This establishes the gradient contract and controlled optimizer-state behavior, not equality of the two natural training trajectories.

Full validation attachments (all tests at 18e06bd, 14 October 2026):

Log	Environment	Coverage
dmuon-pr19-cuda-a800-18e06bd-20261014.log	4x A800-SXM4-80GB / SM80; Python 3.12.13; PyTorch 2.10.0+cu128; CUDA 12.8; driver 580.95.05; NCCL 2.27.5; CUTLASS DSL 4.6.1; pytest 9.1.1	The seven requested entrypoints above, unchanged defaults/tolerances; replicate-gate and storage appendices; CPU/Gloo clipping appendix. All exits 0.
dmuon-pr19-native-ascend950dt-18e06bd-20261014.log	4x Ascend950DT; openEuler aarch64; Python 3.11.6; PyTorch 2.10.0+cpu; torch_npu 2.10.0.post4; CANN 9.2.0-beta.2; driver 25.6.rc1; native HCCL	Current manual native suite and related gates, detailed below.

Native coverage: ascend_snapshot.py for both DDP and FSDP2 (each 48 NS, 48 clipping and five training cases with restore verification); ascend_sharded_accumulation.py (FP32/FP16/BF16, default/no_sync and active-state restore); ascend_distributed.py --topology all (28 calls, rank-local disk restart); and test_hsdp_correctness.py --device npu (default 22-call exact cross-mesh/canonical rank-0 restore gate). The replicate-gate, four-rank storage-ownership and eight NPU clipping-contract cases passed too, with direct exit 0.

pytest tests/native/ recorded one expected skip: the optional legacy migration comparator requires a genuine matching legacy snapshot/bootstrap that was unavailable. This comparison remains N/A, not a pass. The NPU log retains HCCL timeout-configuration and torch_npu serialization/map_location warnings; the exact restart assertions passed. It also retains an initial recorder TypeError that lost the first FSDP2 exit receipt, an admission-only blocked retry with no tests, and the single FSDP2 recovery run with a directly recorded exit 0. Library/test source was unchanged throughout.

Both files contain complete command stdout/stderr, attempt history, versions, timestamps and exits. Only internal machine/user/path identities were substituted; numeric test results, warnings and errors were retained. Embedded raw-log hashes precede those substitutions; public-file hashes are in the accompanying manifest.

Modified-file lint/format checks and the strict English/Chinese documentation build passed. Thank you again for identifying these issues.

[StarrickLiu]

Thanks for addressing every point so thoroughly. Re-verified at 18e06bd:

• CPU unit suite: 349 passed / 55 skipped; ruff check dmuon examples clean.
• 4× H100 (SM90, torch 2.10 + cu128): GPU unit suite 382 passed / 38 skipped (all environmental: quack not installed, fast-clip extension not built, no torch_npu). All requested distributed entrypoints pass (test_hsdp_correctness, test_hsdp_async_correctness, test_grad_accum, test_muon_step, test_checkpoint, test_ddp_correctness), plus test_hsdp_replicate_gate_accum (worst rel consumed-grad diff 1.3e-6) and test_gradient_storage_ownership (shard and HSDP).
• The #20 integration in 59d5b06 and the ownership transfer in 8105157 look correct.

SM90 + Quack remains unverified on both sides; Ascend results are from your logs. Approving.