vi Feature/mtp - #10683
Open
HelloWorldBeginner wants to merge 4 commits into
Open
Conversation
HelloWorldBeginner
force-pushed
the
feature/mtp
branch
from
July 25, 2026 03:15
9e4d434 to
5dcb6d7
Compare
Adapt the FSDP2 MTP implementation from MindSpeed-LLM (mindspeed_llm/fsdp2/models/common/mtp.py) to the LlamaFactory v1 architecture, making it HuggingFace-transformers generic so it can be attached to any Llama/Qwen3/Mistral-style decoder-only causal LM. Components: * MultiTokenPredictionBlock / MultiTokenPredictionLayer (src/llamafactory/v1/plugins/model_plugins/mtp.py): K MTP heads that reuse the base model's decoder layer class, with shared enorm/hnorm norms, e_proj/h_proj projections and a final_layernorm, mirroring MindSpeed-LLM. Head k predicts token p+k+2 from position p. * MTPModelPlugin: grafts the block onto the model as `model.mtp` and patches `model.forward` to emit `mtp_logits` during training. * compute_mtp_loss: loss_weights-weighted per-head CE mean, averaged over heads; total loss = lm_loss + loss_scale * mtp_loss (loss_scale=0.3). * ModelArguments.mtp_config plugin field + ModelEngine wiring. * SFTTrainer MTP-aware compute_loss (non context-parallel path). * FSDP2: MTP inner decoder layers are sharded automatically by the existing prepare_model loop (wraps all modules of the decoder-layer class, including mtp.layers.*.layer). * Example YAML, unit tests, and docs. Verified on llamafactory/tiny-random-qwen2.5 (attach, forward shapes, loss + backward, head-offset convention).
Extend the MTP layer to work under Ulysses context parallelism, following the MTP+CP loss reduction in MindSpeed-LLM's FSDP2 trainer. Changes: * compute_mtp_loss (mtp.py): add a `cp_group` argument and a context-parallel per-head loss path (`_cp_head_loss`). Under CP, each head all-gathers labels / loss_weights / log_probs across the CP group so the per-head loss is computed on the full sequence, mirroring the single-head `sequence_parallel_loss` plugin. Non-CP behavior is unchanged. * sequence_parallel.py: refactor the main-head CP loss into a reusable helper and add a `sequence_parallel_mtp_loss` plugin that computes the main CP loss plus the scaled MTP loss. The MTP decoder layers participate in Ulysses attention automatically through the existing global `_flash_attention_forward` patch. * base_trainer.py: route to `sequence_parallel_mtp_loss` in `fit()` when CP is enabled and the model has an MTP block. * Example YAML (train_full_mtp_ulysses_cp.yaml), CP alignment unit test, and docs. The CP alignment test (2 gloo ranks) verifies that the context-parallel MTP loss reproduces the full-sequence MTP loss exactly. CP requires FSDP2 and flash_attention_2 (not DeepSpeed), same as non-MTP CP.
Upstream hiyouga#10598 replaced the v1 'template' field with apply_chat_template (+ custom_chat_template). Keeping 'template: qwen3_nothink' makes HfArgumentParser reject the config with: ValueError: Some keys are not used by the HfArgumentParser: ['template'] Qwen3 models apply their tokenizer chat template automatically, so the field is simply removed from both MTP example configs.
… mtp_loss logging Address four issues in the v1 MTP (Multi-Token Prediction) implementation: * Layer selection (Q1): pick a full_attention layer index for each MTP head instead of blindly taking the last layer. Hybrid-attention models (Qwen3 mixes full/sliding; Qwen3.5 mixes full/GDN) select the attention type by config.layer_types[layer_idx], and an MTP head needs full self-attention (global context) to predict the token at offset k+2. _select_layer_idx_for_mtp picks the last full_attention idx (fallback to N-1 for all-full Llama/Mistral), so the MTP head always builds a self_attn module that goes through _flash_attention_forward (and is thus covered by the Ulysses CP patch). * Weight save/load (Q2): strip the shared mtp.embed_tokens/output_layer keys before save_pretrained (4 save paths: fsdp2.save_model, fsdp2.save_checkpoint save_ckpt_as_hf, base_trainer non-dist save_model, checkpoint standard save) to avoid the shared-tensors RuntimeError; load_mtp_weights re-reads mtp.* from the checkpoint after apply_mtp (from_pretrained drops them as unexpected). No-op on meta device (FSDP2 meta path loads mtp.* via the HF weight loop). Verified end-to-end: save no longer raises, weights restored bit-identical, shared modules re-tied. * CP forward shift (Q3): shift_input_ids_for_mtp is context-parallel aware — under CP it all-gathers each rank's first token and fills the previous rank's tail with the next rank's first token (only the global last rank is padded), instead of a plain local roll that drops the real next token at every CP boundary. Attention (Ulysses global patch) and loss (_cp_head_loss all-gather) were already correct. * mtp_loss logging: expose the unscaled per-head-mean MTP loss as model._last_mtp_loss in both the non-CP (sft_trainer) and CP (sequence_parallel_mtp_loss) paths; BaseTrainer logs it as `mtp_loss` alongside `loss` so MTP convergence is visible during training. Tests: add test_mtp_save_load (save->reload weight restoration) and test_mtp_shift_input_ids_cp (2-rank gloo shift vs full-sequence roll). All 8 test_mtp.py tests pass; ruff clean. Docs/yaml: add docs/v1_mtp.md updates, docs/v1_mtp_testing.md (end-to-end test flow), scripts/verify_mtp_save_load_e2e.py, and three test yamls (save_load, resume, cp) using Qwen3-0.6B (non-GDN).
HelloWorldBeginner
force-pushed
the
feature/mtp
branch
from
August 11, 2026 09:07
0765a87 to
ff2cd43
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
1 participant
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.
What does this PR do?
Fixes # (issue)
Before submitting