Skip to content

add DFlash2 npu training recipe - #789

Merged
jiapingW merged 3 commits into
sgl-project:mainfrom
curnane-lab:dflash2_npu
Sep 8, 2026
Merged

jiapingW merged 3 commits into
sgl-project:mainfrom
curnane-lab:dflash2_npu

Conversation

@curnane-lab

@curnane-lab curnane-lab commented Aug 29, 2026

Copy link
Copy Markdown
Collaborator

Motivation

Follow-up to #772: make DFlash 2 trainable on Ascend NPU by adding the missing recipe layer — one draft config, one external-server recipe, and a tutorial section. DFlash 2 reuses the DFlash capture path and the NPU disaggregated plumbing from #722 unchanged. Mirrors the Domino/DSpark NPU recipes in #726 / #727.

Modifications

  • configs/qwen3.5-4b-dflash2.json: DFlash2DraftModel config for Qwen3.5-4B — 5 full-attention layers on the same target layers as the 4B DFlash config ([1, 8, 15, 22, 29]), 16, conv_kernel_size: 2, conv_group_size: 16, selector_rank: 256, selector_top_k: 16`.
  • examples/configs/online/disaggregated/external/qwen3.5-4b-dflash2-online-npu.yaml: field-aligned with the existing DFlash NPU recipe, except attention_backend: sdpa (flex_attention is unavailable on Ascend) and dflash2_selector_loss_alpha: 1.0.
  • docs/basic_usage/Ascend/ascend_npu.md: a "Training DFlash 2 drafts" section — capture setup unchanged from DFlash; sdpa mandatory; the convolutions and selector are stock torch ops; serving needs SGLang with DFlash 2 ([Spec] DFlash2: local convolution + candidate selector sglang#35371, merged).

Related Issues

Follow-up to #772; builds on #722; serving side is sgl-project/sglang#35371.

Accuracy Test

Config / recipe / docs only; no model-side code. End-to-end validation on Ascend NPU with the checked-in recipe (Qwen3.5-4B target + DFlash 2 draft, ~1.75k optimizer steps): train/loss ~6.9 → ~1.25, train/acc ~0.09 → ~0.71, selector_loss 2.05 → ~0.25, selector_accuracy ~0.36 → ~0.93, selector_coverage ~0.44 → ~0.90; bounded grad_norm, no NaN or stall.

DFlash2 NPU training curves

Benchmark & Profiling

N/A (no performance-affecting code changes).

Checklist

  • Format your code according to the Code Formatting with Pre-Commit.
  • Add unit tests as outlined in the Running Unit Tests. (Recipe-only change; no new code paths to cover.)
  • Update documentation / docstrings / example tutorials as needed, according to Writing Documentation.
  • Provide throughput / latency benchmark results and accuracy evaluation results as needed, according to Benchmark and Profiling and Accuracy Results. (Training-side NPU validation in Accuracy Test; serving benchmark out of scope.)
  • For reviewers: If you haven't made any contributions to this PR and are only assisting with merging the main branch, please remove yourself as a co-author when merging the PR.
  • Please feel free to join our Slack channel at https://sgl-fru7574.slack.com/archives/C09784E3EN6 to discuss your PR.

@curnane-lab curnane-lab changed the title Dflash2 npu add dFlash2 npu training recipe Aug 29, 2026
@curnane-lab curnane-lab changed the title add dFlash2 npu training recipe add DFlash2 npu training recipe Aug 29, 2026
@curnane-lab
curnane-lab marked this pull request as ready for review August 30, 2026 11:08
@jiapingW
jiapingW merged commit b95365c into sgl-project:main Sep 8, 2026
3 checks passed
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.

3 participants