[generate] Add proper MTP support by Cyrilvallez · Pull Request #46229 · huggingface/transformers

Cyrilvallez · 2026-05-27T05:19:22Z

What does this PR do?

As per the title! Finally add MTP decoding to Transformers!
For models supporting it and having proper trained weights, it can be used as simply as:

from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "zai-org/GLM-4.5-Air"

model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto")
tokenizer = AutoTokenizer.from_pretrained(model_id)

input_text = "Hello, who are you and what can you do?"
inputs = tokenizer(input_text, return_tensors="pt").to(model.device)

# Simply pass the kwarg `use_mtp=True` to use it!
out = model.generate(**inputs, use_mtp=True, do_sample=False, max_new_tokens=100)
output_text = tokenizer.batch_decode(out[:, inputs.input_ids.shape[1]:])[0]
print(output_text)

It will automatically scan the repo corresponding to the model calling generate, and load the weights corresponding to the MTP layers.

Note that it integrates nicely into the existing _assisted_decoding, thus avoiding to add yet another decoding method into generate.

Benchmark and expected perfs

On the previous example, I benchmarked the speed (excluding loading time of the MTP weights of course) and token matching, and obtained an acceptance rate of 73.68% for the drafted tokens from the (unique) MTP layer, resulting in a decoding speedup ratio of about 1.4, or 40% faster decoding compared to not using MTP.
This is of course a single test on a small toy example, so proper benchmarking may be needed, but it gives a good idea of what we can expect in general I believe.

Limitations

Note that this is unfortunately NOT supported for the recent Deepseek v4, as they changed the naming/conventions of the MTP layers (they mostly split the eh_proj layer in e_proj and h_proj, and use quantization on them). It should be relatively straightforward to fix it and make the weights match with simple ConversionOps, either on the general conversion mapping or on a model attribute such as _mtp_conversions or something.

HuggingFaceDocBuilderDev · 2026-05-27T05:31:53Z

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

vasqu

Some initial comments from my side

My main concern is that the whole design is inference centric and I'm not sure if we build technical debt when we want to add this to training in the future. Would love some discussion around this to make sure we are aligned.

Other than that, maybe some concerns about the usage of a cache (seems to be mandatory) and whether hybrid models exist with MTP (e.g. SWA and full attn) as they are not reflected here

vasqu · 2026-05-29T11:35:00Z

+class MTPCandidateGenerator(AssistedCandidateGenerator):
+    requires_model_outputs: bool = True
+    # We always need to pass the hidden states from the main model
+    model_kwargs_overrides: dict[str, Any] = {"output_hidden_states": True}


Is this the same as in gemma4 assistant? (See the note I added there :D)

If we only need the last hidden state, it would maybe make sense to fix the output recorder to only really extract the last one for example or simplify in some other way.

We indeed only need the last one here, but for simplicity it integrates with the usual output_hidden_states which exists for all models. We could definitely tweak the output recorder to allow output_last_hidden_states, but a bit out of scope of this PR I'd say

Yes for sure, sorry was meant as a general note not specific to this PR 🫡

vasqu · 2026-05-29T11:41:51Z

+        self.device = next(x.device for x in main_model.base_model.layers[-1].parameters())  # type: ignore
+        self.mtp_model = MtpLayerStack.from_pretrained(main_model, device_map={"": self.device})


Any reason we really need this heuristic or do you think a .to(device) call when generating the candidate would be to expensive?

Edit: seen we call to device either way. Don't really see any reason why we cant do mtp_model.device instead then with auto or pass the original kwargs from the original model here as well? E.g. The attention implementation should also be toggleable

IMO, it makes a lot of sense to put it on last gpu in general, as it's an additional layer of some sort that would usually be put on last device as well if loaded from the beginning - we can change though if you think otherwise

Ah no I think that's also a fair point. Haven't considered it that way

vasqu · 2026-05-29T11:44:07Z

+            if len(cache) != main_model.config.get_text_config().num_hidden_layers + self.num_mtp_layers:
+                cache.layers.extend([DynamicLayer() for _ in range(self.num_mtp_layers)])
+        else:
+            raise ValueError("No cache yet")


What speaks against no cache? Should be possible as well, no?

It would, but the _assisted_decoding decoding method forces a DynamicCache for now, no cache is not an option yet (due to rollback ops etc)

Argh, ok can weadd this comment at least so it's clear?

vasqu · 2026-05-29T12:12:56Z

+
+        self.post_init()
+
+    def forward(


maybe a generate vs forward split would be nicer. I'm still thinking of training tbh and how that flow into this wholse design

yeah I think drafting training is important as well!

The fact is that it's kind of the same: the usual forward is a generation loop, in the sense that later layers need both the last hidden states from previous layer AND the new drafted token -> so they cannot be separated from one another in general

ArthurZucker

Nice! its elegant! I'd be down to see how it scales with different architectures + training as @vasqu said!

For from pretrained would want to re-use or let it behave normally rather than not go through the general path!

ArthurZucker · 2026-05-29T13:16:26Z

+
+        self.post_init()
+
+    def forward(


yeah I think drafting training is important as well!

ArthurZucker · 2026-05-29T13:22:11Z

+            # Decode one token
+            next_token_logits = logits[:, -1, :].to(device=input_ids.device)
+            if logits_processor is not None and full_input_ids is not None:
+                next_token_scores = logits_processor(full_input_ids, next_token_logits.to(torch.float32))
+            if do_sample:
+                probs = nn.functional.softmax(next_token_scores, dim=-1, dtype=torch.float32)
+                next_mtp_token = torch.multinomial(probs, num_samples=1)
+            else:
+                next_mtp_token = torch.argmax(next_token_scores, dim=-1, keepdim=True)
+            drafted_tokens.append(next_mtp_token)


something we ought to reuse if possible from normal flow

ArthurZucker · 2026-05-29T13:28:34Z

        WeightRenaming("mlp.shared_expert.", "mlp.shared_experts."),
    ]

+    mapping["MtpLayerStack"] = [


this should be model familiy dependant non?

Unless we create a MTP layer for each model, we cannot really do it - or we need to alter a bit the logic of getting the conversion mapping for mtp

we can leave this for the case where the mtp layer is different!

…gface#46229 Wire up Multi-Token Prediction (MTP) training for Qwen3.5 by reusing the inference-side MtpLayer from huggingface#46229 instead of the inline MTP definition from huggingface#45638, and add the matching Trainer hook that combines the auxiliary loss. Model side (Qwen3.5): * CausalLMOutputWithPast gains an mtp_loss field carrying the unweighted MTP loss (averaged over MTP layers and non-ignored tokens). * Qwen3_5TextConfig exposes num_nextn_predict_layers (canonical MTP layer count, aligned with huggingface#46229) and output_mtp_loss (training-only switch). * Qwen3_5ForCausalLM dynamically instantiates MtpLayer copies that mirror the main decoder layer architecture (full / linear attention mix, RMSNorm variant) and recomputes the per-layer loss with properly sliced inputs, position embeddings, and labels. Trainer side: * TrainingArguments gains mtp_loss_coef (default 0.0, BC-safe). * Trainer.compute_loss adds the unweighted output.mtp_loss * mtp_loss_coef to the main loss at the very end, leaving the model forward signature clean and the combination policy centralized in the trainer. Supersedes huggingface#45638's inline-MTP approach in favor of a single shared MtpLayer implementation that the inference path already uses. Refs: huggingface#46229, huggingface#45638

curnane-lab · 2026-06-05T14:20:12Z

@Cyrilvallez — following your add-mtp branch closely. The training side for Qwen3.5 MTP doesn't exist yet in the repo.

A bit of context: I started prototyping this back in #45638 with an inline Qwen3_5MTP module and mtp_num_hidden_layers config — but that design created duplication with your inference-side MtpLayer. So I refactored it to align with #46229:

Uses from ...generation.mtp import MtpLayer (same module as inference)
Config field num_nextn_predict_layers matches your naming
Model returns raw mtp_loss; Trainer applies mtp_loss_coef weighting
~+500 LOC on top of [generate] Add proper MTP support #46229

Cyrilvallez · 2026-06-15T07:07:48Z

Hey @ArthurZucker @vasqu thanks for the reviews! I believe I answered all comments, let me know if something is unclear.
In general, I agree that this is mostly inference-centric, however the forward should be usable for (at least simple cases) of training as well - it would only require adding the proper additional loss, but the forward itself would not change. In general, they go a bit hand-in-hand, as later layers need to see both the last hidden states AND the new drafted token (which is basically inference), so the line is a bit blurry here.
More complicated training would require quite a bit of additional efforts though - for example packed sequence training. Since we need to draft tokens and update the position_ids in the main mtp loop, packing format is not so straightforward.
Detection of exact layer type and norms for the mtp layers is also not so clear all the time, as @vasqu noted in the case of hybrid archs. The only alternative I see here however would be to have per-model mtp implementations, living in the main model file. This would allow explicit classes and variations, but will lead to duplication. Another strategy could be to use a GenericForMTP or something, implementing the forward in general, and then inherit from this and simply add the proper __init__ to describe exact layers and norms in each model.
Let me know what you think!

github-actions · 2026-06-15T07:13:12Z

View the CircleCI Test Summary for this PR:

https://huggingface.co/spaces/transformers-community/circle-ci-viz?pr=46229&sha=862db4

vasqu · 2026-06-15T12:54:09Z

Awesome work overall!!

I think it would be nice to provide the basic training path for MTP then. Iiuc, it's just the "missing" labels in the forward. While they may be similar, I want to avoid that we encounter something unexpected. I don't think packed sequences are a big priority for now tbh so let's keep it simple for sure.

For the specific new patterns (e.g. hybrid linear attention). I think it makes sense to have the general implementation ("normal" attention) which we override to allow more specific patterns (hybrid linear attention). I would expect all hybrids to follow a similar pattern there so that we won't need the model specifics(?) and only variations of MTP if that makes sense

renames mtp_num_hidden_layers to num_mtp_layers (matches huggingface#46229) and removes mtp_loss_weight / output_mtp_loss. forward still returns mtp_loss separately so the trainer can apply its own weighting.

renames Qwen3_5MTPLayer/Qwen3_5MTP to Qwen3_5MtpLayer/Qwen3_5MtpLayerStack and adds the unexpected-key filter used by the generic MtpLayerStack. once huggingface#46229 lands these classes can be replaced by direct imports.

mirrors the entries huggingface#46229 adds for MtpLayerStack, so ckpts trained with this PR can be loaded back by the generic loader once huggingface#46229 lands.

github-actions · 2026-07-07T08:24:13Z

[For maintainers] Suggested jobs to run (before merge)

run-slow: deepseek_v3, glm4_moe, glm4v_moe, mimo_v2_flash, solar_open, youtu

zucchini-nlp · 2026-07-07T09:05:30Z

            "assistant_confidence_threshold": 0.4,
            "assistant_lookbehind": 10,
            "target_lookbehind": 10,
+            "use_mtp": False,


lets not add anything new here, as it was created for BC. The default value can be just None imo

github-actions · 2026-07-07T09:10:12Z

CI recap

Dashboard: View test results in Grafana
Latest run: 28852769836:1
Result: failure | Jobs: 15 | Tests: 171,469 | Failures: 6 | Duration: 17h 2m

ArthurZucker

great PR, its just missing some tests. TY for iterating

ArthurZucker · 2026-07-10T08:40:50Z

+        # just drafted from the main model. We need to slice to get only what the main model just processed, shifted by 1 to the
+        # right to take the new token as well. Say the main model just had token positions [2, 3] as input, the tensors
+        # contains the data for position [0, 1, 2, 3, 4], i.e. full inputs + new drafted token from last position 3 that was processed.
+        # We want to slice to get data for positions [3, 4] for the 1st mtp layer, i.e. same as main model, shifted by 1 to the right.


nice explanation!

ArthurZucker · 2026-07-10T08:42:06Z

        WeightRenaming("mlp.shared_expert.", "mlp.shared_experts."),
    ]

+    mapping["MtpLayerStack"] = [


we can leave this for the case where the mtp layer is different!

Cyrilvallez and others added 2 commits May 20, 2026 11:30

very raw start for later

9f34191

Merge branch 'main' into add-mtp

071e97f

Cyrilvallez and others added 18 commits May 27, 2026 18:57

good progress

17dccc5

close to it

ec80265

Merge branch 'main' into add-mtp

987b772

fix path

853c993

a few fixes

16b4399

continue

fc7932b

fix mask shapes

31ddc07

it works

6b12516

support any drafting method

6a0867c

improve

fdbb234

fixes

c0f8186

fix dtype

29f4174

fix dtype

594278e

add config

85dcd21

Merge branch 'main' into add-mtp

c1cd3aa

add ty ignores

52bca94

skip config check

874766f

avoid changing signatures - put it into generation_config instead

0b74127

Cyrilvallez changed the title ~~Add proper MTP support~~ [generate] Add proper MTP support May 29, 2026

raise error if invalid

3979c49

vasqu reviewed May 29, 2026

View reviewed changes

ArthurZucker reviewed May 29, 2026

View reviewed changes

ArthurZucker mentioned this pull request Jun 1, 2026

Add MTP speculative decoding via MTPCandidateGenerator #45618

Closed

This was referenced Jun 5, 2026

feat(qwen3.5+trainer): MTP training loss plumbing aligned with #46229 curnane-lab/transformers#1

Closed

feat(qwen3.5+trainer): MTP training loss plumbing aligned with #46229 #46433

Closed

Cyrilvallez and others added 6 commits June 15, 2026 14:35

Merge branch 'main' into add-mtp

a29bfd7

change name

5f44288

name changes

8d32ef5

remove pos_ids for mask

f301eee

cache is not strictly needed

c81cac1

fix

862db40

molbap mentioned this pull request Jun 19, 2026

Add LocateAnything #46776

Open

curnane-lab mentioned this pull request Jun 23, 2026

Add MTP training support for Qwen3.5 #45638

Open

Cyrilvallez and others added 6 commits July 7, 2026 14:09

merge main and fix conflioct

b98c9b5

create on meta

60922f8

add loss

8227d2e

move to modeling_layers

24324d9

fix modular

e0e1807

Merge branch 'main' into add-mtp

c49b376

Cyrilvallez added 2 commits July 7, 2026 17:31

fix

03aeb0e

fix

86545d9

zucchini-nlp reviewed Jul 7, 2026

View reviewed changes

curnane-lab mentioned this pull request Jul 8, 2026

[Feature] mtp support sgl-project/SpecForge#662

Open

2 tasks

ArthurZucker approved these changes Jul 10, 2026

View reviewed changes

		self.device = next(x.device for x in main_model.base_model.layers[-1].parameters()) # type: ignore
		self.mtp_model = MtpLayerStack.from_pretrained(main_model, device_map={"": self.device})

Uh oh!

Conversation

Cyrilvallez commented May 27, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Benchmark and expected perfs

Limitations

Uh oh!

HuggingFaceDocBuilderDev commented May 27, 2026

Uh oh!

vasqu left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

ArthurZucker left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

curnane-lab commented Jun 5, 2026

Uh oh!

Cyrilvallez commented Jun 15, 2026

Uh oh!

github-actions Bot commented Jun 15, 2026

Uh oh!

vasqu commented Jun 15, 2026

Uh oh!

github-actions Bot commented Jul 7, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

github-actions Bot commented Jul 7, 2026

CI recap

Uh oh!

ArthurZucker left a comment

Choose a reason for hiding this comment

Cyrilvallez commented May 27, 2026 •

edited by github-actions Bot

Loading