A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 07:15 UTC
[Bugfix] Fix EAGLE vocab embedding construction for Llama 70B (#19033)
Signed-off-by: Benjamin Chislett <benjamin.chislett@centml.ai>
Benjamin Chislett committed
Jun 5, 2025 at 22:10 UTC
3465b87ef8b36d91d8746f29076270d283a5251a
7 files changed
+70
-47
benchmarks/kernels/bench_fp8_gemm.py
+1
-1
@@ -5,11 +5,11 @@ import copy
5
import itertools
6
7
import torch
8
-import triton
8
from weight_shapes import WEIGHT_SHAPES
9
10
from vllm._custom_ops import cutlass_scaled_mm as vllm_scaled_mm
11
from vllm._custom_ops import scaled_fp8_quant as vllm_scaled_fp8_quant
12
+from vllm.triton_utils import triton
13
14
15
@triton.testing.perf_report(
tests/v1/spec_decode/test_eagle.py
+40
-24
@@ -9,6 +9,7 @@ import torch
9
from vllm.config import (CacheConfig, DeviceConfig, LoadConfig, ModelConfig,
10
ParallelConfig, SchedulerConfig, SpeculativeConfig,
11
VllmConfig)
12
+from vllm.model_executor.models.llama import LlamaForCausalLM
13
from vllm.v1.spec_decode.eagle import EagleProposer
14
15
model_dir = "meta-llama/Llama-3.1-8B-Instruct"
@@ -113,21 +114,26 @@ def test_prepare_inputs():
114
assert torch.equal(token_indices, expected_token_indices)
115
116
116
-@pytest.mark.parametrize(
117
- "method,proposer_helper,draft_model_dir,target_attribute_path", [
118
- ("eagle", lambda k: _create_proposer("eagle", k), eagle_dir,
119
- ('lm_head', )),
120
- ("eagle3", lambda k: _create_proposer("eagle3", k), eagle3_dir,
121
- ('model', 'embed_tokens')),
122
- ])
117
+@pytest.mark.parametrize("method,proposer_helper", [
118
+ ("eagle", lambda k: _create_proposer("eagle", k)),
119
+ ("eagle3", lambda k: _create_proposer("eagle3", k)),
120
+])
121
+@pytest.mark.parametrize("pp_size", [1, 2])
122
+@pytest.mark.parametrize("use_distinct_embed_tokens", [True, False])
123
@mock.patch('vllm.v1.spec_decode.eagle.get_pp_group')
124
@mock.patch('vllm.v1.spec_decode.eagle.get_layers_from_vllm_config')
125
@mock.patch('vllm.v1.spec_decode.eagle.get_model')
126
def test_load_model(mock_get_model, mock_get_layers, mock_get_pp_group, method,
127
- proposer_helper, draft_model_dir, target_attribute_path):
128
-
129
- # Setup model mock
127
+ proposer_helper, pp_size, use_distinct_embed_tokens):
128
+ # Setup draft model mock
129
mock_model = mock.MagicMock()
130
+ if use_distinct_embed_tokens:
131
+ # Some models can have a different hidden size than the target model,
132
+ # so we test that their embed_tokens doesn't get overwritten
133
+ mock_model.model.embed_tokens.weight.shape = (131072, 2048)
134
+ else:
135
+ mock_model.model.embed_tokens.weight.shape = (131072, 4096)
136
+
137
mock_get_model.return_value = mock_model
138
139
# Setup mocks for attention layers
@@ -145,22 +151,24 @@ def test_load_model(mock_get_model, mock_get_layers, mock_get_pp_group, method,
151
152
# Setup mock for pp group to return the appropriate value for world size
153
mock_pp_group = mock.MagicMock()
148
- mock_pp_group.world_size = 2 if method == "eagle" else 1
154
+ mock_pp_group.world_size = pp_size
155
mock_get_pp_group.return_value = mock_pp_group
156
151
- # Setup target model with the appropriate attributes
152
- target_model = mock.MagicMock()
157
+ # Setup the target model mock with a custom class so that
158
+ # isinstance() checks match the expected type.
159
+ class _TargetModelStub(LlamaForCausalLM):
160
+ model: mock.MagicMock
161
+ lm_head: mock.MagicMock
162
154
- # Create the necessary attributes on the target model
155
- current_obj = target_model
156
- for i, attr in enumerate(target_attribute_path):
157
- if i == len(target_attribute_path) - 1:
158
- # Set the last attribute in the path to a MagicMock
159
- setattr(current_obj, attr, mock.MagicMock())
160
- else:
161
- # Create intermediate objects if needed
162
- setattr(current_obj, attr, mock.MagicMock())
163
- current_obj = getattr(current_obj, attr)
163
+ target_model = mock.create_autospec(_TargetModelStub, instance=True)
164
+ target_model.model = mock.MagicMock()
165
+ target_model.model.embed_tokens.weight.shape = (131072, 4096)
166
+
167
+ from vllm.model_executor.models import SupportsMultiModal
168
+ assert not isinstance(target_model, SupportsMultiModal)
169
+
170
+ if method == "eagle":
171
+ target_model.lm_head = mock.MagicMock()
172
173
# Create proposer using the helper function
174
proposer = proposer_helper(k=8)
@@ -171,10 +179,18 @@ def test_load_model(mock_get_model, mock_get_layers, mock_get_pp_group, method,
179
# Verify common interactions
180
mock_get_model.assert_called_once()
181
174
- # Verify the specific attribute sharing based on the method
182
+ # Verify that EAGLE models gain the lm head from the target model
183
if method == "eagle":
184
assert proposer.model.lm_head == target_model.lm_head
185
+
186
+ # Verify that the embed tokens are set correctly
187
+ # If pp_size is > 1, the embed tokens should be distinct
188
+ if pp_size > 1 or use_distinct_embed_tokens:
189
+ assert proposer.model.model.embed_tokens != \
190
+ target_model.model.embed_tokens
191
else:
192
+ # When pp_size is 1 and the draft and target models have
193
+ # embed_tokens of the same shape, they should be shared.
194
assert proposer.model.model.embed_tokens == \
195
target_model.model.embed_tokens
196
vllm/model_executor/models/llama_eagle.py
+6
-8
@@ -55,13 +55,11 @@ class LlamaModel(nn.Module):
55
speculative_config.draft_model_config.hf_config
56
self.vocab_size = self.config.vocab_size
57
58
- # if PP disabled then draft will share embed with target
59
- if get_pp_group().world_size > 1:
60
- self.embed_tokens = VocabParallelEmbedding(
61
- self.config.vocab_size,
62
- self.config.hidden_size,
63
- prefix=maybe_prefix(prefix, "embed_tokens"),
64
- )
58
+ self.embed_tokens = VocabParallelEmbedding(
59
+ self.config.vocab_size,
60
+ self.config.hidden_size,
61
+ prefix=maybe_prefix(prefix, "embed_tokens"),
62
+ )
63
64
self.layers = nn.ModuleList([
65
LlamaDecoderLayer(
@@ -164,4 +162,4 @@ class EagleLlamaForCausalLM(LlamaForCausalLM):
162
if "lm_head" not in name:
163
name = "model." + name
164
model_weights[name] = loaded_weight
167
- return loader.load_weights(model_weights.items())
165
+ loader.load_weights(model_weights.items())
vllm/model_executor/models/llama_eagle3.py
+14
-10
@@ -10,7 +10,6 @@ from transformers import LlamaConfig
10
11
from vllm.compilation.decorators import support_torch_compile
12
from vllm.config import VllmConfig
13
-from vllm.distributed.parallel_state import get_pp_group
13
from vllm.logger import init_logger
14
from vllm.model_executor.layers.layernorm import RMSNorm
15
from vllm.model_executor.layers.linear import QKVParallelLinear
@@ -95,13 +94,11 @@ class LlamaModel(nn.Module):
94
speculative_config.draft_model_config.hf_config
95
self.vocab_size = self.config.vocab_size
96
98
- # if PP disabled then draft will share embed with target
99
- if get_pp_group().world_size > 1:
100
- self.embed_tokens = VocabParallelEmbedding(
101
- self.config.vocab_size,
102
- self.config.hidden_size,
103
- prefix=maybe_prefix(prefix, "embed_tokens"),
104
- )
97
+ self.embed_tokens = VocabParallelEmbedding(
98
+ self.config.vocab_size,
99
+ self.config.hidden_size,
100
+ prefix=maybe_prefix(prefix, "embed_tokens"),
101
+ )
102
103
self.layers = nn.ModuleList([
104
LlamaDecoderLayer(
@@ -240,6 +237,7 @@ class Eagle3LlamaForCausalLM(LlamaForCausalLM):
237
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
238
model_weights = {}
239
includes_draft_id_mapping = False
240
+ includes_embed_tokens = False
241
for name, loaded_weight in weights:
242
if "t2d" in name:
243
continue
@@ -248,12 +246,18 @@ class Eagle3LlamaForCausalLM(LlamaForCausalLM):
246
includes_draft_id_mapping = True
247
elif "lm_head" not in name:
248
name = "model." + name
249
+ if "embed_tokens" in name:
250
+ includes_embed_tokens = True
251
model_weights[name] = loaded_weight
252
253
+ skip_substrs = []
254
+ if not includes_draft_id_mapping:
255
+ skip_substrs.append("draft_id_to_target_id")
256
+ if not includes_embed_tokens:
257
+ skip_substrs.append("embed_tokens")
258
loader = AutoWeightsLoader(
259
self,
260
skip_prefixes=None,
256
- skip_substrs=["draft_id_to_target_id"] \
257
- if not includes_draft_id_mapping else None,
261
+ skip_substrs=skip_substrs,
262
)
263
loader.load_weights(model_weights.items())
vllm/platforms/cuda.py
+1
@@ -173,6 +173,7 @@ class CudaPlatformBase(Platform):
173
def get_current_memory_usage(cls,
174
device: Optional[torch.types.Device] = None
175
) -> float:
176
+ torch.cuda.empty_cache()
177
torch.cuda.reset_peak_memory_stats(device)
178
return torch.cuda.max_memory_allocated(device)
179
vllm/utils.py
+1
@@ -899,6 +899,7 @@ class DeviceMemoryProfiler:
899
def current_memory_usage(self) -> float:
900
# Return the memory usage in bytes.
901
from vllm.platforms import current_platform
902
+ gc.collect()
903
return current_platform.get_current_memory_usage(self.device)
904
905
def __enter__(self):
vllm/v1/spec_decode/eagle.py
+7
-4
@@ -330,16 +330,19 @@ class EagleProposer:
330
self.attn_layer_names = list(draft_attn_layer_names)
331
332
# share embed_tokens with the target model if needed
333
- if get_pp_group().world_size == 1:
333
+ if get_pp_group().world_size == 1 \
334
+ and self.model.model.embed_tokens.weight.shape \
335
+ == target_model.model.embed_tokens.weight.shape:
336
logger.info(
335
- "The EAGLE head shares the same vocab embedding" \
337
+ "Assuming the EAGLE head shares the same vocab embedding" \
338
" with the target model."
339
)
340
+ del self.model.model.embed_tokens
341
self.model.model.embed_tokens = target_model.model.embed_tokens
342
else:
343
logger.info(
341
- "Since PP > 1, the EAGLE head loaded its own vocab embedding" \
342
- " weights instead of sharing them with the target model."
344
+ "The EAGLE head's vocab embedding will be loaded separately" \
345
+ " from the target model."
346
)
347
348
# share lm_head with the target model if needed