@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 116 lines 4.58 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Tests for the EAGLE draft ``max_position_embeddings`` override (#48894).
4
5 EAGLE drafts share the target's positional space, but some draft
6 checkpoints (e.g. ``yuhuili/EAGLE3-LLaMA3.1-Instruct-8B``) ship a
7 ``max_position_embeddings`` (2048) far smaller than the target's context.
8 That value sizes the draft's rotary ``cos_sin_cache`` while the proposer
9 feeds positions up to the target's ``max_model_len``, so the cache gather
10 goes out of bounds — a device-side assert under torch.compile and silent
11 garbage reads in eager mode. ``SpeculativeConfig`` must raise the draft's
12 value to the target's ``max_model_len``, with a log, for the eagle/eagle3
13 methods only.
14 """
15
16 import logging
17 from typing import Literal
18
19 import pytest
20 from transformers import PreTrainedConfig
21
22 from vllm.config.model import ModelConfig
23 from vllm.config.parallel import ParallelConfig
24 from vllm.config.speculative import SpeculativeConfig
25
26 # All repos are public; only config/tokenizer-config files are fetched.
27 EAGLE3_DRAFT = "yuhuili/EAGLE3-LLaMA3.1-Instruct-8B" # max_position_embeddings=2048
28 LLAMA3_TARGET = "unsloth/Meta-Llama-3.1-8B-Instruct" # max_position_embeddings=131072
29 AR_MODEL = "JackFram/llama-68m" # max_position_embeddings=2048
30
31 _LOGGER = "vllm.config.speculative"
32 _OVERRIDE_MSG = "Overriding draft model max_position_embeddings"
33
34
35 @pytest.fixture
36 def vllm_caplog(caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch):
37 """Make caplog see vLLM logger records (vLLM sets propagate=False)."""
38 monkeypatch.setattr(logging.getLogger("vllm"), "propagate", True)
39 with caplog.at_level(logging.INFO, logger=_LOGGER):
40 yield caplog
41
42
43 def _override_logged(caplog: pytest.LogCaptureFixture) -> bool:
44 return any(_OVERRIDE_MSG in record.getMessage() for record in caplog.records)
45
46
47 @pytest.mark.cpu_test
48 def test_override_raises_smaller_value(vllm_caplog: pytest.LogCaptureFixture):
49 hf_config = PreTrainedConfig(max_position_embeddings=2048)
50 SpeculativeConfig._maybe_override_draft_max_position_embeddings(
51 hf_config, target_max_model_len=8192
52 )
53 assert hf_config.max_position_embeddings == 8192
54 assert _override_logged(vllm_caplog)
55
56
57 @pytest.mark.cpu_test
58 def test_override_keeps_sufficient_value(vllm_caplog: pytest.LogCaptureFixture):
59 hf_config = PreTrainedConfig(max_position_embeddings=8192)
60 SpeculativeConfig._maybe_override_draft_max_position_embeddings(
61 hf_config, target_max_model_len=8192
62 )
63 assert hf_config.max_position_embeddings == 8192
64 assert not _override_logged(vllm_caplog)
65
66
67 @pytest.mark.cpu_test
68 def test_override_ignores_missing_attribute(vllm_caplog: pytest.LogCaptureFixture):
69 hf_config = PreTrainedConfig()
70 hf_config.__dict__.pop("max_position_embeddings", None)
71 SpeculativeConfig._maybe_override_draft_max_position_embeddings(
72 hf_config, target_max_model_len=8192
73 )
74 assert not hasattr(hf_config, "max_position_embeddings")
75 assert not _override_logged(vllm_caplog)
76
77
78 @pytest.mark.cpu_test
79 @pytest.mark.parametrize("method", ["eagle", "eagle3"])
80 def test_eagle_draft_inherits_target_max_model_len(
81 method: Literal["eagle", "eagle3"], vllm_caplog: pytest.LogCaptureFixture
82 ):
83 target_model_config = ModelConfig(LLAMA3_TARGET)
84 assert target_model_config.max_model_len > 2048
85 speculative_config = SpeculativeConfig(
86 target_model_config=target_model_config,
87 target_parallel_config=ParallelConfig(),
88 model=EAGLE3_DRAFT,
89 method=method,
90 num_speculative_tokens=3,
91 )
92 draft_hf_config = speculative_config.draft_model_config.hf_config
93 assert draft_hf_config.max_position_embeddings == target_model_config.max_model_len
94 assert _override_logged(vllm_caplog)
95
96
97 @pytest.mark.cpu_test
98 def test_independent_draft_model_keeps_its_own_limit(
99 vllm_caplog: pytest.LogCaptureFixture,
100 ):
101 """An independent AR draft may genuinely have a smaller context than the
102 target; its max_position_embeddings must not be resized."""
103 target_model_config = ModelConfig(
104 AR_MODEL, hf_overrides={"max_position_embeddings": 8192}
105 )
106 assert target_model_config.max_model_len == 8192
107 speculative_config = SpeculativeConfig(
108 target_model_config=target_model_config,
109 target_parallel_config=ParallelConfig(),
110 model=AR_MODEL,
111 method="draft_model",
112 num_speculative_tokens=3,
113 )
114 draft_hf_config = speculative_config.draft_model_config.hf_config
115 assert draft_hf_config.max_position_embeddings == 2048
116 assert not _override_logged(vllm_caplog)