Skip to content

feat(mtp): add native MTP conditioning for DSpark and DFlash - #23

Merged
zhouaihui merged 2 commits into
inclusionAI:mainfrom
zhouaihui:feat/qwen35-dmtp-training
Nov 12, 2026
Merged

zhouaihui merged 2 commits into
inclusionAI:mainfrom
zhouaihui:feat/qwen35-dmtp-training

Conversation

@zhouaihui

@zhouaihui zhouaihui commented Nov 11, 2026 •

Copy link
Copy Markdown
Collaborator

Closes #22

Outcome

Add opt-in mtp_conditioning to DSpark and DFlash so a frozen native Qwen3.5 MTP predictor
can condition the Draft through the existing training, recovery and native export workflow.
Token mode supplies the known next-token anchor; hidden mode also adds a normalized,
zero-initialized residual fusion. Inactive features are masked before trainable fusion.

Visible T1 is excluded from supervision. Both algorithms report
train/suffix_simulated_acc_len and eval/suffix_simulated_acc_len, with per-position
suffix_acc_0 / suffix_ploss_0 starting at T2. Checkpoint selection, metadata and reports
retain that suffix scope. These teacher-forced suffix proxies do not measure acceptance
of the native MTP proposal or complete speculative blocks.

Include Qwen3.5-9B Draft configs, a runnable recipe and support documentation. MTP-specific
initialization and binding hooks run only when conditioning is enabled; ordinary
DSpark/DFlash interfaces, supervision and metric names are preserved.

Verification

  • Independent nonzero-fusion equations check logits, loss and trainable parameter gradients
    for both algorithms, reusing the retained Draft backbone. Tests also cover activation
    checkpointing, padded inactive NaNs, frozen feature boundaries, native export/reload
    output equality, checkpoint metric selection and exact-resume history.
  • Add test_dmtp.py, test_training_resume_rng.py and test_evaluation_contract.py to both
    CPU CI and tools.recipes contracts.
  • Isolated Linux CPU/offline validation: CPU CI commands 631 passed and 16 passed;
    recipe contracts 421 passed, 5 skipped; final affected-model/recovery/assets suite
    292 passed, 21 skipped, plus 12 subtests. Lint and local documentation links passed.
  • Bounded CUDA checks on one NVIDIA L20X: 14 passed. The explicit-prefix Dense GQA
    fixture used head dimension 16 and the compiled FlexAttention TRITON path: FP32 independent
    logits/loss/gradient reference, plus FP32/BF16 checkpointing and padded-NaN gradient parity.
  • All 672 tracked source blobs and file modes matched the Linux-tested tree. Runtime:
    Python 3.12.3, Torch 2.13.0+cu130, Transformers 5.12.1.

Supported conditioned execution is Dense full GQA, CP1, MBS1, single-node REPLICATE,
decay objective, without packing or whole-model compilation. These checks make no new
multi-GPU, throughput or end-to-end decoding acceptance claim.

Compatibility and provenance

  • Conditioning is optional. Native assets retain the policy, fusion and frozen predictor
    state; exact resume validates the same-Target binding. Generic SGLang export and the
    ordinary proposal consumer reject conditioned assets pending an MTP-aware consumer.
  • The dual-anchor/fusion design is inspired by DMTP.
    This is an independent integration reusing ASpec's retained native predictor and DFlash
    backbone. Source headers and docs/framework/REFERENCES.md record the method and reuse;
    existing upstream notices and license texts are retained.

@zhouaihui
zhouaihui requested review from Longxmas and icavan November 11, 2026 09:08

Copy link
Copy Markdown
Collaborator

Three suggested follow-ups for 8e7f984:

  1. Suffix metric naming: Conditioned runs still report train/eval/simulated_acc_len, although this only measures the suffix given teacher-forced T1. Please use an explicit suffix name and update checkpoint selection, metadata, and per-position metric labels consistently.

  2. CI coverage: Please add tests/test_dmtp.py and tests/test_training_resume_rng.py to both the CPU CI suite and tools.recipes contracts; neither explicit test list currently includes them.

  3. Fusion validation: Please add an independent forward/backward reference test with nonzero fusion weights, covering logits, loss, and trainable parameter gradients. Also check gradient-checkpointing parity and output equality after native export/reload. For inactive blocks, mask MTP features before fusion and add a padded-NaN regression to verify finite, unchanged gradients on valid data.

Add optional token and hidden conditioning through the existing training,
recovery and native asset workflow, with frozen same-Target predictor bindings.

Validate the supported scope, mask known-token supervision, and cover native
loading, recovery and frozen parameter ownership with reusable checks.
Include recipes and support documentation; conditioned assets remain native-only.
@zhouaihui
zhouaihui force-pushed the feat/qwen35-dmtp-training branch from 8e7f984 to a234ce3 Compare November 12, 2026 07:36
@zhouaihui

zhouaihui commented Nov 12, 2026 •

Copy link
Copy Markdown
Collaborator Author

Addressed in 6ce63d1:

  1. Suffix metric names are consistent across training, evaluation and checkpoint selection.
  2. MTP and resume tests are included in CPU CI and recipe contracts.
  3. Added nonzero-fusion, checkpointing, export/reload and padded-NaN coverage; inactive features are masked before fusion.

CI, CPU regressions and bounded CUDA checks passed. Unconditioned DSpark/DFlash behavior is preserved.

Use explicit suffix metric names for conditioned training, evaluation and
checkpoint selection while preserving ordinary DSpark/DFlash metric names.
Mask inactive features before fusion to prevent non-finite gradients.

Cover nonzero fusion outputs and parameter gradients, activation checkpointing,
native asset roundtrips and padded NaNs. Include MTP and resume contracts in
CPU CI and the recipe contract suite.
@zhouaihui
zhouaihui force-pushed the feat/qwen35-dmtp-training branch from a234ce3 to 6ce63d1 Compare November 12, 2026 07:54

Copy link
Copy Markdown
Collaborator

Thanks for the update. Reviewed a234ce3 and the CI results; all three follow-ups are addressed: suffix metric naming, CI coverage, and fusion validation with inactive-feature masking. Looks good to merge.

@zhouaihui
zhouaihui merged commit 61bf537 into inclusionAI:main Nov 12, 2026
2 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.

Support native MTP conditioning for DSpark and DFlash

2 participants