neuron_model_runner.py 18.6 KB
Newer Older
1
2
# SPDX-License-Identifier: Apache-2.0

3
import os
4
from dataclasses import dataclass
5
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Tuple, Union
6
7

import torch
8
from torch import nn
9

10
from vllm.config import DeviceConfig, VllmConfig
11
from vllm.logger import init_logger
12
13
from vllm.lora.layers import LoRAMapping
from vllm.lora.request import LoRARequest
14
from vllm.model_executor import SamplingMetadata
15
from vllm.model_executor.layers.sampler import SamplerOutput
16
from vllm.model_executor.model_loader.neuron import get_neuron_model
17
18
19
20
from vllm.multimodal import (MULTIMODAL_REGISTRY, BatchedTensorInputs,
                             MultiModalKwargs)
from vllm.platforms import current_platform
from vllm.sampling_params import SamplingParams
21
from vllm.sequence import IntermediateTensors, SequenceGroupMetadata
22
from vllm.utils import is_pin_memory_available, make_tensor_with_pad
23
24
25
26
from vllm.worker.model_runner_base import ModelRunnerBase, ModelRunnerInputBase

if TYPE_CHECKING:
    from vllm.attention.backends.abstract import AttentionBackend
27
28
29
30

logger = init_logger(__name__)


31
32
33
34
35
36
37
38
@dataclass(frozen=True)
class ModelInputForNeuron(ModelRunnerInputBase):
    """
    Used by the NeuronModelRunner.
    """
    input_tokens: Optional[torch.Tensor] = None
    input_positions: Optional[torch.Tensor] = None
    input_block_ids: Optional[torch.Tensor] = None
39
40
    sampling_metadata: SamplingMetadata = None
    multi_modal_kwargs: BatchedTensorInputs = None
41
    adapter_ids: Optional[str] = None
42
43
44

    def as_broadcastable_tensor_dict(
            self) -> Dict[str, Union[int, torch.Tensor]]:
45
46
47
48
49
50
51
        return {
            "input_tokens": self.input_tokens,
            "input_positions": self.input_positions,
            "input_block_ids": self.input_block_ids,
            "sampling_metadata": self.sampling_metadata,
            "multi_modal_kwargs": self.multi_modal_kwargs,
        }
52
53
54
55
56
57
58

    @classmethod
    def from_broadcasted_tensor_dict(
        cls,
        tensor_dict: Dict[str, Any],
        attn_backend: Optional["AttentionBackend"] = None,
    ) -> "ModelInputForNeuron":
59
60
61
62
63
64
65
        return ModelInputForNeuron(
            input_tokens=tensor_dict["input_tokens"],
            input_positions=tensor_dict["input_positions"],
            input_block_ids=tensor_dict["input_block_ids"],
            sampling_metadata=tensor_dict["sampling_metadata"],
            multi_modal_kwargs=tensor_dict["multi_modal_kwargs"],
        )
66
67
68


class NeuronModelRunner(ModelRunnerBase[ModelInputForNeuron]):
69
    """A model runner for AWS Neuron hardware"""
70

71
72
73
    # NEURON has an upper limit on the top_k
    _MAX_NEURON_SAMPLING_TOP_K = 256

74
75
    def __init__(
        self,
76
        vllm_config: VllmConfig,
77
    ):
78
        ModelRunnerBase.__init__(self, vllm_config)
79
80
81

        if (self.model_config is not None
                and self.model_config.get_sliding_window()):
82
83
            logger.warning("Sliding window is not supported on Neuron. "
                           "The model will run without sliding window.")
84
85
        self.device_config = (self.device_config if self.device_config
                              is not None else DeviceConfig())
86
        self.lora_config = vllm_config.lora_config
87
88
89
        self.device = self.device_config.device
        self.pin_memory = is_pin_memory_available()

90
91
92
93
        # Multi-modal data support
        self.multi_modal_input_mapper = MULTIMODAL_REGISTRY \
            .create_input_mapper(self.model_config)

94
95
96
        # Lazy initialization.
        self.model: nn.Module  # initialize after load_model.

97
98
99
100
101
102
103
104
105
106
107
        # Once NEURON_ON_DEVICE_SAMPLING_DISABLED is set to a non-zero value,
        # turn off on-device sampling.
        self._on_device_sampling_disabled = int(
            os.getenv("NEURON_ON_DEVICE_SAMPLING_DISABLED", "0"))

        # NEURON needs to update sampling parameters when request IDs change
        # across batches. This variable stores the previous batch's request IDs
        # to determine if an update is needed.
        self._previous_batch_request_ids: List[str] = []

        if not self._on_device_sampling_disabled:
108
            self._init_neuron_sampling()
109

110
111
112
    def _init_neuron_sampling(self) -> None:
        if current_platform.use_transformers_neuronx():
            from transformers_neuronx.config import GenerationConfig
113
        else:
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
            from transformers import GenerationConfig
        logger.warning(
            "On-device sampling is turned on in Neuron by default, only "
            "top_k, top_p, and temperature are current supported sampling "
            "parameters. To turn off the on-device sampling, please set "
            "the environment variable NEURON_ON_DEVICE_SAMPLING_DISABLED=1.")
        self.model_config.neuron_sampling_params = GenerationConfig(
            max_length=self.scheduler_config.max_model_len,
            do_sample=True,
            per_batch_line=True,
            top_k=[self._MAX_NEURON_SAMPLING_TOP_K] \
                  * self.scheduler_config.max_num_seqs,
            top_p=[1.0] * self.scheduler_config.max_num_seqs,
            temperature=[1.0] * self.scheduler_config.max_num_seqs,
            dynamic=True,
            global_top_k=self._MAX_NEURON_SAMPLING_TOP_K)

    def load_model(self) -> None:
        self.model = get_neuron_model(self.model_config,
                                      parallel_config=self.parallel_config,
                                      scheduler_config=self.scheduler_config)
135

136
137
138
    def get_model(self) -> nn.Module:
        return self.model

139
140
141
    def _prepare_prompt(
        self,
        seq_group_metadata_list: List[SequenceGroupMetadata],
142
143
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, List[int],
               BatchedTensorInputs]:
144
145
146
147
148
        assert len(seq_group_metadata_list) > 0
        input_tokens: List[List[int]] = []
        input_positions: List[List[int]] = []
        input_block_ids: List[int] = []

149
        seq_lens: List[int] = []
150
        multi_modal_kwargs_list: List[MultiModalKwargs] = []
151
152
153
154
155
156
157
158
        for seq_group_metadata in seq_group_metadata_list:
            assert seq_group_metadata.is_prompt
            seq_ids = list(seq_group_metadata.seq_data.keys())
            assert len(seq_ids) == 1
            seq_id = seq_ids[0]

            seq_data = seq_group_metadata.seq_data[seq_id]
            prompt_tokens = seq_data.get_token_ids()
159
160
            seq_len = len(prompt_tokens)
            seq_lens.append(seq_len)
161
162

            input_tokens.append(prompt_tokens)
163
            input_positions.append(list(range(seq_len)))
164
165
166
167
168
169

            assert seq_group_metadata.block_tables is not None
            block_table = seq_group_metadata.block_tables[seq_id]
            assert len(block_table) == 1
            input_block_ids.append(block_table[0])

170
171
            mm_kwargs = seq_group_metadata.multi_modal_data
            if mm_kwargs:
172
                multi_modal_kwargs_list.append(mm_kwargs)
173

174
175
        max_seq_len = max(seq_lens)
        assert max_seq_len > 0
176
177
        input_tokens = make_tensor_with_pad(input_tokens,
                                            pad=0,
178
                                            max_len=max_seq_len,
179
180
181
182
                                            dtype=torch.long,
                                            device=self.device)
        input_positions = make_tensor_with_pad(input_positions,
                                               pad=0,
183
                                               max_len=max_seq_len,
184
185
186
187
188
189
                                               dtype=torch.long,
                                               device=self.device)
        input_block_ids = torch.tensor(input_block_ids,
                                       dtype=torch.long,
                                       device=self.device)

190
        multi_modal_kwargs = MultiModalKwargs.batch(multi_modal_kwargs_list)
191
192
193

        return (input_tokens, input_positions, input_block_ids, seq_lens,
                multi_modal_kwargs)
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

    def _prepare_decode(
        self,
        seq_group_metadata_list: List[SequenceGroupMetadata],
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        assert len(seq_group_metadata_list) > 0
        input_tokens: List[List[int]] = []
        input_positions: List[List[int]] = []
        input_block_ids: List[int] = []
        context_lens: List[int] = []

        for seq_group_metadata in seq_group_metadata_list:
            assert not seq_group_metadata.is_prompt

            seq_ids = list(seq_group_metadata.seq_data.keys())

            for seq_id in seq_ids:
                seq_data = seq_group_metadata.seq_data[seq_id]
                generation_token = seq_data.get_last_token_id()
                input_tokens.append([generation_token])

                seq_len = seq_data.get_len()
                position = seq_len - 1
                input_positions.append([position])
                context_lens.append(seq_len)

                assert seq_group_metadata.block_tables is not None
                block_table = seq_group_metadata.block_tables[seq_id]
                assert len(block_table) == 1
                input_block_ids.append(block_table[0])

        input_tokens = make_tensor_with_pad(input_tokens,
                                            pad=0,
227
                                            max_len=1,
228
229
230
231
                                            dtype=torch.long,
                                            device=self.device)
        input_positions = make_tensor_with_pad(input_positions,
                                               pad=0,
232
                                               max_len=1,
233
234
235
236
237
238
239
240
241
242
243
                                               dtype=torch.long,
                                               device=self.device)
        context_lens = torch.tensor(context_lens,
                                    dtype=torch.int,
                                    device=self.device)
        input_block_ids = torch.tensor(input_block_ids,
                                       dtype=torch.long,
                                       device=self.device)

        return input_tokens, input_positions, input_block_ids

244
245
246
247
248
    def make_model_input_from_broadcasted_tensor_dict(
            self, tensor_dict: Dict[str, Any]) -> ModelInputForNeuron:
        return ModelInputForNeuron.from_broadcasted_tensor_dict(tensor_dict)

    def prepare_model_input(
249
        self,
250
        seq_group_metadata_list: List[SequenceGroupMetadata],
251
        virtual_engine: int = 0,
Mor Zusman's avatar
Mor Zusman committed
252
        finished_requests_ids: Optional[List[str]] = None
253
    ) -> ModelInputForNeuron:
254
        multi_modal_kwargs = None
255
256
257
258
259
        # NOTE: We assume that all sequences in the group are all prompts or
        # all decodes.
        is_prompt = seq_group_metadata_list[0].is_prompt
        # Prepare input tensors.
        if is_prompt:
260
261
262
            (input_tokens, input_positions, input_block_ids, seq_lens,
             multi_modal_kwargs
             ) = self._prepare_prompt(seq_group_metadata_list)
263
264
265
        else:
            (input_tokens, input_positions,
             input_block_ids) = self._prepare_decode(seq_group_metadata_list)
266
            seq_lens = None
267
268
269
270
271
272
273
274
275
276

        if not self._on_device_sampling_disabled:
            for seq_group_metadata in seq_group_metadata_list:
                sampling_params = seq_group_metadata.sampling_params
                top_k, top_p, temperature = (
                    self._convert_to_neuron_sampling_params(sampling_params))
                sampling_params.top_k = top_k
                sampling_params.top_p = top_p
                sampling_params.temperature = temperature

277
278
        sampling_metadata = SamplingMetadata.prepare(
            seq_group_metadata_list,
279
280
            seq_lens,
            # query_lens is not needed if chunked prefill is not
281
            # supported. Since neuron worker doesn't support chunked prefill
282
283
            # just use seq_lens instead.
            seq_lens,
284
            self.device,
285
286
            self.pin_memory,
            generators=self.get_generators(finished_requests_ids))
287

288
289
        if current_platform.use_transformers_neuronx(
        ) and not self._on_device_sampling_disabled:
290
291
292
293
294
295
296
            # Once the request IDs are changed in current iteration, we will
            # update the on-device sampling parameters.
            current_batch_request_ids = [
                seq_group_meta_data.request_id
                for seq_group_meta_data in seq_group_metadata_list
            ]
            if current_batch_request_ids != self._previous_batch_request_ids:
297
                self._update_neuron_sampling_params(seq_group_metadata_list)
298
299
                self._previous_batch_request_ids = current_batch_request_ids

300
301
302
        return ModelInputForNeuron(input_tokens=input_tokens,
                                   input_positions=input_positions,
                                   input_block_ids=input_block_ids,
303
304
                                   sampling_metadata=sampling_metadata,
                                   multi_modal_kwargs=multi_modal_kwargs)
305

306
307
    def _update_neuron_sampling_params(
            self, seq_group_metadata_list: List[SequenceGroupMetadata]):
308
309
310
311
312
313
        # Update Neuron sampling parameters (GenerationConfig in Neuron)
        current_sampling_params = self.model_config.neuron_sampling_params
        assert current_sampling_params is not None, (
            f"Failed to update sampling_params, "
            f"current sampling params is {current_sampling_params}")

314
315
        is_update_needed = False

316
317
318
319
        top_k = current_sampling_params.top_k
        top_p = current_sampling_params.top_p
        temperature = current_sampling_params.temperature

320
321
322
323
324
        # The index of a sequence's sampling parameters in neuron is equal to
        # its index in `input_block_ids`.
        for seq_group_metadata in seq_group_metadata_list:
            seq_ids = list(seq_group_metadata.seq_data.keys())
            sampling_params = seq_group_metadata.sampling_params
325

326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
            seq_group_top_k = sampling_params.top_k
            seq_group_top_p = sampling_params.top_p
            seq_group_temperature = sampling_params.temperature

            for seq_id in seq_ids:
                index = seq_group_metadata.block_tables[seq_id][0]
                if (top_k[index] != seq_group_top_k
                        or top_p[index] != seq_group_top_p
                        or temperature[index] != seq_group_temperature):
                    is_update_needed = True

                top_k[index] = seq_group_top_k
                top_p[index] = seq_group_top_p
                temperature[index] = seq_group_temperature

        # update_generation_config is only available in transformers-neuronx
        if is_update_needed and current_platform.use_transformers_neuronx():
            self.model.model.update_generation_config(current_sampling_params)

    def _convert_to_neuron_sampling_params(
            self, sampling_params: SamplingParams) -> Tuple[int, float, float]:
        # Returns the top_k, top_p and temperature parameters for neuron.
        top_k = sampling_params.top_k
        top_p = sampling_params.top_p
        temperature = sampling_params.temperature

        if temperature == 0.0:
            # Enable greedy sampling on zero temperature
            return (1, 1.0, 1.0)
355
        if top_k < 1 or top_k > self._MAX_NEURON_SAMPLING_TOP_K:
356
357
358
            top_k = self._MAX_NEURON_SAMPLING_TOP_K

        return (top_k, top_p, temperature)
359

360
361
362
    @torch.inference_mode()
    def execute_model(
        self,
363
364
        model_input: ModelInputForNeuron,
        kv_caches: Optional[List[torch.Tensor]] = None,
365
        intermediate_tensors: Optional[IntermediateTensors] = None,
366
367
368
369
370
371
        num_steps: int = 1,
    ) -> Optional[List[SamplerOutput]]:
        if num_steps > 1:
            raise ValueError(
                "NeuronModelRunner does not support multi-step execution.")

372
373
374
375
376
377
378
379
380
381
382
383
384
        # extract top_k, top_p and temperature from model_input for neuron
        # forward call
        sampling_params = (torch.tensor([[
            seq_group.sampling_params.top_k, seq_group.sampling_params.top_p,
            seq_group.sampling_params.temperature
        ] for seq_group in model_input.sampling_metadata.seq_groups]))

        if current_platform.use_neuronx_distributed():
            hidden_states = self.model(
                input_ids=model_input.input_tokens,
                positions=model_input.input_positions,
                input_block_ids=model_input.input_block_ids,
                sampling_params=sampling_params,
385
                adapter_ids=model_input.adapter_ids,
386
387
388
389
390
                **MultiModalKwargs.as_kwargs(
                    model_input.multi_modal_kwargs or {},
                    dtype=self.model_config.dtype,
                    device=self.device,
                ),
391
392
393
394
            )
        elif current_platform.use_transformers_neuronx():
            # [TODO] validate on-device sampling
            # The model signature may need change for on-device sampling
395
396
397
398
            hidden_states = self.model(
                input_ids=model_input.input_tokens,
                positions=model_input.input_positions,
                input_block_ids=model_input.input_block_ids,
399
400
401
402
403
                **MultiModalKwargs.as_kwargs(
                    model_input.multi_modal_kwargs or {},
                    dtype=self.model_config.dtype,
                    device=self.device,
                ),
404
            )
405

406
407
408
409
410
411
412
        # Compute the logits only if the on-device sampling is turned off as
        # on-device sampling outputs the token ids.
        if self._on_device_sampling_disabled:
            logits = self.model.compute_logits(hidden_states,
                                               model_input.sampling_metadata)
        else:
            logits = hidden_states
413
414
415
416

        # Sample the next token.
        output = self.model.sample(
            logits=logits,
417
            sampling_metadata=model_input.sampling_metadata,
418
        )
419
        return [output]
420
421
422
423

    @property
    def vocab_size(self) -> int:
        return self.model_config.get_vocab_size()
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448

    def remove_all_loras(self):
        raise NotImplementedError(
            "LoRAs are not supported for Transformers NeuronX framework")

    def set_active_loras(self, lora_requests: Set[LoRARequest],
                         lora_mapping: LoRAMapping) -> None:
        raise NotImplementedError(
            "LoRAs are not supported for Transformers NeuronX framework")

    def add_lora(self, lora_request: LoRARequest):
        raise NotImplementedError(
            "LoRAs are not supported for Transformers NeuronX framework")

    def remove_lora(self, lora_id: int) -> bool:
        raise NotImplementedError(
            "LoRAs are not supported for Transformers NeuronX framework")

    def pin_lora(self, lora_id: int) -> bool:
        raise NotImplementedError(
            "LoRAs are not supported for Transformers NeuronX framework")

    def list_loras(self) -> Set[int]:
        raise NotImplementedError(
            "LoRAs are not supported for Transformers NeuronX framework")