feat: add mtp support#667
Open
curnane-lab wants to merge 28 commits into
Open
Conversation
- Add Qwen3_5MTPDraftModel (specforge/modeling/draft/mtp.py) - fc fusion of normalized input embeddings + target last hidden states - 1-layer Qwen3 transformer + shared lm_head - Weight layout matches SGLang Qwen3_5ForCausalLMMTP (mtp.* prefix) - Add OnlineMTPModel (specforge/core/mtp.py) - Performs next-token shift internally on raw input_ids - Returns per-layer raw loss and per-position acc_corrects/acc_denoms - Add generate_mtp_data to Eagle3TargetModel backends - SGLang: extend(return_last_hidden_states=True, aux=False, logits=False) - HF/Custom: output_hidden_states=True and take hidden_states[-1] - Returns raw input_ids (no pre-shift) for DFlash-style MTP semantics - Add scripts/train_mtp.py - Reuses build_eagle3_dataset and FSDP training skeleton - Shares/freezes target embed_tokens and lm_head with draft model - Saves checkpoints with mtp.* weight keys for SGLang inference - Add configs/qwen3.5-4b-mtp.json for Qwen3.5-4B MTP training - Register new classes in specforge.modeling and specforge.core
P0-1: Nest the 1-layer transformer under mtp.model (new Qwen3MTPInnerModel) so the saved keys become mtp.model.layers.0.* / mtp.model.norm.weight. Previously the draft produced mtp.layers.0.* / mtp.norm.weight, which SGLang's load_weights (mtp.->model. remap) maps to model.layers.0.* but SGLang's Qwen3_5ForCausalLMMTP expects model.model.layers.0.* (its inner Qwen3_5ForCausalLM.model.layers) - so the transformer + final norm would silently fail to load at inference. P0-2: Set config architectures to Qwen3_5ForConditionalGeneration so SGLang's is_draft_model switching maps it to Qwen3_5ForCausalLMMTP. The previous value Qwen3_5MTPDraftModel is a SpecForge training class unrecognized by SGLang.
Drop @AbstractMethod from Eagle3TargetModel.generate_mtp_data and give it a default body that raises NotImplementedError. generate_eagle3_data stays abstract (core capability), while MTP is now an optional capability: external Eagle3TargetModel subclasses remain constructible without overriding generate_mtp_data, and only fail at call time if MTP is actually used. The three in-tree backends (HF/SGLang/Custom) still override it, so existing MTP behavior is unchanged.
Add examples/run_qwen3.5_4b_mtp_online_npu.sh, mirroring the DFlash NPU script layout but using scripts/train_mtp.py and configs/qwen3.5-4b-mtp.json. Key settings: - target-model-backend: hf (same as DFlash NPU script) - chat-template: qwen3.5 - attention-backend: sdpa (default), overridable via - 8 GPUs default, overridable via - batch-size 2 + accumulation-steps 4
- Move ATTENTION_BACKEND_CHOICES import into SGLangBackendArgs.add_args so that specforge/args.py no longer imports sglang at module load time. - In train_mtp.py, only add SGLang CLI args when --target-model-backend=sglang is selected, allowing NPU runs with --target-model-backend=hf and --attention-backend=sdpa to run without sglang installed.
Make init_distributed detect the accelerator via get_device_type(): - CUDA -> backend='nccl', torch.cuda.set_device - NPU -> backend='hccl', torch.npu.set_device - Replace hard-coded 'cuda' device_type in device meshes with the detected type. This lets NPU training scripts run without manually changing the backend.
Replace hard-coded device='cuda' in get_eagle3_target_model with the auto-detected device_type (cuda/npu). This allows the HF backend to load the target model on NPU without hitting 'Torch not compiled with CUDA enabled'.
- HFEagle3TargetModel.generate_mtp_data now captures only the last layer output via a forward hook instead of materializing the full hidden_states tuple for all layers. - Forward through self.model.model (the transformer backbone) instead of self.model to skip the final lm_head, saving logits memory/compute. - Put target_model in eval mode in train_mtp.py since it is frozen and only used for hidden-state generation. These changes significantly lower per-rank peak memory on NPU/GPU.
Single-rank HF backend now loads the model on CPU (with low_cpu_mem_usage) and then explicitly moves it to the target device, matching the DFlash HF backend. This avoids transformers' device_map / caching_allocator_warmup path on NPU, which can OOM when the PyTorch NPU allocator pool is small even though total device memory is large.
HFEagle3TargetModel inherits from ABC, not nn.Module, so it has no eval(). Eval the inner self.model instead (with a safe fallback for other backends).
TargetEmbeddingsAndHead now falls back through common embedding and lm_head key names when the default 'model.embed_tokens.weight' / 'lm_head.weight' are not present. This supports plain LLMs, MLLMs, and renamed checkpoints without requiring manual --embedding-key / --lm-head-key arguments.
Include 'model.language_model.embed_tokens.weight' and 'language_model.embed_tokens.weight' in the auto-detection fallback list. These variants are used by Qwen3.5-A3B and similar MLLMs (see Domino NPU example's --embedding-key).
_get_transformer_layers() only handled causal LMs (model.model.layers / model.layers / model.transformer.h), so HFEagle3TargetModel.generate_mtp_data crashed with "Could not locate transformer layers" on multimodal targets like Qwen3_5ForConditionalGeneration, whose text decoder is nested under model.language_model (embedding key resolves to model.language_model.embed_tokens.weight). Add _get_language_model() returning the inner transformer (without lm_head) for both causal LMs (model.model) and multimodal models (model.language_model or .language_model.model), and route both _get_transformer_layers and the generate_mtp_data forward call through it. The embedding-key resolution was already VLM-aware; this makes layer lookup and the cheap forward VLM-aware too. Also the generate_mtp_data forward used self.model.model(...) which is wrong for VLMs; now uses _get_language_model() which picks model.language_model for VLMs. Non-VLM behavior unchanged (Qwen3ForCausalLM still resolves to model.model). Also benefits generate_eagle3_data layer lookup on VLM HF backend.
My previous fix (8f6f1dc) assumed the text decoder lives at model.language_model, but the crash persisted on the actual Qwen3.5 checkpoint, meaning the real nesting differs. Make _get_language_model fall back to a recursive search: scan the module tree for any module whose direct child is a `layers`/`h` nn.ModuleList, preferring a `language_model` path and the largest layer count. This copes with arbitrary nesting (e.g. model.model.language_model, or a bare text CausalLM loaded by AutoModelForCausalLM). Also make the _get_transformer_layers error diagnostic: it now prints the model class and top-level children, so any future miss reports the actual structure instead of failing opaquely. generate_mtp_data forward call is unchanged (already routes through _get_language_model, which now resolves the correct text transformer).
…into base model The new script takes a trained MTP checkpoint (from train_mtp.py) and merges the MTP weights back into the original Qwen3.5 base model checkpoint. This produces a single self-contained model directory that can be served directly. Supports two output key layouts: - sglang: keeps mtp.model.layers.* (SGLang internal remap) - hf: converts to mtp.layers.* (native HuggingFace / vLLM layout)
When the base Qwen3.5 checkpoint already contains native MTP weights, the merge script now removes the old MTP entries before adding the trained MTP weights. This avoids duplicate keys and correctly updates the model rather than appending extra MTP tensors.
…vert Qwen3MTPInnerModel nesting; SGLang ForCausalLM is flat not nested)
The draft config and attention were copied from Qwen3 (head_dim=128, num_heads=32, no output gate, full rotary), which is incompatible with the official Qwen3.5-4B architecture and produces weights that cannot be loaded by vLLM/SGLang. Changes: - configs/qwen3.5-4b-mtp.json: update all architecture values to match official Qwen3.5-4B text_config (head_dim=256, num_attention_heads=16, num_key_value_heads=4, intermediate_size=9216, vocab_size=248320, attn_output_gate=true, partial_rotary_factor=0.25, rope_theta=1e7) - mtp.py: add attn_output_gate support (q_proj outputs 2x for q+gate, sigmoid gate applied before o_proj) - mtp.py: add PartialRotaryEmbedding (inv_freq for only int(head_dim * partial_rotary_factor) dimensions) - mtp.py: update apply_rotary_pos_emb to split q/k into rotary and non-rotary parts, matching official Qwen3_5 modeling code
…Forge into add_mtp_support
…rom Eagle3TargetModel
curnane-lab
requested review from
FlamingoPg,
FrankLeeeee,
shuaills and
sleepcoo
as code owners
Contributor
|
Warning You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again! |
curnane-lab
force-pushed
the
add_mtp_support
branch
from
67c54ac to
c849d4b
Compare
curnane-lab
force-pushed
the
add_mtp_support
branch
from
644adcf to
f965680
Compare
curnane-lab
force-pushed
the
add_mtp_support
branch
from
f965680 to
3580c7a
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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.pymodules, 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 flatmtp.layers.*layout required by SGLang/vLLM.specforge/core/mtp.py:OnlineMTPModeltraining 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 fromEagle3TargetModel.specforge/modeling/target/eagle3_target_model.py: minimal changes — CPU-first HF target loading for NPU memory, andcapture_aux_hidden_statesflag for SGLangextend().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
Accuracy Test
N/A — training utility and model wrapper; no change to existing inference model math.
Benchmark & Profiling
--target-model-backend hf.Checklist