Skip to content

MTP#27

Open
tomiock wants to merge 4 commits into
mainfrom
mtp
Open

MTP#27
tomiock wants to merge 4 commits into
mainfrom
mtp

Conversation

@tomiock

@tomiock tomiock commented Jun 9, 2026

Copy link
Copy Markdown
Member

Implementation of Multi Token Prediction on Qwen3-VL and Qwen3.5

  • training loss converges early
  • compilation time is worse

Needs stronger verification via a proper training run

tomiock and others added 4 commits June 9, 2026 10:30
Depth-1 sequential MTP head (DeepSeek-V3 / Qwen3-Next style): a top-level
`mtp` module predicts token i+2 from the trunk hidden + embed(token i+1),
sharing embed_tokens and the tied lm_head. Param layout matches the released
Qwen3.5 checkpoints (mtp.fc / pre_fc_norm_{embedding,hidden} / layers / norm);
fc concat order is [embedding, hidden].

- Config-gated by text_config.mtp_num_hidden_layers (0 = disabled, fully
  backward-compatible; base Qwen3-VL checkpoints unaffected).
- Forward returns total = main_ce + mtp_loss_weight * mtp_ce; build_mtp_targets
  masks packed-segment boundaries so no supervision leaks across samples.
- Training wiring: train_mtp / random_init_mtp / lr_mtp / mtp_loss_weight,
  init_mtp, MTP optimizer group, FSDP/TP/compile, FLOPs accounting, warning_once
  on architecture-vs-config mismatch.
- eval/benchmark_mtp_decode.py: self-speculative greedy decode benchmark
  (verified ~1.61x decode speedup, ~0.77 acceptance, lossless on Qwen3.5-2B).
- models/tests/test_mtp.py: masking, forward/backward, checkpoint-load parity,
  and a draft-quality regression guard.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
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.

1 participant