test_sampler.py 12.7 KB
Newer Older
1
import random
2
from typing import List, Optional, Tuple
3
4
from unittest.mock import patch

Woosuk Kwon's avatar
Woosuk Kwon committed
5
import pytest
6
import torch
7
from transformers import GenerationConfig, GenerationMixin
8
9
10
11

from vllm.model_executor.layers.sampler import Sampler
from vllm.model_executor.utils import set_random_seed
from vllm.sequence import SamplingParams, SequenceData, SequenceGroupMetadata
Woosuk Kwon's avatar
Woosuk Kwon committed
12
from vllm.worker.model_runner import ModelRunner
13
14
15
16


class MockLogitsSampler(Sampler):

17
18
    def __init__(self, fake_logits: torch.Tensor):
        super().__init__()
19
20
21
        self.fake_logits = fake_logits

    def forward(self, *args, **kwargs):
22
        return super().forward(*args, **kwargs)
23
24
25
26


def _prepare_test(
    batch_size: int
Woosuk Kwon's avatar
Woosuk Kwon committed
27
) -> Tuple[torch.Tensor, torch.Tensor, MockLogitsSampler, ModelRunner]:
28
    vocab_size = 32000
29
    input_tensor = torch.rand((batch_size, 1024), dtype=torch.float16)
30
31
32
    fake_logits = torch.full((batch_size, vocab_size),
                             1e-2,
                             dtype=input_tensor.dtype)
33
    sampler = MockLogitsSampler(fake_logits)
34
    model_runner = ModelRunner(None, None, None, None, None)
Woosuk Kwon's avatar
Woosuk Kwon committed
35
    return input_tensor, fake_logits, sampler, model_runner
36
37
38


RANDOM_SEEDS = list(range(128))
39
40
41
CUDA_DEVICES = [
    f"cuda:{i}" for i in range(1 if torch.cuda.device_count() == 1 else 2)
]
42
43


Nick Hill's avatar
Nick Hill committed
44
45
46
47
48
49
50
def _do_sample(
    batch_size: int,
    input_tensor: torch.Tensor,
    sampler: MockLogitsSampler,
    model_runner: ModelRunner,
    sampling_params: SamplingParams,
):
51
    seq_group_metadata_list = []
Woosuk Kwon's avatar
Woosuk Kwon committed
52
    prompt_lens = []
53
54
55
56
57
58
    for i in range(batch_size):
        seq_group_metadata_list.append(
            SequenceGroupMetadata(
                request_id=f"test_{i}",
                is_prompt=True,
                seq_data={0: SequenceData([1, 2, 3])},
Nick Hill's avatar
Nick Hill committed
59
                sampling_params=sampling_params,
60
61
                block_tables={0: [1]},
            ))
Woosuk Kwon's avatar
Woosuk Kwon committed
62
        prompt_lens.append(seq_group_metadata_list[-1].seq_data[0].get_len())
63

Woosuk Kwon's avatar
Woosuk Kwon committed
64
    sampling_metadata = model_runner._prepare_sample(seq_group_metadata_list,
65
66
                                                     prompt_lens,
                                                     subquery_lens=prompt_lens)
67
    return sampler(logits=input_tensor, sampling_metadata=sampling_metadata)
Nick Hill's avatar
Nick Hill committed
68
69
70
71
72
73
74
75
76
77
78
79


@pytest.mark.parametrize("seed", RANDOM_SEEDS)
@pytest.mark.parametrize("device", CUDA_DEVICES)
def test_sampler_all_greedy(seed: int, device: str):
    set_random_seed(seed)
    torch.set_default_device(device)
    batch_size = random.randint(1, 256)
    input_tensor, fake_logits, sampler, model_runner = _prepare_test(
        batch_size)

    sampling_params = SamplingParams(temperature=0)
80
81
    sampler_output = _do_sample(batch_size, fake_logits, sampler, model_runner,
                                sampling_params)
82
83
    expected = torch.argmax(fake_logits, dim=-1)
    for i, sequence_output in enumerate(sampler_output):
Woosuk Kwon's avatar
Woosuk Kwon committed
84
        for nth_output in sequence_output.samples:
85
86
            assert nth_output.output_token == expected[i].item()

Simon Mo's avatar
Simon Mo committed
87
88
    del model_runner

89
90

@pytest.mark.parametrize("seed", RANDOM_SEEDS)
91
92
@pytest.mark.parametrize("device", CUDA_DEVICES)
def test_sampler_all_random(seed: int, device: str):
93
    set_random_seed(seed)
94
    torch.set_default_device(device)
95
    batch_size = random.randint(1, 256)
Woosuk Kwon's avatar
Woosuk Kwon committed
96
97
    input_tensor, fake_logits, sampler, model_runner = _prepare_test(
        batch_size)
98
99
100
101

    for i in range(batch_size):
        fake_logits[i, i] = 1e2

Nick Hill's avatar
Nick Hill committed
102
103
104
105
    sampling_params = SamplingParams(
        temperature=1.0,
        n=random.randint(1, 10),
    )
106
107
    sampler_output = _do_sample(batch_size, fake_logits, sampler, model_runner,
                                sampling_params)
Nick Hill's avatar
Nick Hill committed
108
109
110
111
112
113
114
115
116
117
118
119
120
121

    for i, sequence_output in enumerate(sampler_output):
        for nth_output in sequence_output.samples:
            assert nth_output.output_token == i

    del model_runner


@pytest.mark.parametrize("seed", RANDOM_SEEDS)
@pytest.mark.parametrize("device", CUDA_DEVICES)
def test_sampler_all_random_seed(seed: int, device: str):
    set_random_seed(seed)
    torch.set_default_device(device)
    batch_size = random.randint(1, 256)
122
    _, fake_logits, sampler, model_runner = _prepare_test(batch_size)
Nick Hill's avatar
Nick Hill committed
123

124
    for i in range(batch_size):
Nick Hill's avatar
Nick Hill committed
125
126
127
128
129
130
131
        fake_logits[i, i] = 1e2

    sampling_params = SamplingParams(
        temperature=1.0,
        n=random.randint(1, 10),
        seed=random.randint(0, 10000),
    )
132
133
    sampler_output = _do_sample(batch_size, fake_logits, sampler, model_runner,
                                sampling_params)
134
135

    for i, sequence_output in enumerate(sampler_output):
Woosuk Kwon's avatar
Woosuk Kwon committed
136
        for nth_output in sequence_output.samples:
137
138
            assert nth_output.output_token == i

Simon Mo's avatar
Simon Mo committed
139
140
    del model_runner

141

Nick Hill's avatar
Nick Hill committed
142
143
144
145
146
147
@pytest.mark.parametrize("seed", RANDOM_SEEDS)
@pytest.mark.parametrize("device", CUDA_DEVICES)
def test_sampler_all_random_seed_deterministic(seed: int, device: str):
    set_random_seed(seed)
    torch.set_default_device(device)
    batch_size = random.randint(1, 256)
148
    _, fake_logits, sampler, model_runner = _prepare_test(batch_size)
Nick Hill's avatar
Nick Hill committed
149
150
151
152
153
154

    sampling_params = SamplingParams(
        temperature=1.0,
        n=random.randint(1, 10),
        seed=random.randint(0, 10000),
    )
155
    first_sampler_output = _do_sample(batch_size, fake_logits, sampler,
Nick Hill's avatar
Nick Hill committed
156
157
                                      model_runner, sampling_params)

158
    second_sampler_output = _do_sample(batch_size, fake_logits, sampler,
Nick Hill's avatar
Nick Hill committed
159
160
161
162
163
164
165
                                       model_runner, sampling_params)

    assert first_sampler_output == second_sampler_output

    del model_runner


166
@pytest.mark.parametrize("seed", RANDOM_SEEDS)
167
168
@pytest.mark.parametrize("device", CUDA_DEVICES)
def test_sampler_all_beam(seed: int, device: str):
169
    set_random_seed(seed)
170
    torch.set_default_device(device)
171
    batch_size = random.randint(1, 256)
172
    _, fake_logits, sampler, model_runner = _prepare_test(batch_size)
173

Nick Hill's avatar
Nick Hill committed
174
175
176
177
178
    sampling_params = SamplingParams(
        temperature=0,
        best_of=2,
        use_beam_search=True,
    )
179
    _do_sample(batch_size, fake_logits, sampler, model_runner, sampling_params)
180
181
182
183
    # no assertion here as I am not sure how to determine whether
    # the outputs are expected - in other words, this just tests
    # whether there are no exceptions in the sampler
    # when handling an all-beam search case.
Simon Mo's avatar
Simon Mo committed
184
    del model_runner
185
186
187


@pytest.mark.parametrize("seed", RANDOM_SEEDS)
188
189
@pytest.mark.parametrize("device", CUDA_DEVICES)
def test_sampler_mixed(seed: int, device: str):
190
    set_random_seed(seed)
191
    torch.set_default_device(device)
192
    batch_size = random.randint(1, 256)
Woosuk Kwon's avatar
Woosuk Kwon committed
193
194
    input_tensor, fake_logits, sampler, model_runner = _prepare_test(
        batch_size)
195
196

    seq_group_metadata_list = []
Nick Hill's avatar
Nick Hill committed
197
    expected_tokens: List[Optional[List[int]]] = []
Woosuk Kwon's avatar
Woosuk Kwon committed
198
    prompt_lens = []
199
    for i in range(batch_size):
Nick Hill's avatar
Nick Hill committed
200
201
        expected: Optional[List[int]] = None
        sampling_type = random.randint(0, 3)
202
203
        if sampling_type == 0:
            sampling_params = SamplingParams(temperature=0)
Nick Hill's avatar
Nick Hill committed
204
205
            expected = [torch.argmax(fake_logits[i], dim=-1).item()]
        elif sampling_type in (1, 2):
206
207
208
209
210
211
212
213
            n = random.randint(1, 10)
            sampling_params = SamplingParams(
                temperature=random.random() + 0.1,
                top_p=min(random.random() + 0.1, 1),
                top_k=random.randint(0, 10) or -1,
                n=n,
                presence_penalty=random.randint(0, 1),
            )
Nick Hill's avatar
Nick Hill committed
214
215
216
217
218
219
            if sampling_type == 2:
                sampling_params.seed = random.randint(0, 10000)
            else:
                for idx in range(n):
                    fake_logits[i, i + idx] = 1e2
                expected = list(range(i, i + n))
220
221
222
223
        else:
            sampling_params = SamplingParams(temperature=0,
                                             use_beam_search=True,
                                             best_of=2)
Nick Hill's avatar
Nick Hill committed
224
        expected_tokens.append(expected)
225
226
227
228
229
230
231
232
        seq_group_metadata_list.append(
            SequenceGroupMetadata(
                request_id=f"test_{i}",
                is_prompt=True,
                seq_data={0: SequenceData([1, 2, 3])},
                sampling_params=sampling_params,
                block_tables={0: [1]},
            ))
Woosuk Kwon's avatar
Woosuk Kwon committed
233
        prompt_lens.append(seq_group_metadata_list[-1].seq_data[0].get_len())
234

Nick Hill's avatar
Nick Hill committed
235
236
237
    def test_sampling(model_runner: ModelRunner):
        sampling_metadata = model_runner._prepare_sample(
            seq_group_metadata_list, prompt_lens, subquery_lens=prompt_lens)
238
        sampler_output = sampler(logits=fake_logits,
Nick Hill's avatar
Nick Hill committed
239
240
241
242
243
244
245
                                 sampling_metadata=sampling_metadata)

        for i, (sequence_output, metadata) in enumerate(
                zip(sampler_output, seq_group_metadata_list)):
            if metadata.sampling_params.use_beam_search:
                continue

246
247
248
249
            if (metadata.sampling_params.seed is not None
                    and expected_tokens[i] is None):
                # Record seeded random result to compare with results of
                # second invocation
Nick Hill's avatar
Nick Hill committed
250
251
252
253
254
255
256
                expected_tokens[i] = [
                    nth_output.output_token
                    for nth_output in sequence_output.samples
                ]
                continue

            for n, nth_output in enumerate(sequence_output.samples):
257
258
                if (metadata.sampling_params.temperature == 0
                        or metadata.sampling_params.seed is not None):
Nick Hill's avatar
Nick Hill committed
259
260
261
                    # Ensure exact matches for greedy or random with seed
                    assert nth_output.output_token == expected_tokens[i][n]
                else:
262
263
                    # For non-seeded random check that one of the high-logit
                    # tokens were chosen
Nick Hill's avatar
Nick Hill committed
264
265
266
267
268
269
270
271
272
273
274
275
276
277
                    assert nth_output.output_token in expected_tokens[i]

    # Test batch
    test_sampling(model_runner)

    # Shuffle the batch and resample
    target_index = list(range(batch_size))
    for list_to_shuffle in (target_index, seq_group_metadata_list,
                            expected_tokens, prompt_lens):
        random.Random(seed).shuffle(list_to_shuffle)
    target_index = torch.tensor(target_index)
    input_tensor.data = input_tensor.index_select(0, target_index)
    fake_logits.data = fake_logits.index_select(0, target_index)

278
279
    # This time, results of seeded random samples will be compared with
    # the corresponding sample in the pre-shuffled batch
Nick Hill's avatar
Nick Hill committed
280
    test_sampling(model_runner)
281

Simon Mo's avatar
Simon Mo committed
282
283
    del model_runner

284

285
@pytest.mark.parametrize("seed", RANDOM_SEEDS)
286
287
@pytest.mark.parametrize("device", CUDA_DEVICES)
def test_sampler_top_k_top_p(seed: int, device: str):
288
289
290
291
292
293
    set_random_seed(seed)
    batch_size = random.randint(1, 256)
    top_k = random.randint(100, 500)
    top_p = random.random() * 0.1
    vocab_size = 32000
    input_tensor = torch.rand((batch_size, 1024),
294
                              device=device,
295
296
297
298
299
300
                              dtype=torch.float16)
    fake_logits = torch.normal(0,
                               5,
                               size=(batch_size, vocab_size),
                               device=input_tensor.device,
                               dtype=input_tensor.dtype)
301
    sampler = MockLogitsSampler(fake_logits)
302
    model_runner = ModelRunner(None, None, None, None, None)
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328

    generation_model = GenerationMixin()
    generation_config = GenerationConfig(top_k=top_k,
                                         top_p=top_p,
                                         do_sample=True)
    warpers = generation_model._get_logits_warper(generation_config)
    assert len(warpers) == 2  # top_p and top_k

    seq_group_metadata_list = []
    prompt_lens = []
    for i in range(batch_size):
        seq_group_metadata_list.append(
            SequenceGroupMetadata(
                request_id=f"test_{i}",
                is_prompt=True,
                seq_data={0: SequenceData([1, 2, 3])},
                sampling_params=SamplingParams(
                    temperature=1,
                    top_k=top_k,
                    top_p=top_p,
                ),
                block_tables={0: [1]},
            ))
        prompt_lens.append(seq_group_metadata_list[-1].seq_data[0].get_len())

    sampling_metadata = model_runner._prepare_sample(seq_group_metadata_list,
329
330
                                                     prompt_lens,
                                                     subquery_lens=prompt_lens)
331
332
333

    sample_probs = None

334
    def mock_sample(probs, *args, **kwargs):
335
336
337
338
339
        nonlocal sample_probs
        sample_probs = probs
        return [[prob.topk(1, dim=-1).indices.tolist(), [0]] for prob in probs]

    with patch("vllm.model_executor.layers.sampler._sample", mock_sample):
340
        sampler(logits=fake_logits, sampling_metadata=sampling_metadata)
341
342
343
344
    hf_probs = warpers(torch.zeros_like(fake_logits), fake_logits.clone())
    hf_probs = torch.softmax(hf_probs, dim=-1, dtype=torch.float)
    assert torch.allclose(hf_probs, sample_probs, atol=1e-5)
    assert torch.equal(hf_probs.eq(0), sample_probs.eq(0))
Simon Mo's avatar
Simon Mo committed
345
346

    del model_runner