[Bugfix][MiniMax M3] Share target embeddings with MTP under PP (#58648)
Signed-off-by: Peiyuan Zhou <peiyuanzhou1994@gmail.com> Co-authored-by: Codex <noreply@openai.com> Co-authored-by: Woosuk Kwon <woosuk@inferact.ai> Co-authored-by: Jee Jee Li <pandaleefree@gmail.com>
xiaozhoupy committed
Sep 30, 2026 at 11:04 UTC
a8e069f81b4842723a04a84fad2cd5f757820af3
6 files changed
+229
-6
tests/v1/worker/test_spec_decode_embed_sharing_pp.py
+188
@@ -1,15 +1,20 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import importlib
4
from types import SimpleNamespace
5
6
import pytest
7
import torch
8
import torch.nn as nn
9
10
+from vllm.model_executor.model_loader.mtp_validation import (
11
+ disable_mtp_completeness_check,
12
+)
13
from vllm.model_executor.models.utils import (
14
PPMissingLayer,
15
spec_decode_needs_target_embed,
16
)
17
+from vllm.platforms import current_platform
18
from vllm.v1.worker.gpu.spec_decode.eagle import utils as eagle_utils
19
20
VOCAB, HIDDEN = 32, 8
@@ -108,3 +113,186 @@ def test_target_embedding_provisioning(
113
speculative_config = None if method is None else SimpleNamespace(method=method)
114
config = SimpleNamespace(speculative_config=speculative_config)
115
assert spec_decode_needs_target_embed(config) is expected
116
+
117
+
118
+@pytest.mark.parametrize(
119
+ "pp_size,is_last_rank,expected",
120
+ [(2, True, True), (2, False, False), (1, True, False)],
121
+)
122
+def test_mtp_target_embedding_requires_model_opt_in(
123
+ monkeypatch, pp_size, is_last_rank, expected
124
+):
125
+ monkeypatch.setattr(
126
+ "vllm.distributed.parallel_state.get_pp_group",
127
+ _fake_pp(pp_size, is_last_rank),
128
+ )
129
+ config = SimpleNamespace(speculative_config=SimpleNamespace(method="mtp"))
130
+ assert spec_decode_needs_target_embed(config, include_mtp=True) is expected
131
+ assert not spec_decode_needs_target_embed(config)
132
+
133
+
134
+@pytest.fixture
135
+def m3_modules():
136
+ if not current_platform.is_cuda_alike():
137
+ pytest.skip("M3 model imports require CUDA or ROCm dependencies")
138
+ backend = "amd" if current_platform.is_rocm() else "nvidia"
139
+ root = f"vllm.models.minimax_m3.{backend}"
140
+ return importlib.import_module(f"{root}.mtp"), importlib.import_module(
141
+ f"{root}.model"
142
+ )
143
+
144
+
145
+@pytest.fixture
146
+def m3_config():
147
+ hf_config = SimpleNamespace(
148
+ vocab_size=VOCAB,
149
+ hidden_size=HIDDEN,
150
+ num_local_experts=0,
151
+ num_hidden_layers=0,
152
+ rms_norm_eps=1e-5,
153
+ )
154
+ return SimpleNamespace(
155
+ use_v2_model_runner=True,
156
+ load_config=SimpleNamespace(load_format="auto"),
157
+ model_config=SimpleNamespace(hf_text_config=hf_config),
158
+ speculative_config=SimpleNamespace(
159
+ method="mtp", draft_model_config=SimpleNamespace(hf_config=hf_config)
160
+ ),
161
+ quant_config=None,
162
+ cache_config=None,
163
+ )
164
+
165
+
166
+@pytest.fixture
167
+def make_m3_draft(monkeypatch, m3_modules, m3_config):
168
+ mtp, _ = m3_modules
169
+ inner = _inner(_embed())
170
+ inner.layers = nn.ModuleDict({"0": nn.Linear(1, 1, bias=False)})
171
+ inner.num_mtp_layers = 1
172
+ monkeypatch.setattr(mtp, "MiniMaxM3MultiTokenPredictor", lambda **_: inner)
173
+ monkeypatch.setattr(mtp, "ParallelLMHead", lambda *_, **kw: _embed())
174
+ monkeypatch.setattr(mtp, "LogitsProcessor", lambda *_: None)
175
+ monkeypatch.setattr(
176
+ mtp, "fused_moe_make_expert_params_mapping", lambda *_, **kw: []
177
+ )
178
+ monkeypatch.setattr(mtp, "get_pp_group", _fake_pp(2))
179
+ monkeypatch.setattr(eagle_utils, "get_pp_group", _fake_pp(2))
180
+ return lambda: mtp.MiniMaxM3MTP(vllm_config=m3_config)
181
+
182
+
183
+@pytest.mark.parametrize(
184
+ "key",
185
+ [None, "model.embed_tokens.weight", "language_model.model.embed_tokens.weight"],
186
+)
187
+def test_m3_checkpoint_embedding_ownership(make_m3_draft, key):
188
+ """Actual M3 loading determines sharing and survives an unrelated reload."""
189
+ draft = make_m3_draft()
190
+ assert draft.has_own_embed_tokens
191
+ weights = [("model.mtp.layers.0.weight", torch.ones(1, 1))]
192
+ if key is not None:
193
+ weights.append((key, torch.ones(VOCAB, HIDDEN)))
194
+ loaded = draft.load_weights(weights)
195
+ assert ("model.embed_tokens.weight" in loaded) is (key is not None)
196
+ with disable_mtp_completeness_check():
197
+ draft.load_weights([("lm_head.weight", torch.zeros(VOCAB, HIDDEN))])
198
+ assert draft.has_own_embed_tokens is (key is not None)
199
+ target = _inner(_embed(fill=2.0))
200
+ original = draft.model.embed_tokens
201
+ eagle_utils.maybe_share_target_embed(draft, draft.model, target)
202
+ if key is None:
203
+ assert draft.model.embed_tokens is target.embed_tokens
204
+ else:
205
+ assert draft.model.embed_tokens is original
206
+ torch.testing.assert_close(original.weight, torch.ones(VOCAB, HIDDEN))
207
+
208
+
209
+@pytest.mark.parametrize(
210
+ "load_format", ["dummy", "sharded_state", "runai_streamer_sharded", "tensorizer"]
211
+)
212
+def test_m3_direct_loader_embedding_ownership(make_m3_draft, m3_config, load_format):
213
+ """Directly loaded embeddings survive transfers and share only when equal."""
214
+ m3_config.load_config.load_format = load_format
215
+ draft = make_m3_draft()
216
+ draft.load_state_dict(
217
+ {key: torch.ones_like(value) for key, value in draft.state_dict().items()}
218
+ )
219
+ with disable_mtp_completeness_check():
220
+ draft.load_weights([("model.mtp.layers.0.weight", torch.zeros(1, 1))])
221
+ target = _inner(_embed(fill=2.0))
222
+ original = draft.model.embed_tokens
223
+ eagle_utils.maybe_share_target_embed(draft, draft.model, target)
224
+ assert draft.model.embed_tokens is (
225
+ target.embed_tokens if load_format == "dummy" else original
226
+ )
227
+ if load_format != "dummy":
228
+ equal_target = _inner(_embed(fill=1.0))
229
+ eagle_utils.maybe_share_target_embed(draft, draft.model, equal_target)
230
+ assert draft.model.embed_tokens is equal_target.embed_tokens
231
+
232
+
233
+@pytest.mark.parametrize("use_v2,pp_size", [(True, 1), (False, 1), (False, 2)])
234
+def test_m3_preserves_legacy_embedding_sharing(
235
+ monkeypatch, m3_modules, m3_config, make_m3_draft, use_v2, pp_size
236
+):
237
+ mtp, _ = m3_modules
238
+ m3_config.use_v2_model_runner = use_v2
239
+ monkeypatch.setattr(mtp, "get_pp_group", _fake_pp(pp_size))
240
+ draft = make_m3_draft()
241
+ draft.load_weights(
242
+ [
243
+ ("model.mtp.layers.0.weight", torch.ones(1, 1)),
244
+ ("model.embed_tokens.weight", torch.ones(VOCAB, HIDDEN)),
245
+ ]
246
+ )
247
+ assert not hasattr(draft, "has_own_embed_tokens")
248
+
249
+
250
+@pytest.mark.parametrize(
251
+ "rank,use_v2,method,expected",
252
+ [
253
+ (0, True, "mtp", True),
254
+ (1, True, "mtp", False),
255
+ (2, True, "mtp", True),
256
+ (2, False, "mtp", False),
257
+ (2, True, None, False),
258
+ ],
259
+)
260
+def test_m3_target_allocates_embedding_on_required_stage(
261
+ monkeypatch, m3_modules, m3_config, rank, use_v2, method, expected
262
+):
263
+ _, model = m3_modules
264
+ pp = lambda: SimpleNamespace(
265
+ world_size=3, is_first_rank=rank == 0, is_last_rank=rank == 2
266
+ )
267
+ monkeypatch.setattr(model, "get_pp_group", pp)
268
+ monkeypatch.setattr("vllm.distributed.parallel_state.get_pp_group", pp)
269
+ monkeypatch.setattr(model, "VocabParallelEmbedding", lambda *_, **kw: _embed())
270
+ monkeypatch.setattr(model, "make_layers", lambda *_, **kw: (0, 0, nn.ModuleList()))
271
+ monkeypatch.setattr(model, "MiniMAXGemmaRMSNorm", lambda *_, **kw: nn.Identity())
272
+ m3_config.use_v2_model_runner = use_v2
273
+ if method is None:
274
+ m3_config.speculative_config = None
275
+ target = model.MiniMaxM3Model(vllm_config=m3_config)
276
+ assert isinstance(target.embed_tokens, nn.Embedding) is expected
277
+
278
+
279
+@pytest.mark.parametrize("initial_embed", [False, True])
280
+def test_m3_weight_transfer_updates_embedding_ownership(make_m3_draft, initial_embed):
281
+ """Layer transfers retain ownership; newly loaded embeddings become owned."""
282
+ draft = make_m3_draft()
283
+ weights = [("model.mtp.layers.0.weight", torch.ones(1, 1))]
284
+ if initial_embed:
285
+ weights.append(("model.embed_tokens.weight", torch.ones(VOCAB, HIDDEN)))
286
+ draft.load_weights(weights)
287
+ with disable_mtp_completeness_check():
288
+ draft.load_weights([("model.mtp.layers.0.weight", torch.zeros(1, 1))])
289
+ assert draft.has_own_embed_tokens is initial_embed
290
+ draft.load_weights(
291
+ [("model.embed_tokens.weight", torch.full((VOCAB, HIDDEN), 3.0))]
292
+ )
293
+ assert draft.has_own_embed_tokens
294
+ target = _inner(_embed(fill=2.0))
295
+ original = draft.model.embed_tokens
296
+ eagle_utils.maybe_share_target_embed(draft, draft.model, target)
297
+ assert draft.model.embed_tokens is original
298
+ torch.testing.assert_close(original.weight, torch.full((VOCAB, HIDDEN), 3.0))
vllm/model_executor/models/utils.py
+11
-4
@@ -844,17 +844,24 @@ class PPMissingLayer(torch.nn.Identity):
844
return args[0] if args else next(iter(kwargs.values()))
845
846
847
-def spec_decode_needs_target_embed(vllm_config: VllmConfig) -> bool:
848
- """Whether the last PP rank needs the target input embedding."""
847
+def spec_decode_needs_target_embed(
848
+ vllm_config: VllmConfig, *, include_mtp: bool = False
849
+) -> bool:
850
+ """Whether the last PP rank needs the target input embedding.
851
+
852
+ MTP targets opt in when their drafter supports sharing this embedding.
853
+ """
854
from vllm.distributed.parallel_state import get_pp_group
855
856
speculative_config = vllm_config.speculative_config
852
- if speculative_config is None or speculative_config.method not in (
857
+ if speculative_config is None:
858
+ return False
859
+ if speculative_config.method not in (
860
"eagle",
861
"eagle3",
862
"dflash",
863
"dspark",
857
- ):
864
+ ) and not (include_mtp and speculative_config.method == "mtp"):
865
return False
866
pp = get_pp_group()
867
return pp.world_size > 1 and pp.is_last_rank
vllm/models/minimax_m3/amd/model.py
+4
-1
@@ -85,6 +85,7 @@ from vllm.model_executor.models.utils import (
85
make_empty_intermediate_tensors_factory,
86
make_layers,
87
maybe_prefix,
88
+ spec_decode_needs_target_embed,
89
)
90
from vllm.model_executor.models.vision import run_dp_sharded_mrope_vision_model
91
from vllm.models.minimax_m3.amd.indexer_aiter import (
@@ -1370,7 +1371,9 @@ class MiniMaxM3Model(nn.Module, EagleModelMixin):
1371
1372
self.vocab_size = config.vocab_size
1373
1373
- if get_pp_group().is_first_rank:
1374
+ if get_pp_group().is_first_rank or spec_decode_needs_target_embed(
1375
+ vllm_config, include_mtp=vllm_config.use_v2_model_runner
1376
+ ):
1377
self.embed_tokens = VocabParallelEmbedding(
1378
config.vocab_size,
1379
config.hidden_size,
vllm/models/minimax_m3/amd/mtp.py
+11
@@ -24,6 +24,7 @@ import torch
24
import torch.nn as nn
25
26
from vllm.config import VllmConfig
27
+from vllm.distributed import get_pp_group
28
from vllm.model_executor.layers.fused_moe import (
29
fused_moe_make_expert_params_mapping,
30
)
@@ -164,6 +165,10 @@ class MiniMaxM3MTP(nn.Module):
165
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
166
super().__init__()
167
168
+ if vllm_config.use_v2_model_runner and get_pp_group().world_size > 1:
169
+ # Preserve embeddings from loaders that bypass load_weights.
170
+ self.has_own_embed_tokens = vllm_config.load_config.load_format != "dummy"
171
+
172
assert vllm_config.speculative_config is not None
173
self.config = vllm_config.speculative_config.draft_model_config.hf_config
174
self.quant_config = vllm_config.quant_config
@@ -333,4 +338,10 @@ class MiniMaxM3MTP(nn.Module):
338
f"Failed to load MTP layer {layer_idx} weights from checkpoint."
339
)
340
341
+ if hasattr(self, "has_own_embed_tokens"):
342
+ if "model.embed_tokens.weight" in loaded_params:
343
+ self.has_own_embed_tokens = True
344
+ elif is_mtp_completeness_check_enabled():
345
+ # A complete checkpoint without embeddings borrows the target's.
346
+ self.has_own_embed_tokens = False
347
return loaded_params
vllm/models/minimax_m3/nvidia/model.py
+4
-1
@@ -68,6 +68,7 @@ from vllm.model_executor.models.utils import (
68
make_empty_intermediate_tensors_factory,
69
make_layers,
70
maybe_prefix,
71
+ spec_decode_needs_target_embed,
72
)
73
from vllm.model_executor.models.vision import run_dp_sharded_mrope_vision_model
74
from vllm.models.minimax_m3.common.encoder_cudagraph import (
@@ -803,7 +804,9 @@ class MiniMaxM3Model(nn.Module, EagleModelMixin):
804
805
self.vocab_size = config.vocab_size
806
806
- if get_pp_group().is_first_rank:
807
+ if get_pp_group().is_first_rank or spec_decode_needs_target_embed(
808
+ vllm_config, include_mtp=vllm_config.use_v2_model_runner
809
+ ):
810
self.embed_tokens = VocabParallelEmbedding(
811
config.vocab_size,
812
config.hidden_size,
vllm/models/minimax_m3/nvidia/mtp.py
+11
@@ -8,6 +8,7 @@ import torch
8
import torch.nn as nn
9
10
from vllm.config import VllmConfig
11
+from vllm.distributed import get_pp_group
12
from vllm.model_executor.layers.fused_moe import (
13
fused_moe_make_expert_params_mapping,
14
)
@@ -146,6 +147,10 @@ class MiniMaxM3MTP(nn.Module):
147
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
148
super().__init__()
149
150
+ if vllm_config.use_v2_model_runner and get_pp_group().world_size > 1:
151
+ # Preserve embeddings from loaders that bypass load_weights.
152
+ self.has_own_embed_tokens = vllm_config.load_config.load_format != "dummy"
153
+
154
assert vllm_config.speculative_config is not None
155
self.config = vllm_config.speculative_config.draft_model_config.hf_config
156
self.quant_config = vllm_config.quant_config
@@ -315,4 +320,10 @@ class MiniMaxM3MTP(nn.Module):
320
f"Failed to load MTP layer {layer_idx} weights from checkpoint."
321
)
322
323
+ if hasattr(self, "has_own_embed_tokens"):
324
+ if "model.embed_tokens.weight" in loaded_params:
325
+ self.has_own_embed_tokens = True
326
+ elif is_mtp_completeness_check_enabled():
327
+ # A complete checkpoint without embeddings borrows the target's.
328
+ self.has_own_embed_tokens = False
329
return loaded_params