@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 160 lines 5.27 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from unittest.mock import Mock, patch
5
6 import pytest
7 import torch
8
9 from vllm.config import LoadConfig, ModelConfig, SpeculativeConfig, VllmConfig
10 from vllm.model_executor.models.utils import get_draft_quant_config
11 from vllm.platforms import current_platform
12
13 DEVICE_TYPE = current_platform.device_type
14 DEVICES = (
15 [f"{DEVICE_TYPE}:{i}" for i in range(min(torch.accelerator.device_count(), 2))]
16 if not current_platform.is_cpu()
17 else ["cpu"]
18 )
19
20
21 def test_get_draft_quant_config_with_draft_model():
22 mock_draft_model_config = Mock(spec=ModelConfig)
23 mock_load_config = Mock(spec=LoadConfig)
24 mock_speculative_config = Mock(spec=SpeculativeConfig)
25 mock_speculative_config.draft_model_config = mock_draft_model_config
26
27 mock_vllm_config = Mock(spec=VllmConfig)
28 mock_vllm_config.speculative_config = mock_speculative_config
29 mock_vllm_config.load_config = mock_load_config
30
31 mock_quant_config = Mock()
32 with patch.object(
33 VllmConfig, "get_quantization_config", return_value=mock_quant_config
34 ):
35 result = get_draft_quant_config(mock_vllm_config)
36
37 # Verify the function calls get_quantization_config with draft model config
38 VllmConfig.get_quantization_config.assert_called_once_with(
39 mock_draft_model_config, mock_load_config
40 )
41 assert result == mock_quant_config
42
43
44 def test_get_draft_quant_config_without_draft_model():
45 mock_speculative_config = Mock(spec=SpeculativeConfig)
46 mock_speculative_config.draft_model_config = None
47
48 mock_vllm_config = Mock(spec=VllmConfig)
49 mock_vllm_config.speculative_config = mock_speculative_config
50 mock_vllm_config.load_config = Mock(spec=LoadConfig)
51
52 result = get_draft_quant_config(mock_vllm_config)
53
54 assert result is None
55
56
57 @torch.inference_mode()
58 @pytest.mark.parametrize("device", DEVICES)
59 def test_fc_layer_quant_config_usage(default_vllm_config, dist_init, device) -> None:
60 import torch
61
62 from vllm.model_executor.layers.linear import ReplicatedLinear
63
64 if current_platform.is_cuda_alike():
65 torch.accelerator.set_device_index(device)
66
67 torch.set_default_device(device)
68
69 input_size = 256
70 output_size = 128
71
72 fc_no_quant = ReplicatedLinear(
73 input_size=input_size,
74 output_size=output_size,
75 bias=False,
76 params_dtype=torch.float16,
77 quant_config=None,
78 prefix="fc",
79 )
80
81 assert fc_no_quant.quant_config is None
82 assert fc_no_quant.input_size == input_size
83 assert fc_no_quant.output_size == output_size
84
85 mock_quant_config = Mock()
86 mock_quant_config.online_quantization_config = None
87 fc_with_quant = ReplicatedLinear(
88 input_size=input_size,
89 output_size=output_size,
90 bias=False,
91 params_dtype=torch.float16,
92 quant_config=mock_quant_config,
93 prefix="fc",
94 )
95
96 assert fc_with_quant.quant_config == mock_quant_config
97
98 # Check forward pass
99 x = torch.randn(2, input_size, dtype=torch.float16)
100 output, _ = fc_no_quant(x)
101 assert output.shape == (2, output_size)
102
103
104 def test_maybe_remap_kv_scale_name():
105 from vllm.model_executor.model_loader.weight_utils import maybe_remap_kv_scale_name
106
107 params_dict = {
108 "layers.0.self_attn.kv_scale": Mock(),
109 "layers.1.self_attn.kv_scale": Mock(),
110 }
111
112 name = "layers.0.self_attn.some_scale"
113 remapped = maybe_remap_kv_scale_name(name, params_dict)
114
115 assert remapped in params_dict or remapped == name or remapped is None
116
117
118 def test_eagle3_lm_head_receives_quant_config():
119 """Eagle3LlamaForCausalLM must pass quant_config to ParallelLMHead.
120
121 Without quant_config, quantized lm_head weights (e.g. INT8 per-channel)
122 in Eagle3 drafter checkpoints fail to load because ParallelLMHead doesn't
123 expect weight_packed tensors.
124 """
125 from vllm.model_executor.models.llama_eagle3 import Eagle3LlamaForCausalLM
126
127 mock_quant_config = Mock()
128
129 mock_hf_config = Mock()
130 mock_hf_config.draft_vocab_size = 1000
131 mock_hf_config.hidden_size = 256
132 mock_hf_config.vocab_size = 32000
133 mock_hf_config.logit_scale = 1.0
134
135 mock_vllm_config = Mock()
136 mock_vllm_config.speculative_config.draft_model_config.hf_config = mock_hf_config
137 mock_vllm_config.model_config.get_num_layers.return_value = 32
138 mock_vllm_config.speculative_config.parallel_drafting = False
139
140 with (
141 patch("vllm.model_executor.models.llama_eagle3.LlamaModel") as MockModel,
142 patch("vllm.model_executor.models.llama_eagle3.ParallelLMHead") as MockLMHead,
143 patch("vllm.model_executor.models.llama_eagle3.LogitsProcessor"),
144 patch(
145 "vllm.model_executor.models.llama_eagle3.get_draft_quant_config",
146 return_value=mock_quant_config,
147 ),
148 ):
149 MockModel.return_value.use_aux_hidden_state = True
150
151 Eagle3LlamaForCausalLM(vllm_config=mock_vllm_config)
152
153 MockLMHead.assert_called_once()
154 call_kwargs = MockLMHead.call_args.kwargs
155 assert "quant_config" in call_kwargs, (
156 "ParallelLMHead must receive quant_config for quantized lm_head weights"
157 )
158 assert call_kwargs["quant_config"] is mock_quant_config, (
159 "ParallelLMHead must receive the draft model's quant_config"
160 )