File size: 10,937 Bytes
4a28d4d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
"""Test quant_policy=QuantPolicy.TURBO_QUANT (K=4bit, V=2bit mixed precision)
for PytorchEngine.

This module tests both functional correctness and accuracy of quant_policy=QuantPolicy.TURBO_QUANT against a non-
quantized (quant_policy=QuantPolicy.NONE) baseline.
"""

import gc

import pytest
import torch

from lmdeploy import GenerationConfig, PytorchEngineConfig, pipeline
from lmdeploy.messages import QuantPolicy, Response

# Use smaller model to avoid OOM when running both quant_policy=QuantPolicy.NONE
# and quant_policy=QuantPolicy.TURBO_QUANT
MODEL_ID = 'Qwen/Qwen3-8B'


# =============================================================================
# Shared Fixtures
# =============================================================================

@pytest.fixture(scope='session')
def model_id():
    """Model ID for testing."""
    return MODEL_ID


@pytest.fixture(scope='session')
def pipe_no_quant(model_id):
    """Create pipeline without quantization (baseline).

    This fixture has session scope to avoid reloading the model for each test. Caller is responsible for cleanup.
    """
    engine_config = PytorchEngineConfig(
        tp=1,
        cache_max_entry_count=0.05,
        quant_policy=QuantPolicy.NONE,  # No quantization
    )
    pipe = pipeline(model_id, backend_config=engine_config, log_level='INFO')
    yield pipe
    # Cleanup
    pipe.close()
    del pipe
    gc.collect()
    if torch.cuda.is_available() and torch.cuda.device_count() > 0:
        torch.cuda.empty_cache()


@pytest.fixture(scope='session')
def pipe_quant_42(model_id):
    """Create pipeline with quant_policy=QuantPolicy.TURBO_QUANT.

    This fixture has session scope to avoid reloading the model for each test. Caller is responsible for cleanup.
    """
    engine_config = PytorchEngineConfig(
        tp=1,
        cache_max_entry_count=0.05,
        quant_policy=QuantPolicy.TURBO_QUANT,  # K=4bit, V=2bit mixed precision
    )
    pipe = pipeline(model_id, backend_config=engine_config, log_level='INFO')
    yield pipe
    # Cleanup
    pipe.close()
    del pipe
    gc.collect()
    if torch.cuda.is_available() and torch.cuda.device_count() > 0:
        torch.cuda.empty_cache()


# =============================================================================
# Basic Functional Tests (quant_policy=QuantPolicy.TURBO_QUANT only)
# =============================================================================

class TestQuantPolicy42Basic:
    """Basic functional tests for quant_policy=QuantPolicy.TURBO_QUANT.

    These tests verify that the quantized model can perform basic inference without errors. They test single prompt,
    batch prompts, and generation config.
    """

    @pytest.fixture(scope='class')
    def pipe(self):
        """Create pipeline with quant_policy=QuantPolicy.TURBO_QUANT."""
        engine_config = PytorchEngineConfig(
            tp=1,
            cache_max_entry_count=0.1,
            quant_policy=QuantPolicy.TURBO_QUANT,
        )
        pipe = pipeline(MODEL_ID, backend_config=engine_config, log_level='INFO')
        yield pipe
        pipe.close()
        del pipe
        gc.collect()
        if torch.cuda.is_available() and torch.cuda.device_count() > 0:
            torch.cuda.empty_cache()

    def test_infer_single_prompt(self, pipe):
        """Test single prompt inference with
        quant_policy=QuantPolicy.TURBO_QUANT."""
        prompt = 'Hello, how are you?'
        response = pipe.infer(prompt, max_new_tokens=30)

        assert isinstance(response, Response)
        assert hasattr(response, 'text')
        assert len(response.text) > 0
        assert len(response.text.strip()) > 0

    def test_infer_batch_prompts(self, pipe):
        """Test batch inference with quant_policy=QuantPolicy.TURBO_QUANT."""
        prompts = ['What is AI?', 'Hello!']
        responses = pipe.infer(prompts, max_new_tokens=20)

        assert isinstance(responses, list)
        assert len(responses) == len(prompts)
        for resp in responses:
            assert isinstance(resp, Response)
            assert len(resp.text) > 0

    def test_infer_with_generation_config(self, pipe):
        """Test inference with GenerationConfig."""
        gen_config = GenerationConfig(max_new_tokens=20, temperature=0.7)
        prompt = 'Tell me a short joke'
        response = pipe.infer(prompt, gen_config=gen_config)

        assert isinstance(response, Response)
        assert len(response.text) > 0


# =============================================================================
# Accuracy Tests (quant_policy=QuantPolicy.NONE vs quant_policy=QuantPolicy.TURBO_QUANT)
# =============================================================================

class TestQuantPolicy42Accuracy:
    """Accuracy tests comparing quant_policy=QuantPolicy.TURBO_QUANT against
    non-quantized baseline.

    These tests verify the numerical accuracy/precision of quant_policy=QuantPolicy.TURBO_QUANT
    (K=4bit, V=2bit mixed precision) by comparing against quant_policy=QuantPolicy.NONE.

    Error thresholds are relaxed due to aggressive quantization:
    - MAE < 0.1 on logits
    - Max AE < 0.5 on logits
    """

    def test_logits_accuracy(self, pipe_no_quant, pipe_quant_42):
        """Test logits accuracy by comparing output logits.

        Compares logits between quantized and non-quantized models.
        Uses deterministic generation settings for reproducibility.

        Thresholds:
        - Mean absolute error (MAE) < 0.1
        - Max absolute error < 0.5
        """
        gen_config = GenerationConfig(
            max_new_tokens=0,  # Required for logits output
            temperature=0.0,
            top_p=1.0,
            top_k=1,
            output_logits='all',
        )

        prompt = 'The capital of France is'

        response_no_quant = pipe_no_quant.infer(prompt, gen_config=gen_config)
        response_quant = pipe_quant_42.infer(prompt, gen_config=gen_config)

        assert isinstance(response_no_quant, Response)
        assert isinstance(response_quant, Response)

        if response_no_quant.logits is not None and response_quant.logits is not None:
            logits_no_quant = response_no_quant.logits
            logits_quant = response_quant.logits

            assert logits_no_quant.shape == logits_quant.shape, \
                f'Logits shape mismatch: {logits_no_quant.shape} vs {logits_quant.shape}'

            abs_error = (logits_no_quant - logits_quant).abs()
            mean_abs_error = abs_error.mean().item()
            max_abs_error = abs_error.max().item()

            print('\nLogits accuracy metrics:')
            print(f'  Mean absolute error: {mean_abs_error:.6f}')
            print(f'  Max absolute error: {max_abs_error:.6f}')

            assert mean_abs_error < 0.1, \
                f'Mean absolute error {mean_abs_error:.6f} exceeds threshold 0.1'
            assert max_abs_error < 0.5, \
                f'Max absolute error {max_abs_error:.6f} exceeds threshold 0.5'
        else:
            pytest.skip('Logits not available for comparison')

    def test_token_accuracy(self, pipe_no_quant, pipe_quant_42):
        """Test token-level accuracy by comparing output token IDs.

        Checks that both models generate output and compares token match rate.
        Note: With aggressive quantization (K=4bit, V=2bit), token match rate
        can be low - this is expected behavior.
        """
        gen_config = GenerationConfig(
            max_new_tokens=20,
            temperature=0.0,
            top_p=1.0,
            top_k=1,
        )

        prompt = 'Hello, how are you?'

        response_no_quant = pipe_no_quant.infer(prompt, gen_config=gen_config)
        response_quant = pipe_quant_42.infer(prompt, gen_config=gen_config)

        assert isinstance(response_no_quant, Response)
        assert isinstance(response_quant, Response)

        tokens_no_quant = response_no_quant.token_ids
        tokens_quant = response_quant.token_ids

        min_len = min(len(tokens_no_quant), len(tokens_quant))
        if min_len > 0:
            matching_tokens = sum(1 for i in range(min_len)
                                  if tokens_no_quant[i] == tokens_quant[i])
            match_rate = matching_tokens / min_len

            print('\nToken accuracy metrics:')
            print(f'  Baseline tokens: {len(tokens_no_quant)}')
            print(f'  Quantized tokens: {len(tokens_quant)}')
            print(f'  Matching tokens: {matching_tokens}/{min_len}')
            print(f'  Match rate: {match_rate:.2%}')

            # Basic sanity check - both models should produce output
            assert len(tokens_no_quant) > 0, 'Baseline produced no tokens'
            assert len(tokens_quant) > 0, 'Quantized model produced no tokens'
        else:
            pytest.skip('No tokens generated for comparison')

    def test_text_quality(self, pipe_no_quant, pipe_quant_42):
        """Test that quantized output is still meaningful text.

        Verifies the quantized model produces coherent text output, even if not exactly matching the non-quantized
        baseline.
        """
        gen_config = GenerationConfig(
            max_new_tokens=30,
            temperature=0.7,
            top_p=0.9,
        )

        prompt = 'Write a short story about a robot.'

        response_no_quant = pipe_no_quant.infer(prompt, gen_config=gen_config)
        response_quant = pipe_quant_42.infer(prompt, gen_config=gen_config)

        assert isinstance(response_no_quant, Response)
        assert isinstance(response_quant, Response)

        assert len(response_no_quant.text.strip()) > 0, 'Baseline output is empty'
        assert len(response_quant.text.strip()) > 0, 'Quantized output is empty'

        print('\nText quality metrics:')
        print(f'  Baseline text length: {len(response_no_quant.text)}')
        print(f'  Quantized text length: {len(response_quant.text)}')

    def test_logprobs_sanity(self, pipe_no_quant, pipe_quant_42):
        """Test that logprobs are reasonable when available."""
        gen_config = GenerationConfig(
            max_new_tokens=10,
            temperature=0.0,
            top_p=1.0,
            top_k=1,
            logprobs=1,
        )

        prompt = 'What is 2+2?'

        response_no_quant = pipe_no_quant.infer(prompt, gen_config=gen_config)
        response_quant = pipe_quant_42.infer(prompt, gen_config=gen_config)

        assert isinstance(response_no_quant, Response)
        assert isinstance(response_quant, Response)

        if response_no_quant.logprobs is not None and response_quant.logprobs is not None:
            print('\nLogprobs available for both models')
            assert isinstance(response_no_quant.logprobs, list)
            assert isinstance(response_quant.logprobs, list)
        else:
            print('\nLogprobs not available (this is expected for some configurations)')