[Bugfix][SpecDecode] Scope MTP completeness checks outside bucketed updates (#49178)
Signed-off-by: aoshen02 <aoshen02@users.noreply.github.com> Co-authored-by: aoshen02 <aoshen02@users.noreply.github.com>
aoshen02 committed
Jul 22, 2026 at 09:56 UTC
ba1892907940c6baa50f4b385ae4f21de5452c4d
14 files changed
+131
-37
tests/model_executor/model_loader/test_mtp_validation.py
new
+19
@@ -0,0 +1,19 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+
6
+from vllm.model_executor.model_loader.mtp_validation import (
7
+ disable_mtp_completeness_check,
8
+ is_mtp_completeness_check_enabled,
9
+)
10
+
11
+
12
+def test_disable_mtp_completeness_check_is_scoped():
13
+ assert is_mtp_completeness_check_enabled()
14
+
15
+ with pytest.raises(RuntimeError), disable_mtp_completeness_check():
16
+ assert not is_mtp_completeness_check_enabled()
17
+ raise RuntimeError
18
+
19
+ assert is_mtp_completeness_check_enabled()
vllm/distributed/weight_transfer/ipc_engine.py
+6
-1
@@ -274,7 +274,12 @@ class IPCWeightTransferEngine(
274
weight = rebuild_cuda_tensor(*list_args)
275
weights.append((name, weight))
276
277
- self.model.load_weights(weights)
277
+ from vllm.model_executor.model_loader.mtp_validation import (
278
+ disable_mtp_completeness_check,
279
+ )
280
+
281
+ with disable_mtp_completeness_check():
282
+ self.model.load_weights(weights)
283
284
def shutdown(self) -> None:
285
pass
vllm/distributed/weight_transfer/nccl_engine.py
+31
-26
@@ -168,36 +168,41 @@ class NCCLWeightTransferEngine(
168
"Call init_transfer_engine() first."
169
)
170
171
- if update_info.packed:
172
- # Build iterator of (name, (shape, dtype)) from update_info
173
- def state_dict_info_iterator():
171
+ from vllm.model_executor.model_loader.mtp_validation import (
172
+ disable_mtp_completeness_check,
173
+ )
174
+
175
+ with disable_mtp_completeness_check():
176
+ if update_info.packed:
177
+ # Build iterator of (name, (shape, dtype)) from update_info
178
+ def state_dict_info_iterator():
179
+ for name, dtype_name, shape in zip(
180
+ update_info.names, update_info.dtype_names, update_info.shapes
181
+ ):
182
+ dtype = getattr(torch, dtype_name)
183
+ yield (name, (shape, dtype))
184
+
185
+ packed_nccl_broadcast_consumer(
186
+ iterator=state_dict_info_iterator(),
187
+ group=self.model_update_group,
188
+ src=0,
189
+ post_unpack_func=self.model.load_weights,
190
+ buffer_size_bytes=update_info.packed_buffer_size_bytes,
191
+ num_buffers=update_info.packed_num_buffers,
192
+ device=self.device,
193
+ )
194
+ else:
195
+ # Use simple one-by-one broadcasting
196
for name, dtype_name, shape in zip(
197
update_info.names, update_info.dtype_names, update_info.shapes
198
):
199
dtype = getattr(torch, dtype_name)
178
- yield (name, (shape, dtype))
179
-
180
- packed_nccl_broadcast_consumer(
181
- iterator=state_dict_info_iterator(),
182
- group=self.model_update_group,
183
- src=0,
184
- post_unpack_func=self.model.load_weights,
185
- buffer_size_bytes=update_info.packed_buffer_size_bytes,
186
- num_buffers=update_info.packed_num_buffers,
187
- device=self.device,
188
- )
189
- else:
190
- # Use simple one-by-one broadcasting
191
- for name, dtype_name, shape in zip(
192
- update_info.names, update_info.dtype_names, update_info.shapes
193
- ):
194
- dtype = getattr(torch, dtype_name)
195
- weight = torch.empty(shape, dtype=dtype, device=self.device)
196
- self.model_update_group.broadcast(
197
- weight, src=0, stream=torch.cuda.current_stream()
198
- )
199
- self.model.load_weights([(name, weight)])
200
- del weight
200
+ weight = torch.empty(shape, dtype=dtype, device=self.device)
201
+ self.model_update_group.broadcast(
202
+ weight, src=0, stream=torch.cuda.current_stream()
203
+ )
204
+ self.model.load_weights([(name, weight)])
205
+ del weight
206
207
def shutdown(self) -> None:
208
if self.model_update_group is not None:
vllm/model_executor/model_loader/mtp_validation.py
new
+26
@@ -0,0 +1,26 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Scoped controls for MTP checkpoint completeness validation."""
4
+
5
+from collections.abc import Iterator
6
+from contextlib import contextmanager
7
+from contextvars import ContextVar
8
+
9
+_mtp_completeness_check_enabled: ContextVar[bool] = ContextVar(
10
+ "mtp_completeness_check_enabled", default=True
11
+)
12
+
13
+
14
+def is_mtp_completeness_check_enabled() -> bool:
15
+ """Return whether MTP completeness validation is enabled in this scope."""
16
+ return _mtp_completeness_check_enabled.get()
17
+
18
+
19
+@contextmanager
20
+def disable_mtp_completeness_check() -> Iterator[None]:
21
+ """Temporarily disable MTP completeness validation for one weight load."""
22
+ token = _mtp_completeness_check_enabled.set(False)
23
+ try:
24
+ yield
25
+ finally:
26
+ _mtp_completeness_check_enabled.reset(token)
vllm/model_executor/models/bailing_moe_mtp.py
+7
-1
@@ -20,6 +20,9 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
20
ParallelLMHead,
21
VocabParallelEmbedding,
22
)
23
+from vllm.model_executor.model_loader.mtp_validation import (
24
+ is_mtp_completeness_check_enabled,
25
+)
26
from vllm.model_executor.model_loader.weight_utils import (
27
default_weight_loader,
28
maybe_remap_kv_scale_name,
@@ -370,7 +373,10 @@ class BailingMoeV25MTPModel(nn.Module):
373
self.model.mtp_start_layer_idx,
374
self.model.mtp_start_layer_idx + self.model.num_mtp_layers,
375
):
373
- if layer_idx not in loaded_mtp_layers:
376
+ if (
377
+ layer_idx not in loaded_mtp_layers
378
+ and is_mtp_completeness_check_enabled()
379
+ ):
380
raise ValueError(
381
f"Bailing MTP speculative decoding layer {layer_idx} "
382
"weights are missing from checkpoint. Use a checkpoint "
vllm/model_executor/models/deepseek_mtp.py
+4
-1
@@ -21,6 +21,9 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
21
ParallelLMHead,
22
VocabParallelEmbedding,
23
)
24
+from vllm.model_executor.model_loader.mtp_validation import (
25
+ is_mtp_completeness_check_enabled,
26
+)
27
from vllm.model_executor.model_loader.weight_utils import (
28
default_weight_loader,
29
maybe_remap_kv_scale_name,
@@ -516,7 +519,7 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts):
519
self.model.mtp_start_layer_idx,
520
self.model.mtp_start_layer_idx + self.model.num_mtp_layers,
521
):
519
- if layer_idx not in loaded_layers:
522
+ if layer_idx not in loaded_layers and is_mtp_completeness_check_enabled():
523
raise ValueError(
524
f"MTP speculative decoding layer {layer_idx} weights "
525
f"missing from checkpoint. The checkpoint may have "
vllm/model_executor/models/step3p5_mtp.py
+4
-1
@@ -15,6 +15,9 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
15
ParallelLMHead,
16
VocabParallelEmbedding,
17
)
18
+from vllm.model_executor.model_loader.mtp_validation import (
19
+ is_mtp_completeness_check_enabled,
20
+)
21
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
22
from vllm.sequence import IntermediateTensors
23
@@ -283,7 +286,7 @@ class Step3p5MTP(nn.Module):
286
and getattr(param, "requires_grad", False) is False
287
}
288
params_need_to_load -= optional_params
286
- if params_need_to_load != loaded_params:
289
+ if params_need_to_load != loaded_params and is_mtp_completeness_check_enabled():
290
missing_params = list(params_need_to_load - loaded_params)
291
param_name_example = missing_params[0]
292
raise RuntimeError(
vllm/models/deepseek_v32/nvidia/mtp.py
+4
-1
@@ -17,6 +17,9 @@ from vllm.model_executor.layers.logits_processor import LogitsProcessor
17
from vllm.model_executor.layers.vocab_parallel_embedding import (
18
VocabParallelEmbedding,
19
)
20
+from vllm.model_executor.model_loader.mtp_validation import (
21
+ is_mtp_completeness_check_enabled,
22
+)
23
from vllm.model_executor.model_loader.weight_utils import (
24
default_weight_loader,
25
maybe_remap_kv_scale_name,
@@ -416,7 +419,7 @@ class DeepseekV32MTP(nn.Module, DeepseekV2MixtureOfExperts):
419
self.model.mtp_start_layer_idx,
420
self.model.mtp_start_layer_idx + self.model.num_mtp_layers,
421
):
419
- if layer_idx not in loaded_layers:
422
+ if layer_idx not in loaded_layers and is_mtp_completeness_check_enabled():
423
raise ValueError(
424
f"MTP speculative decoding layer {layer_idx} weights "
425
f"missing from checkpoint."
vllm/models/deepseek_v4/amd/mtp.py
+4
-1
@@ -32,6 +32,9 @@ from vllm.model_executor.layers.mhc import HCHeadOp
32
from vllm.model_executor.layers.vocab_parallel_embedding import (
33
VocabParallelEmbedding,
34
)
35
+from vllm.model_executor.model_loader.mtp_validation import (
36
+ is_mtp_completeness_check_enabled,
37
+)
38
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
39
from vllm.model_executor.models.deepseek_mtp import SharedHead
40
from vllm.model_executor.models.deepseek_v2 import get_spec_layer_idx_from_weight_name
@@ -461,7 +464,7 @@ class DeepSeekV4MTP(nn.Module):
464
self.model.mtp_start_layer_idx,
465
self.model.mtp_start_layer_idx + self.model.num_mtp_layers,
466
):
464
- if layer_idx not in loaded_layers:
467
+ if layer_idx not in loaded_layers and is_mtp_completeness_check_enabled():
468
raise ValueError(
469
f"MTP speculative decoding layer {layer_idx} weights "
470
f"missing from checkpoint. The checkpoint may have "
vllm/models/deepseek_v4/nvidia/mtp.py
+4
-1
@@ -37,6 +37,9 @@ from vllm.model_executor.layers.logits_processor import LogitsProcessor
37
from vllm.model_executor.layers.vocab_parallel_embedding import (
38
VocabParallelEmbedding,
39
)
40
+from vllm.model_executor.model_loader.mtp_validation import (
41
+ is_mtp_completeness_check_enabled,
42
+)
43
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
44
from vllm.model_executor.models.deepseek_mtp import SharedHead
45
from vllm.model_executor.models.deepseek_v2 import get_spec_layer_idx_from_weight_name
@@ -467,7 +470,7 @@ class DeepSeekV4MTP(nn.Module):
470
self.model.mtp_start_layer_idx,
471
self.model.mtp_start_layer_idx + self.model.num_mtp_layers,
472
):
470
- if layer_idx not in loaded_layers:
473
+ if layer_idx not in loaded_layers and is_mtp_completeness_check_enabled():
474
raise ValueError(
475
f"MTP speculative decoding layer {layer_idx} weights "
476
f"missing from checkpoint. The checkpoint may have "
vllm/models/deepseek_v4/xpu/mtp.py
+4
-1
@@ -32,6 +32,9 @@ from vllm.model_executor.layers.mhc import HCHeadOp
32
from vllm.model_executor.layers.vocab_parallel_embedding import (
33
VocabParallelEmbedding,
34
)
35
+from vllm.model_executor.model_loader.mtp_validation import (
36
+ is_mtp_completeness_check_enabled,
37
+)
38
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
39
from vllm.model_executor.models.deepseek_mtp import SharedHead
40
from vllm.model_executor.models.deepseek_v2 import get_spec_layer_idx_from_weight_name
@@ -469,7 +472,7 @@ class DeepSeekV4MTP(nn.Module):
472
self.model.mtp_start_layer_idx,
473
self.model.mtp_start_layer_idx + self.model.num_mtp_layers,
474
):
472
- if layer_idx not in loaded_layers:
475
+ if layer_idx not in loaded_layers and is_mtp_completeness_check_enabled():
476
raise ValueError(
477
f"MTP speculative decoding layer {layer_idx} weights "
478
f"missing from checkpoint. The checkpoint may have "
vllm/models/inkling/nvidia/mtp.py
+4
-1
@@ -25,6 +25,9 @@ from vllm.config import VllmConfig
25
from vllm.model_executor.layers.linear import ReplicatedLinear
26
from vllm.model_executor.layers.logits_processor import LogitsProcessor
27
from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
28
+from vllm.model_executor.model_loader.mtp_validation import (
29
+ is_mtp_completeness_check_enabled,
30
+)
31
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
32
from vllm.model_executor.models.utils import maybe_prefix
33
from vllm.sequence import IntermediateTensors
@@ -396,7 +399,7 @@ def _load_inkling_mtp_weights(
399
for name in params
400
if name.startswith("model.layers.") or name.startswith("model.chain_norm.")
401
}
399
- if missing := sorted(required - loaded):
402
+ if (missing := sorted(required - loaded)) and is_mtp_completeness_check_enabled():
403
raise ValueError(
404
"Inkling MTP checkpoint is missing required parameters: "
405
+ ", ".join(missing)
vllm/models/minimax_m3/amd/mtp.py
+7
-1
@@ -35,6 +35,9 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
35
ParallelLMHead,
36
VocabParallelEmbedding,
37
)
38
+from vllm.model_executor.model_loader.mtp_validation import (
39
+ is_mtp_completeness_check_enabled,
40
+)
41
from vllm.model_executor.model_loader.weight_utils import (
42
default_weight_loader,
43
maybe_remap_kv_scale_name,
@@ -322,7 +325,10 @@ class MiniMaxM3MTP(nn.Module):
325
326
# Validate that weights were loaded for each MTP layer.
327
for layer_idx in range(self.model.num_mtp_layers):
325
- if layer_idx not in loaded_mtp_layers:
328
+ if (
329
+ layer_idx not in loaded_mtp_layers
330
+ and is_mtp_completeness_check_enabled()
331
+ ):
332
raise ValueError(
333
f"Failed to load MTP layer {layer_idx} weights from checkpoint."
334
)
vllm/models/minimax_m3/nvidia/mtp.py
+7
-1
@@ -19,6 +19,9 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
19
ParallelLMHead,
20
VocabParallelEmbedding,
21
)
22
+from vllm.model_executor.model_loader.mtp_validation import (
23
+ is_mtp_completeness_check_enabled,
24
+)
25
from vllm.model_executor.model_loader.weight_utils import (
26
default_weight_loader,
27
maybe_remap_kv_scale_name,
@@ -304,7 +307,10 @@ class MiniMaxM3MTP(nn.Module):
307
308
# Validate that weights were loaded for each MTP layer.
309
for layer_idx in range(self.model.num_mtp_layers):
307
- if layer_idx not in loaded_mtp_layers:
310
+ if (
311
+ layer_idx not in loaded_mtp_layers
312
+ and is_mtp_completeness_check_enabled()
313
+ ):
314
raise ValueError(
315
f"Failed to load MTP layer {layer_idx} weights from checkpoint."
316
)