feat: add mtp support

[curnane-lab]

Motivation

This PR adds training support for Multi-Token Prediction (MTP) to SpecForge, targeting Qwen3.5-4B. MTP is natively supported by Qwen3.5 and by SGLang/vLLM through per-target *_mtp.py modules, so the trained weights must match the target's native checkpoint layout to be loadable at inference time.

Modifications

• specforge/modeling/draft/mtp.py: Qwen3.5 MTP draft model (Qwen3_5MTPDraftModel) matching the official Qwen3.5-4B architecture (head_dim=256, attn_output_gate, partial rotary). Saves weights in the flat mtp.layers.* layout required by SGLang/vLLM.
• specforge/core/mtp.py: OnlineMTPModel training wrapper (next-token shift, CE loss, per-position accuracy metrics).
• scripts/train_mtp.py: end-to-end MTP training script with FSDP, BF16 optimizer, optional native-MTP finetune init, and frozen/shared embed/lm_head.
• specforge/modeling/target/mtp_target_model.py: standalone MTP data generators for SGLang, HF, and custom target backends, decoupled from Eagle3TargetModel.
• specforge/modeling/target/eagle3_target_model.py: minimal changes — CPU-first HF target loading for NPU memory, and capture_aux_hidden_states flag for SGLang extend().
• scripts/merge_mtp_to_base.py: merges trained MTP weights back into the base model checkpoint for direct serving.
• configs/qwen3.5-4b-mtp.json & examples/run_qwen3.5_4b_mtp_online_npu.sh: draft config and NPU training example.

Related Issues

• Related to [Feature] mtp support #662
• Related to [Feature] will support qwen3.5-mtp training and qwen3.5-eagle3 training? #497

Accuracy & Performance Evaluation

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 MTP head.

Training Configuration

SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
ROOT_DIR=$(dirname $SCRIPT_DIR)

export HF_DATASETS_CACHE=$ROOT_DIR/cache/hf_datasets
export TORCHINDUCTOR_CACHE_DIR=$ROOT_DIR/cache/compiled_kernels
export SPECFORGE_DATA_NUM_PROC=64

ATTENTION_BACKEND=${2:-sdpa}
NUM_GPUS=${1:-8}

torchrun \
    --standalone \
    --nproc_per_node $NUM_GPUS \
    $ROOT_DIR/scripts/train_mtp.py \
    --target-model-path PATH/TO/Qwen3.5-4B \
    --draft-config-path $ROOT_DIR/configs/qwen3.5-4b-mtp.json \
    --train-data-path PATH/TO/TRAIN_DATA.jsonl \
    --output-dir $ROOT_DIR/outputs/qwen3.5-4b-mtp \
    --num-epochs 10 \
    --batch-size 2 \
    --accumulation-steps 4 \
    --learning-rate 6e-4 \
    --warmup-ratio 0.04 \
    --max-grad-norm 1.0 \
    --max-length 3072 \
    --chat-template qwen3.5 \
    --attention-backend $ATTENTION_BACKEND \
    --ploss-decay 1.0 \
    --log-interval 50 \
    --save-interval 10000 \
    --report-to tensorboard \
    --target-model-backend hf \
    --trust-remote-code \
    --embedding-key model.language_model.embed_tokens.weight \
    --lm-head-key model.language_model.embed_tokens.weight \

Evaluation was performed on GSM8K (500 samples, max_new_tokens=512, temperature=0.0) for both the fine-tuned MTP head and the original Qwen3.5-4B native MTP head, served with qwen3_5_mtp speculative decoding (num_speculative_tokens=1, enforce_eager=true).

Metric	Trained MTP (A)	Original MTP (B)	A / B
Avg TTFT (s)	0.111	0.111	1.00x
Avg Decode Time (s)	5.198	5.086	0.98x
Avg Decode Speed (t/s)	98.22	100.38	0.98x

The trained MTP head — initialized and trained from scratch with --no-init-from-native-mtp — achieves on-par end-to-end latency and throughput with the original native MTP head (within 2%), confirming that the randomly-initialized MTP head can be trained to match the native head and that the trained weights are correctly loaded and functionally aligned with the serving path.

Checklist

• Format your code according to the Code Formatting with Pre-Commit.
• Add unit tests as outlined in the Running Unit Tests.
• 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.
• 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.

[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!

[jiapingW]

@curnane-lab Great work! We believe that fine-tuning MTP is now a very important task. The current Qwen3.5 and later models have very high MTP performance. Modifying the current PR to fine-tune MTP and integrating it into the current main branch would be very meaningful. However, we believe that the amount of data required for MTP pretraining is too large, similar to Eagle3. No dedicated implementation is needed. We look forward to your feedback.

[curnane-lab]

@curnane-lab Great work! We believe that fine-tuning MTP is now a very important task. The current Qwen3.5 and later models have very high MTP performance. Modifying the current PR to fine-tune MTP and integrating it into the current main branch would be very meaningful. However, we believe that the amount of data required for MTP pretraining is too large, similar to Eagle3. No dedicated implementation is needed. We look forward to your feedback.

Agreed — fine-tuning is the useful path here. I'll drop the from-scratch pretraining part, keep only native-MTP init, and rework it as a strategy in the unified specforge train flow (typed YAML + disaggregated data plane) instead of the standalone script. Will push the rework after rebasing onto latest main.

[curnane-lab]

@curnane-lab Great work! We believe that fine-tuning MTP is now a very important task. The current Qwen3.5 and later models have very high MTP performance. Modifying the current PR to fine-tune MTP and integrating it into the current main branch would be very meaningful. However, we believe that the amount of data required for MTP pretraining is too large, similar to Eagle3. No dedicated implementation is needed. We look forward to your feedback.

Agreed — fine-tuning is the useful path here. I'll drop the from-scratch pretraining part, keep only native-MTP init, and rework it as a strategy in the unified specforge train flow (typed YAML + disaggregated data plane) instead of the standalone script. Will push the rework after rebasing onto latest main.

Hi @jiapingW , the rework you suggested is done: MTP is now a strategy in the unified specforge train flow (typed YAML + disaggregated data plane), the from-scratch pretraining path has been dropped, and training always fine-tunes the native mtp.* head. Could you take another look when available?

Online managed-local training on Ascend NPU (16x A3)

Qwen3.5-4B MTP head trained end-to-end on the managed-local NPU stack — a single specforge train command owns Mooncake, 6 TP=1 capture servers (devices 0-5), and 10 trainer ranks (devices 6-15). Recipe: examples/configs/qwen3.5-4b-mtp-disaggregated-npu.yaml.

• Data: 10K Open-PerfectBlend regenerated prompts × 10 epochs, max_length 32768; global batch 80/step, ~1.25K optimizer steps, ~1.2h total.
• The head was initialized from the native mtp.* weights shipped in the Qwen3.5-4B checkpoint (native-MTP fine-tune); embed_tokens / lm_head are shared with the target and frozen.
• Convergence: draft accuracy 0.67 -> ~0.97, loss ~1.45 -> ~0.05; grad norm settles below 2 after warmup; cosine LR 6e-4 -> 0. No NaN, no stall.
• Throughput: ~3.3s/optimizer step (~1.05K steps/hour, ~22-25 samples/s); durable ack ~26ms. Train compute is ~1.3s/step against ~2-4.5s data wait, so the pipeline is capture-bound — more capture servers is the next lever.

[jiapingW]

Can you merge the commits to less?

[curnane-lab]

Can you merge the commits to less?

Done.

[curnane-lab]

Modifications

• specforge/algorithms/mtp/: register MTP as a built-in algorithm in the unified specforge train flow — draft builder, native-mtp.* weight initialization from the target checkpoint, shared/frozen target embeddings, capture-layer resolution, and the AlgorithmSpec/providers wiring.
• specforge/training/strategies/base.py: MTPTrainStrategy adapting OnlineMTPModel to the trainer StepOutput contract (requires input_ids / loss_mask / target_last_hidden_states; per-position accuracy metrics; checkpoint filter persists stripped embed_tokens.* + mtp.* keys matching the native serving layout).
• specforge/modeling/draft/mtp.py: Qwen3_5MTPDraftModel matching the official Qwen3.5-4B architecture (head_dim=256, attn_output_gate, partial_rotary_factor=0.25). Weights are saved in the flat mtp.layers.* layout required by SGLang/vLLM.
• specforge/core/mtp.py: OnlineMTPModel training wrapper (next-token shift, CE loss, per-position accuracy metrics).
• specforge/algorithms/common/dflash_family_data.py: MTP offline reader, normalizer, and collator for the dflash-family feature pipeline.
• specforge/algorithms/builtin.py, specforge/runtime/contracts.py, specforge/modeling/draft/__init__.py: register the mtp strategy (builtin registry, DraftStrategyName literal, draft-model export).
• scripts/merge_mtp_to_base.py: merge trained MTP weights back into the base checkpoint for direct serving (shared-embedding handling, old MTP removal, text_config patching).
• configs/qwen3.5-4b-mtp.json: draft config matching the official Qwen3.5-4B MTP head.
• examples/configs/qwen3.5-4b-mtp-disaggregated-npu.yaml + examples/README.md: managed-local NPU recipe — one specforge train command owns Mooncake, a patched SGLang capture server (device 0), and a 14-rank trainer (devices 2-15) on 16x A3; 8-card fallback noted in comments. Registered in the recipe table.
• Tests: test_builtin_providers.py (mtp added to the builtin provider contract), test_launch_topology.py / test_unified_feature_reachability.py (golden topology + reachability registration for the new recipe).

[jiapingW]

@curnane-lab We therefore do not plan to introduce separate strategies such as qwen_mtp or glm_mtp, nor a separate "family adapter" layer. And we want to rename algorithms/common/dflash_family_data.py to algorithms/common/hidden_states_data.py.

The intended structure is approximately:

specforge/
├── algorithms/
│   └── mtp/
│       └── providers.py
├── modeling/
│   ├── draft/
│   │   └── mtp/
│   │       ├── base.py
│   │       ├── qwen3_5.py
│   │       └── glm_moe_dsa.py
│   └── target/
│       ├── checkpoint.py
│       └── target_utils.py
├── core/
│   └── mtp.py
├── training/
│   └── strategies/
│       └── base.py
└── export/
    └── mtp.py

The responsibilities would be:

- algorithms/mtp/providers.py
    - Declares the MTP feature contract and compatible draft architectures.
    - Wires together draft construction, data providers, and MTPTrainStrategy.
    - Should remain model-family agnostic.
    - Owns the MTP reader, normalizer, and collator.
    - The common feature contract remains input_ids, loss_mask, and target_last_hidden_states.

- core/mtp.py
    - Owns the architecture-independent token/label shifting, CE loss, and metrics.
    - It should accept a generic MTP draft module and should not import Qwen3_5MTPDraftModel.

- modeling/draft/mtp/<architecture>.py
    - Owns the actual trainable MTP network and native weight layout for one architecture family.
    - Qwen3.5 and Qwen3.6 should share the same Qwen3.5-family implementation when their native MTP structures are
    compatible.

    - GLM-5.2 would use a separate registered architecture because its DSA/MoE layer and model.layers.<N>.* checkpoint
    layout are different.

- modeling/target/checkpoint.py
    - Provides model-independent selective loading from local or Hugging Face sharded checkpoints.
    - It should not contain Qwen- or GLM-specific key knowledge.
    - Owns merging the trained MTP state back into the target checkpoint.
    - Architecture-specific draft modules provide the native state mapping, while the exporter owns checkpoint/index
    writing.

- training and runtime
    - Continue to use one MTPTrainStrategy, the existing FSDP backend, and the existing feature transport.
    - We should not add per-model trainers, consumers, or producers.

For this PR, Qwen3.5 can remain the first supported MTP draft architecture. However, we would like the common MTP objective/provider code to avoid Qwen-specific assumptions
so that the next architecture can be added through:

1. A new registered draft model under modeling/draft/mtp/.
2. A corresponding draft config and recipe.
3. Adding its class name to the MTP provider’s compatible architectures.
4. Native-weight and export parity tests.

We would also prefer native MTP initialization to be strict by default: if trainable native weights are missing, training
should fail instead of silently continuing from random initialization. An explicit random-init option can be added
separately if needed.

The key verification criteria for each architecture should be:

- Complete native-weight loading coverage.
- Correct common MTP shift/loss behavior.
- Round-trip native checkpoint mapping.
- Numerical parity with the corresponding serving implementation.
- Successful loading of the merged checkpoint by the pinned SGLang version.

[curnane-lab]

@curnane-lab We therefore do not plan to introduce separate strategies such as qwen_mtp or glm_mtp, nor a separate "family adapter" layer. And we want to rename algorithms/common/dflash_family_data.py to algorithms/common/hidden_states_data.py.

The intended structure is approximately:

specforge/
├── algorithms/
│   └── mtp/
│       └── providers.py
├── modeling/
│   ├── draft/
│   │   └── mtp/
│   │       ├── base.py
│   │       ├── qwen3_5.py
│   │       └── glm_moe_dsa.py
│   └── target/
│       ├── checkpoint.py
│       └── target_utils.py
├── core/
│   └── mtp.py
├── training/
│   └── strategies/
│       └── base.py
└── export/
    └── mtp.py

The responsibilities would be:

- algorithms/mtp/providers.py
    - Declares the MTP feature contract and compatible draft architectures.
    - Wires together draft construction, data providers, and MTPTrainStrategy.
    - Should remain model-family agnostic.
    - Owns the MTP reader, normalizer, and collator.
    - The common feature contract remains input_ids, loss_mask, and target_last_hidden_states.

- core/mtp.py
    - Owns the architecture-independent token/label shifting, CE loss, and metrics.
    - It should accept a generic MTP draft module and should not import Qwen3_5MTPDraftModel.

- modeling/draft/mtp/<architecture>.py
    - Owns the actual trainable MTP network and native weight layout for one architecture family.
    - Qwen3.5 and Qwen3.6 should share the same Qwen3.5-family implementation when their native MTP structures are
    compatible.

    - GLM-5.2 would use a separate registered architecture because its DSA/MoE layer and model.layers.<N>.* checkpoint
    layout are different.

- modeling/target/checkpoint.py
    - Provides model-independent selective loading from local or Hugging Face sharded checkpoints.
    - It should not contain Qwen- or GLM-specific key knowledge.
    - Owns merging the trained MTP state back into the target checkpoint.
    - Architecture-specific draft modules provide the native state mapping, while the exporter owns checkpoint/index
    writing.

- training and runtime
    - Continue to use one MTPTrainStrategy, the existing FSDP backend, and the existing feature transport.
    - We should not add per-model trainers, consumers, or producers.

For this PR, Qwen3.5 can remain the first supported MTP draft architecture. However, we would like the common MTP objective/provider code to avoid Qwen-specific assumptions
so that the next architecture can be added through:

1. A new registered draft model under modeling/draft/mtp/.
2. A corresponding draft config and recipe.
3. Adding its class name to the MTP provider’s compatible architectures.
4. Native-weight and export parity tests.

We would also prefer native MTP initialization to be strict by default: if trainable native weights are missing, training
should fail instead of silently continuing from random initialization. An explicit random-init option can be added
separately if needed.

The key verification criteria for each architecture should be:

- Complete native-weight loading coverage.
- Correct common MTP shift/loss behavior.
- Round-trip native checkpoint mapping.
- Numerical parity with the corresponding serving implementation.
- Successful loading of the merged checkpoint by the pinned SGLang version.

Hi @jiapingW, thanks for the review. The branch has been reworked accordingly.

Architecture now follows your sketch: one mtp algorithm and one MTPTrainStrategy, family differences resolve only through the draft registry (modeling/draft/mtp/base.py + qwen3_5.py, no adapter layer). dflash_family_data.py is renamed to hidden_states_data.py. algorithms/mtp/providers.py and core/mtp.py contain no Qwen-specific assumptions. Native init is strict: missing mtp.* weights fail construction instead of silently random-initing (warm start via model.draft_checkpoint_path still works).

One thing to confirm: the review assigns "owns merging" to checkpoint.py but "checkpoint/index writing" to the exporter. I implemented it as generic merge mechanism incl. writing in checkpoint.py (merge_state_into_checkpoint), MTP-specific policy (key normalization, embedding backfill, config patch) in export/mtp.py. Happy to move it if you prefer the other split.

CPU tests (tests/test_utils/test_mtp.py) cover full native-weight loading, shift/loss alignment, selective loading, and the merge round-trip. Serving parity and merged-checkpoint loading on the pinned SGLang version are being re-run on the NPU stack; training numbers are in the comment above.

[jiapingW]

Could you grant me permission to submit patches? I am making some relevant fixes, including RoPE alignment and model export with merged MTP weights. @curnane-lab

[curnane-lab]

Could you grant me permission to submit patches? I am making some relevant fixes, including RoPE alignment and model export with merged MTP weights. @curnane-lab

Done, you've been added as a collaborator on the fork. Feel free to push the RoPE alignment and export fixes directly to the PR branch. Thanks for picking these up!

[jiapingW]

I have updated it. You can see the last commit. If you think is OK, you can help use the qwen3.5 to run a simple training pipeline to test it. Then I'll merge it. Thanks. @curnane-lab