Skip to content

AutoEP: generate large NPU permutation indices on device - #8822

Closed
hazelduan wants to merge 1 commit into
deepspeedai:masterfrom
hazelduan:codex/npu-autoep-permutation-20261014
Closed

hazelduan wants to merge 1 commit into
deepspeedai:masterfrom
hazelduan:codex/npu-autoep-permutation-20261014

Conversation

@hazelduan

Copy link
Copy Markdown

AutoEP's NPU permutation path needs a device implementation instead of invoking CUDA Triton or copying large source/expert count groups to the host. This change builds NPU permutation indices with PyTorch tensor operations when NPU-resident counts contain at least 64 source/expert groups. Smaller groups and host counts use the CPU reference path; indices and aligned counts are transferred to the token device before reordering.

It preserves ordering, padding and token gradients, including empty experts and strided inputs. The threshold is internal; no configuration or optional package is added. Other dispatch steps can still synchronize with the host.

Split from #8784; no Qwen or HiFloat8 dependency. Related RFC: #8777. Base: 3b80932; validated source: b7bee89. An open-PR search for NPU permutation returned no match on 2026-10-14. PR #8784 retains an NPU host correctness fallback, so dispatch selection overlaps: when combining, retain this threshold selection and the shared non-CPU device-transfer guard.

Validation:

  • macOS CPU, Python 3.12.2: 10 passed, 37 device-only skips. Changed-file pre-commit and git diff --check passed.
  • One Ascend 950PR NPU (driver 25.7.rc1, CANN 9.1, Python 3.11.15, PyTorch/TorchNPU 2.11.0): standalone suite 21 passed, 26 CUDA-only skips. Exact CPU-reference indices, aligned sizes, forward values and input gradients cover empty groups, strided input and host counts. Command: PYTHONPATH=<isolated-permutation-repo> python -m pytest tests/unit/v1/moe/test_autoep_permute_rows.py -q -rs.

The private combined-stack validation also completed Qwen3.5-35B-A3B text-only EP4/ZeRO2 for 30 matched steps. All 40 layers execute; only routed/shared experts in layers 0/39 train (12 tensors/rank), batch 1/GAS1, sequence length 256, eager attention/experts, no activation checkpointing or checkpoint save, seed 42, AdamW 1e-5/cosine, same model weights and 400-row data. Every rank has exact input hashes, loss, gradient norms, native operator counts and all final trainable tensors. Maximum-rank steady-step median (excluding 5 warm-up steps) is 4.1886s baseline vs 4.1906s candidate (+0.046%); training NPU peak is 33,487,014,400 bytes and initialization peak is 86,776,612,352 bytes in both. This combined-stack result does not attribute an independent speedup or numerical improvement to this split.

Native commands and logs are retained at /data/dxq/hif8_a5_20260914/experiments/quantize_convergence_20261014_01: bash native_gates.sh, bash paired_ep4.sh, qwen_native_stack.log, permutation_native.log, comparison_ep4_baseline_r1_ep4_candidate_r1.json. The runner bind-mounts the original Conda environment (dxq_hif8_torchao_py311) and CANN/driver; it changes no global packages or device permissions. Standalone runtime file hashes match each immutable source revision below. The model run uses a private composition of main PR and split runtime hunks, with the overlapping dispatch hunks reconciled; this composition is not the main PR contents.

Both model comparison variants include permutation optimization, so the paired run does not measure its independent benefit. The standalone tests establish index generation and local reordering. CUDA/other accelerator execution, distributed/offload combinations, activation checkpointing and isolated end-to-end performance/memory effects remain untested. No speedup is claimed.

AI assistance was used to implement and review this change; the human submitter owns review and validation.

Use tensor index generation for NPU-resident counts with at least 64 source/expert groups. Keep the CPU reference for smaller and host count groups, and move indices to non-CPU token devices before reordering.

Ascend 950PR standalone reference and gradient checks: 21 passed, 26 CUDA-only skips. Changed-file pre-commit passed. No end-to-end speedup claim.

Signed-off-by: hazelduan <duanxueqi@outlook.com>
@hazelduan hazelduan closed this Oct 14, 2026
@hazelduan
hazelduan deleted the codex/npu-autoep-permutation-20261014 branch October 14, 2026 14:13
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant