Add Ascend NPU support for the Domino training pipeline (Qwen3.5-4B)

[curnane-lab]

Summary

This PR brings the Domino training pipeline up on Ascend NPU. The adaptation is intentionally narrow: it keeps the existing CUDA path byte-for-byte identical and only swaps in NPU equivalents when the active accelerator is detected as npu. No new public APIs are added to the user-facing training entrypoint.

Scope

Area	Change
specforge/modeling/target/target_head.py	Accelerator-aware device + text_config fallback
specforge/core/domino.py, specforge/core/dflash.py	NPU bf16 GRU workaround, FlexAttention guard
scripts/train_domino.py	Replace hard-coded cuda literals
configs/qwen3.5-4b-domino.json, examples/run_qwen3.5_4b_domino_npu.sh	Qwen3.5-4B Domino draft config and Ascend launcher

Design notes

Device / backend resolution

A small helper pair lives in specforge.utils:

get_device_type() -> str          # cuda | npu | cpu
get_local_device() -> torch.device

Resolution order is SPECFORGE_DEVICE env var ➜ torch.cuda ➜ torch.npu ➜ cpu. init_distributed() consumes these helpers and selects nccl / hccl / gloo accordingly; init_device_mesh and DeviceMesh.from_group use the resolved device type instead of a hard-coded "cuda". A SPECFORGE_DIST_BACKEND env var is honored as an explicit override.

yunchang is imported softly because the package is currently CUDA only. When it is unavailable the seq-parallel groups fall back to the draft SP group; set_seq_parallel_pg becomes a no-op.

Domino bf16 GRU workaround

Ascend's DynamicGRU kernel does not yet support bfloat16. To keep training in bf16 end-to-end without breaking FSDP's mixed-dtype guard, OnlineDominoModel constructs a private float16 nn.GRU sandbox and attaches it via object.__setattr__ so it is not registered as a submodule. On every forward we sync the prefix-GRU weights into the sandbox, run inference in fp16, and cast back to the input dtype. The overhead is a single weight copy per step.

FlexAttention on NPU

FLEX_ATTENTION_AVAILABLE is forced to False at module-import time when running on NPU. OnlineDFlashModel.forward already raises a descriptive error when the backend is not available, which now also fires on NPU and tells the user to choose --attention-backend sdpa or --attention-backend eager.

TargetHead

TargetHead.from_pretrained now uses get_local_device() instead of .cuda(). The class also restores the upstream text_config / hidden_size / vocab_size fallback so VLM-style configs continue to load (the fallback is a no-op for plain text models such as Qwen3.5).

Qwen3.5-4B example

configs/qwen3.5-4b-domino.json matches the public Qwen3.5-4B architecture (intermediate_size: 9216, target_layer_ids: [3, 11, 19, 27, 31] — the four full-attention layers plus the final layer). The sample launcher examples/run_qwen3.5_4b_domino_npu.sh exposes TARGET_MODEL_PATH / TRAIN_DATA_PATH as required env vars, exports ASCEND_RT_VISIBLE_DEVICES and a conservative PYTORCH_NPU_ALLOC_CONF, and passes --attention-backend sdpa.

Compatibility

• CUDA path is unchanged. get_device_type() returns cuda whenever torch.cuda.is_available() is true, and every code path that used to read "cuda" literals continues to do so.
• The PR does not modify any DFlash-only training code path, configs, or launchers.
• No new dependencies are introduced.

How to test

# CUDA (smoke)
bash examples/run_qwen3_8b_domino_online.sh

# Ascend NPU (Qwen3.5-4B)
export TARGET_MODEL_PATH=/path/to/Qwen3.5-4B
export TRAIN_DATA_PATH=/path/to/dataset
bash examples/run_qwen3.5_4b_domino_npu.sh

Accuracy Test

Hardware

8× Ascend 910C NPU

Target model

Qwen3.5-4B

Training data preparation

Qwen3.5-4B (target model) with the HF backend is used to regenerate responses for 10K samples from Open-PerfectBlend. The regenerated QA pairs serve as training data for both the DFlash and Domino draft models, keeping the data pipeline identical across the two baselines.

Baseline: DFlash on NPU (PR #562 + #559)

The DFlash NPU baseline uses the existing qwen3.5-4b-dflash.json config and the HF backend with SDPA attention.

This PR: Domino on NPU

The Domino NPU run uses configs/qwen3.5-4b-domino.json added in this PR, with the same target model, regenerated 10K dataset, and comparable hyperparameters.

Head-to-head comparison

The following figure overlays the accuracy / loss / learning rate curves of Domino NPU (this PR) and DFlash NPU (baseline) on the same axes.

Summary

Under the same 10K regenerated Open-PerfectBlend data, Qwen3.5-4B target model, and identical LR schedule, Domino NPU (red) consistently outperforms the DFlash NPU baseline (blue): it achieves a lower final training loss (~1.4 vs ~2.1) and a substantially higher accuracy (~0.65 vs ~0.45) after ~5,500 steps.

Commits

1. feat(target): use accelerator-aware device for TargetHead loading
2. feat(domino): work around Ascend NPU bf16 GRU and flex_attention
3. feat(scripts): make train_domino device-aware
4. feat(qwen3.5): add 4B Domino draft config and Ascend launcher

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[curnane-lab]

Hi @claude @jiapingW @wenqf11 , would you be up for reviewing this PR?