[Feature] NPU (Ascend) training performance optimizations for Qwen3.5 dense text model #2118

Motivation

Training Qwen3.5-class models (GatedDeltaNet layers, dense MLPs, a 154,880-vocab lm head) on Ascend NPUs currently hits four gaps in the v1 ops/engine stack:

1. GDN stack dead ends. get_zero_centered_rms_norm_fn() in xtuner/v1/ops/rms_norm raises NotImplementedError on NPU, so the in-repo zero-centered RMSNorm route is unusable on Ascend; and the GatedDeltaNet op stack (causal conv1d, chunked gated delta rule, gated RMSNorm) falls back to slow fp32 torch decompositions — causal conv1d has no NPU kernel at all in the package. The native fp32 decomposition OOMs at step 1 on all 16 ranks at the default GA=4/mbs=2 config.
2. Eager activation chain. The dense-MLP activation runs as a two-pass silu(gate) + mul forward and a two-pass backward — four full-tensor memory passes where fused one-pass kernels need one each.
3. No NPU liger loss. loss_cfg.mode == "liger" imports liger_kernel's stock CUDA class unconditionally, so the fused-linear-cross-entropy path (which avoids materializing full logits over the 154,880-token vocab) has no working NPU implementation.
4. FSDP2 all-gather memory strand. torch_npu's caching allocator segregates cached blocks per stream, so FSDP2's full-param all-gather on the separate all_gather_stream strands reserved memory that empty_cache cannot reclaim — pure per-stream fragmentation, at no benefit.
5. No host-side CPU affinity. Every rank drives its NPU from the host — aclnn op dispatch, tensor metadata, and the device's companion host threads are all host CPU work; unpinned ranks can land on NUMA nodes remote from their NPU, and a single straggler rank paces the synchronized step.

Proposal

Six independent PRs move the production NPU implementations into the package — one per gap, each opt-in/env-gated, each leaving non-NPU behavior unchanged — plus two stacked follow-up PRs (#2240, #2241):

1. [Feature] NPU fused zero-centered RMSNorm (XTUNER_NPU_FUSED_ZC_NORM) #2119 — NPU fused zero-centered RMSNorm: a single npu_rms_norm pass with the (1 + w) scale folded into the fp32 gamma (bf16 output rounds once); fp32-decomposition fallback; gated by XTUNER_NPU_FUSED_ZC_NORM (default on).
2. [Feature] NPU AscendC op stack for GatedDeltaNet (causal conv1d / chunk delta rule / fused gated norm) #2120 — NPU AscendC op stack for GatedDeltaNet: causal conv1d with three-level in-module dispatch, the chunked gated delta rule kernel wired to fla_npu AscendC ops plus mindspeed triton helpers with a device-side chunk-indices cache, and a fused gated RMSNorm. Dispatch is gated on get_device() == 'npu' plus XTUNER_NPU_* env knobs (all default on).
3. [Feature] Fused SwiGLU Triton kernels for MLP activation (XTUNER_NPU_FUSED_SWIGLU) #2128 — Fused SwiGLU Triton kernels: one-pass silu(gate)·up forward and backward for the dense-MLP activation. Device-agnostic plain Triton (no torch_npu dependency), opt-in via XTUNER_NPU_FUSED_SWIGLU (default off), read live so runtime injection works.
4. [Feature] NPU liger fused-linear-cross-entropy loss (XTUNER_LOSS_MODE=liger) #2129 — NPU liger fused-linear-cross-entropy loss: a get_flce_loss_cls() device dispatch — NPU returns an in-package NPUFusedLinearCrossEntropyLoss driving liger-ascend's triton kernels (the stock ascend plain-CE kernel miscomputes loss_1d in multi-rank training, so the forward computes CE via torch on the matmul logits), CUDA keeps liger_kernel's stock class. Chunked forward behind a full-logits byte gate; no device env knob by design — the fallback arm is numerically wrong.
5. [Feature] Route FSDP2 unshard to the compute stream (XTUNER_FSDP_UNSHARD_ASYNC_OP) #2130 — Route FSDP2 unshard to the compute stream: XTUNER_FSDP_UNSHARD_ASYNC_OP=1 (default off) issues FSDP2's parameter all-gather on the compute stream so the per-stream strand is never created; wait_for_unshard + explicit forward prefetch keep the overlap.
6. [Feature] Bind NPU training ranks to NUMA-local CPU slices (XTUNER_NPU_CPU_BIND) #2131 — Bind NPU training ranks to NUMA-local CPU slices: a stdlib-only npu_cpu_binder pins each rank process and its device's host threads (acl/release/hccp) to an 8-core NUMA-local slice, wired into _try_bind_numa so all training modes bind; XTUNER_NPU_CPU_BIND (default on), bind failures degrade to a log line.
7. [Fix] fla_npu >= 26.10 compatibility for the GDN AscendC stack (chunk loader + causal conv backends) #2240 — fla_npu ≥ 26.10 compatibility for the GDN AscendC stack (stacked on [Feature] NPU AscendC op stack for GatedDeltaNet (causal conv1d / chunk delta rule / fused gated norm) #2120): fla_npu's 26.10 line splits its API surface — stable raw ops move to fla_npu.ops.ascendc.* and the legacy torch.ops.npu.* registrations need an explicit load_legacy_torch_ops(). The causal-conv backend becomes a dual-backend resolver decided once per process (stable API on 26.10, legacy torch.ops on 26.6), the chunk loader degrades cleanly when the embedded legacy extension is absent, and the startup self-probe upgrades from isfinite-only to a numeric check against an F.conv1d autograd reference, so "finite but wrong" kernels are rejected into the unmasked-correction fallback instead of surfacing as mid-training NaN. XTUNER_NPU_GDN_API no longer selects the conv backend. Behavior on fla_npu 26.6 is unchanged.
8. [Feature] Optimize the NPU liger FLCE backward and add a fused LM z-loss #2241 — NPU liger FLCE backward optimization + fused LM z-loss (stacked on [Feature] NPU liger fused-linear-cross-entropy loss (XTUNER_LOSS_MODE=liger) #2129): backward scratch-buffer reuse and vocabulary-slab grad_weight accumulation remove the two large logits-path transients (~1.9 GiB fresh alloc per chunk + the ~2 GiB [V, H] temporary; −4.6 GB max_memory e2e), the backward recompute chunk ties to LOSS_CHUNK_SIZE, a latent reduction='none' multi-chunk grad_output indexing bug is fixed (chunks > 0 read wrong rows; ce_loss.py uses sum, so latent only), the LM z-loss (CELossConfig.z_loss_coefficient, default off) fuses CE+LSE² into one forward kernel, and the no-weight backward call gains 0.8.1/0.8.5 dual support via HAS_TAIL signature detection.

Dependency (NPU only)

• flash-linear-attention-npu (fla_npu) v26.6.0, compiled and installed from source ([Feature] NPU fused zero-centered RMSNorm (XTUNER_NPU_FUSED_ZC_NORM) #2119 / [Feature] NPU AscendC op stack for GatedDeltaNet (causal conv1d / chunk delta rule / fused gated norm) #2120); the 26.10 line (decoupled runtime, stable fla_npu.ops.ascendc.* raw ops + explicit legacy loader) is supported by [Fix] fla_npu >= 26.10 compatibility for the GDN AscendC stack (chunk loader + causal conv backends) #2240.
• liger-kernel 0.8.1 ([Feature] NPU liger fused-linear-cross-entropy loss (XTUNER_LOSS_MODE=liger) #2129) and 0.8.5 ([Feature] Optimize the NPU liger FLCE backward and add a fused LM z-loss #2241 — HAS_TAIL signature detection; the fused z-loss kernels need >= 0.8.4) — the versions whose liger_kernel.ops.backends._ascend ships the ascend triton kernels used by the NPU FLCE; CUDA installs keep the version they already use.

Both are imported lazily inside the NPU dispatch branches only; non-NPU environments need no new dependency.

Results (Qwen3.5 dense text model, 16 Ascend NPUs, frozen launcher recipe)

• Step-1/step-2 losses bitwise-identical across A/B runs (anchor gate), 0 NaN; 1000-step run exits 0 with 0 NaN.

• Global throughput 179.2K tok/s on the default GA=4/mbs=2 config with the GDN stack enabled; the fused RMSNorm is +10.6% over the native fp32 decomposition at GA=2/mbs=1, and is what makes the default config fit in memory at all (the native decomposition OOMs at step 1 on all 16 ranks).

• CPU binding (XTUNER_NPU_CPU_BIND) is scheduling-only — step-1/step-2 losses bitwise-identical, peak memory unchanged — and adds +1.7% global tok/s at 256 ranks (16 nodes × 16 Ascend NPUs, 100-step A/B).

• fla_npu 26.10 compat ([Fix] fla_npu >= 26.10 compatibility for the GDN AscendC stack (chunk loader + causal conv backends) #2240): on a 26.10 dev wheel (256 Ascend NPUs = 16 nodes × 16, Qwen3.5 dense text, 10-step benchmark, mbz=1/gas=2/gbs=512) the GDN chunk route resolves via the stable API, 10/10 steps, 0 NaN, benchmark success; the broken stable conv adapter is rejected by the numeric startup probe and the unmasked-correction fallback engages (error band fwd 0.094 / dx 0.125 / dw 1.0). On fla_npu 26.6 the conv route is bit-exact at dispatcher level (dx/dw within the same few-ulp band as [Feature] NPU AscendC op stack for GatedDeltaNet (causal conv1d / chunk delta rule / fused gated norm) #2120).

• Liger FLCE backward optimization ([Feature] Optimize the NPU liger FLCE backward and add a fused LM z-loss #2241; same implementation A/B'd on an internal fork, GLM-5.3-Flash 30B VL, 16 NPUs, 10 steps): step-1 loss bitwise identical to baseline, max_memory −4.6 GB sustained (27.6 vs 32.2 GB steady), steady tok/s +4.2% (trend only).

Env knobs (per PR)

Notes: #2129 selects its NPU path via the loss config (loss_cfg.mode == "liger", in our runs injected as XTUNER_LOSS_MODE=liger) — the device dispatch itself has no env knob by design; its three knobs above only tune the chunked forward. #2120 additionally reuses the existing XTUNER_HF_IMPL knob for bitwise-parity comparison against the fla/HF reference implementations. #2241 additionally adds the config field CELossConfig.z_loss_coefficient (default 0.0) for the fused LM z-loss — a config, not an env knob; in liger mode the z-loss enters the gradient and logging (lm_z_loss) but not the loss value. #2240 narrows XTUNER_NPU_GDN_API to the GDN chunk op resolution only (it no longer selects the causal-conv backend, which is auto-resolved per fla_npu version instead).

PRs

• [Feature] NPU fused zero-centered RMSNorm (XTUNER_NPU_FUSED_ZC_NORM) #2119 — NPU fused zero-centered RMSNorm (XTUNER_NPU_FUSED_ZC_NORM)

• [Feature] NPU AscendC op stack for GatedDeltaNet (causal conv1d / chunk delta rule / fused gated norm) #2120 — NPU AscendC op stack for GatedDeltaNet (causal conv1d / chunk delta rule / fused gated norm)

• [Feature] Fused SwiGLU Triton kernels for MLP activation (XTUNER_NPU_FUSED_SWIGLU) #2128 — Fused SwiGLU Triton kernels for MLP activation (XTUNER_NPU_FUSED_SWIGLU)

• [Feature] NPU liger fused-linear-cross-entropy loss (XTUNER_LOSS_MODE=liger) #2129 — NPU liger fused-linear-cross-entropy loss (XTUNER_LOSS_MODE=liger)

• [Feature] Route FSDP2 unshard to the compute stream (XTUNER_FSDP_UNSHARD_ASYNC_OP) #2130 — Route FSDP2 unshard to the compute stream (XTUNER_FSDP_UNSHARD_ASYNC_OP)

• [Feature] Bind NPU training ranks to NUMA-local CPU slices (XTUNER_NPU_CPU_BIND) #2131 — Bind NPU training ranks to NUMA-local CPU slices (XTUNER_NPU_CPU_BIND)

• [Fix] fla_npu >= 26.10 compatibility for the GDN AscendC stack (chunk loader + causal conv backends) #2240 — fla_npu >= 26.10 compatibility for the GDN AscendC stack (chunk loader + causal conv backends; stacked on [Feature] NPU AscendC op stack for GatedDeltaNet (causal conv1d / chunk delta rule / fused gated norm) #2120)

• [Feature] Optimize the NPU liger FLCE backward and add a fused LM z-loss #2241 — NPU liger FLCE backward optimization + fused LM z-loss (XTUNER_LIGER_GRAD_WEIGHT_CHUNK, CELossConfig.z_loss_coefficient; stacked on [Feature] NPU liger fused-linear-cross-entropy loss (XTUNER_LOSS_MODE=liger) #2129)