[Bugfix] Forward callable hf_overrides to the draft model config (#45352)
Signed-off-by: HumphreySun98 <humphreysun98@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
Humphrey committed
Jul 6, 2026 at 23:12 UTC
cbe9c40f998f13975b967773ac7e7920e115387f
3 files changed
+148
-3
tests/config/test_speculative_draft_hf_overrides.py
new
+106
@@ -0,0 +1,106 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Tests for SpeculativeConfig.compose_draft_hf_overrides.
4
+
5
+Callable ``hf_overrides`` on the target model config (e.g. the
6
+``dummy_hf_overrides`` shrink used by ``tests/models/test_initialization.py``)
7
+must also be applied when building the draft ``ModelConfig``. Otherwise a
8
+draft belonging to a large target model is instantiated at full size even
9
+when the target itself is shrunk — which is what kept spec-decode archs like
10
+``EagleMistralLarge3ForCausalLM`` stuck at ``is_available_online=False``
11
+("TODO: revert once figuring out OOM in CI").
12
+"""
13
+
14
+import functools
15
+
16
+import pytest
17
+from transformers import PretrainedConfig
18
+
19
+from vllm.config.speculative import SpeculativeConfig
20
+
21
+
22
+def _make_hf_config(**kwargs) -> PretrainedConfig:
23
+ defaults = dict(
24
+ architectures=["LlamaForCausalLM"],
25
+ model_type="llama",
26
+ num_hidden_layers=64,
27
+ )
28
+ defaults.update(kwargs)
29
+ return PretrainedConfig(**defaults)
30
+
31
+
32
+@pytest.mark.cpu_test
33
+def test_dict_overrides_are_not_forwarded_to_draft():
34
+ """Dict overrides are target-specific key patches; the draft must get
35
+ only the architecture-mapping override."""
36
+ composed = SpeculativeConfig.compose_draft_hf_overrides(
37
+ {"max_position_embeddings": 1234}
38
+ )
39
+ assert composed is SpeculativeConfig.hf_config_override
40
+
41
+
42
+@pytest.mark.cpu_test
43
+def test_none_overrides_fall_back_to_arch_mapping():
44
+ composed = SpeculativeConfig.compose_draft_hf_overrides(None)
45
+ assert composed is SpeculativeConfig.hf_config_override
46
+
47
+
48
+@pytest.mark.cpu_test
49
+def test_callable_overrides_reach_the_draft_config():
50
+ """A callable override (config-to-config transform) composes with the
51
+ architecture-mapping override and is applied to the draft config."""
52
+
53
+ def shrink(hf_config: PretrainedConfig) -> PretrainedConfig:
54
+ hf_config.num_hidden_layers = 1
55
+ return hf_config
56
+
57
+ composed = SpeculativeConfig.compose_draft_hf_overrides(shrink)
58
+ assert composed is not SpeculativeConfig.hf_config_override
59
+
60
+ out = composed(_make_hf_config())
61
+ # The shrink transform must have been applied to the draft config.
62
+ assert out.num_hidden_layers == 1
63
+
64
+
65
+@pytest.mark.cpu_test
66
+def test_arch_mapping_applies_before_callable_override():
67
+ """The static arch-mapping override runs first, so the user callable
68
+ observes (and may adjust) the post-mapping config."""
69
+ seen_architectures: list[str] = []
70
+
71
+ def record(hf_config: PretrainedConfig) -> PretrainedConfig:
72
+ seen_architectures.append(hf_config.architectures[0])
73
+ return hf_config
74
+
75
+ composed = SpeculativeConfig.compose_draft_hf_overrides(record)
76
+
77
+ # MiMo is one of the arch-mapped model types: hf_config_override
78
+ # rewrites architectures to ["MiMoMTPModel"].
79
+ mimo = _make_hf_config(
80
+ architectures=["MiMoForCausalLM"],
81
+ model_type="mimo",
82
+ num_nextn_predict_layers=1,
83
+ )
84
+ composed(mimo)
85
+ assert seen_architectures == ["MiMoMTPModel"]
86
+
87
+
88
+def _module_level_shrink(hf_config: PretrainedConfig) -> PretrainedConfig:
89
+ hf_config.num_hidden_layers = 1
90
+ return hf_config
91
+
92
+
93
+@pytest.mark.cpu_test
94
+def test_composed_override_is_picklable():
95
+ """The draft ``ModelConfig`` is sent to spawned engine-core processes, so
96
+ the composed override must be picklable. A nested local closure is not
97
+ (it raised ``Can't get local object`` on DFlashDraftModel); a
98
+ ``functools.partial`` over a module-referenceable static method is.
99
+ Guard against regressing to a closure."""
100
+ composed = SpeculativeConfig.compose_draft_hf_overrides(_module_level_shrink)
101
+
102
+ assert isinstance(composed, functools.partial)
103
+ assert composed.func is SpeculativeConfig._apply_composed_hf_override
104
+
105
+ out = composed(_make_hf_config())
106
+ assert out.num_hidden_layers == 1
tests/models/registry.py
-2
@@ -1492,8 +1492,6 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
1492
"EagleMistralLarge3ForCausalLM": _HfExamplesInfo(
1493
"mistralai/Mistral-Large-3-675B-Instruct-2512",
1494
speculative_model="mistralai/Mistral-Large-3-675B-Instruct-2512-Eagle",
1495
- # TODO: revert once figuring out OOM in CI
1496
- is_available_online=False,
1495
),
1496
"LlamaForCausalLMEagle3": _HfExamplesInfo(
1497
"Qwen/Qwen3-8B",
vllm/config/speculative.py
+42
-1
@@ -2,6 +2,8 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import copy
5
+import functools
6
+from collections.abc import Callable
7
from typing import TYPE_CHECKING, Any, Literal, get_args
8
9
from pydantic import Field, SkipValidation, field_validator, model_validator
@@ -588,6 +590,40 @@ class SpeculativeConfig:
590
591
return hf_config
592
593
+ @staticmethod
594
+ def _apply_composed_hf_override(
595
+ target_hf_overrides: Callable[[PretrainedConfig], PretrainedConfig],
596
+ hf_config: PretrainedConfig,
597
+ ) -> PretrainedConfig:
598
+ hf_config = SpeculativeConfig.hf_config_override(hf_config)
599
+ return target_hf_overrides(hf_config)
600
+
601
+ @staticmethod
602
+ def compose_draft_hf_overrides(
603
+ target_hf_overrides: HfOverrides | None,
604
+ ) -> Callable[[PretrainedConfig], PretrainedConfig]:
605
+ """Build the ``hf_overrides`` for the draft ``ModelConfig``.
606
+
607
+ Callable overrides on the target are config-to-config transforms
608
+ (e.g. test harnesses shrinking ``num_hidden_layers``) and must also
609
+ reach the draft config — otherwise a draft belonging to a large
610
+ target is instantiated at full size even when the target is shrunk.
611
+ Dict overrides are target-specific key patches and are not applied
612
+ to the draft.
613
+
614
+ The composed override must stay picklable: the draft ``ModelConfig``
615
+ is sent to spawned engine-core processes, so a local closure would
616
+ fail with ``Can't get local object`` during pickling. Bind the
617
+ target via ``functools.partial`` over a module-referenceable static
618
+ method instead.
619
+ """
620
+ if not callable(target_hf_overrides):
621
+ return SpeculativeConfig.hf_config_override
622
+
623
+ return functools.partial(
624
+ SpeculativeConfig._apply_composed_hf_override, target_hf_overrides
625
+ )
626
+
627
def __post_init__(self):
628
# Note: "method" is a new parameter that helps to extend the
629
# configuration of non-model-based proposers, and the "model" parameter
@@ -752,7 +788,12 @@ class SpeculativeConfig:
788
if self.method == "medusa":
789
draft_hf_overrides = {"model_type": "medusa"}
790
else:
755
- draft_hf_overrides = SpeculativeConfig.hf_config_override
791
+ # Compose any callable hf_overrides set on the target so the
792
+ # draft config receives the same transform (e.g. the test
793
+ # shrink). Dict overrides stay target-only.
794
+ draft_hf_overrides = SpeculativeConfig.compose_draft_hf_overrides(
795
+ self.target_model_config.hf_overrides
796
+ )
797
self.draft_model_config = ModelConfig(
798
model=self.model,
799
runner="draft",