@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 043e4c4955
Fork 0

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, 11:00 UTC

Add NeuronxDistributedInference support, Speculative Decoding, Dynamic on-device sampling (#16357)

Signed-off-by: Satyajith Chilappagari <satchill@amazon.com> Co-authored-by: Aaron Dou <yzdou@amazon.com> Co-authored-by: Shashwat Srijan <sssrijan@amazon.com> Co-authored-by: Chongming Ni <chongmni@amazon.com> Co-authored-by: Amulya Ballakur <amulyaab@amazon.com> Co-authored-by: Patrick Lange <patlange@amazon.com> Co-authored-by: Elaine Zhao <elaineyz@amazon.com> Co-authored-by: Lin Lin Pan <tailinpa@amazon.com> Co-authored-by: Navyadhara Gogineni <navyadha@amazon.com> Co-authored-by: Yishan McNabb <yishanm@amazon.com> Co-authored-by: Mrinal Shukla <181322398+mrinalks@users.noreply.github.com>

Satyajith Chilappagari committed May 7, 2025 at 00:07 UTC 043e4c495511d51841a66980e8503ee7da943aa1
15 files changed +1623 -102
examples/offline_inference/neuron_eagle.py new
+54
@@ -0,0 +1,54 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +"""
3 +This example shows how to run offline inference with an EAGLE speculative
4 +decoding model on neuron. To use EAGLE speculative decoding, you must use
5 +a draft model that is specifically fine-tuned for EAGLE speculation.
6 +Additionally, to use EAGLE with NxD Inference, the draft model must include
7 +the LM head weights from the target model. These weights are shared between
8 +the draft and target model.
9 +"""
10 +
11 +from vllm import LLM, SamplingParams
12 +
13 +# Sample prompts.
14 +prompts = [
15 + "What is annapurna labs?",
16 +]
17 +
18 +# Create a sampling params object.
19 +sampling_params = SamplingParams(top_k=1, max_tokens=500, ignore_eos=True)
20 +
21 +# Create an LLM.
22 +llm = LLM(
23 + model="/home/ubuntu/model_hf/Meta-Llama-3.1-70B-Instruct",
24 + speculative_config={
25 + "model": "/home/ubuntu/model_hf/Llama-3.1-70B-Instruct-EAGLE-Draft",
26 + "num_speculative_tokens": 5,
27 + "max_model_len": 2048
28 + },
29 + max_num_seqs=4,
30 + # The max_model_len and block_size arguments are required to be same as
31 + # max sequence length when targeting neuron device.
32 + # Currently, this is a known limitation in continuous batching support
33 + # in neuronx-distributed-inference.
34 + max_model_len=2048,
35 + block_size=2048,
36 + # The device can be automatically detected when AWS Neuron SDK is installed.
37 + # The device argument can be either unspecified for automated detection,
38 + # or explicitly assigned.
39 + device="neuron",
40 + tensor_parallel_size=32,
41 + override_neuron_config={
42 + "enable_eagle_speculation": True,
43 + "enable_fused_speculation": True
44 + },
45 +)
46 +
47 +# Generate texts from the prompts. The output is a list of RequestOutput objects
48 +# that contain the prompt, generated text, and other information.
49 +outputs = llm.generate(prompts, sampling_params)
50 +# Print the outputs.
51 +for output in outputs:
52 + prompt = output.prompt
53 + generated_text = output.outputs[0].text
54 + print(f"Prompt: {prompt!r}, \n\n\n\ Generated text: {generated_text!r}")
examples/offline_inference/neuron_speculation.py new
+64
@@ -0,0 +1,64 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +"""
3 +This example shows how to run offline inference with a speculative
4 +decoding model on neuron.
5 +"""
6 +
7 +import os
8 +
9 +from vllm import LLM, SamplingParams
10 +
11 +# Sample prompts.
12 +prompts = [
13 + "Hello, I am a language model and I can help",
14 + "The president of the United States is",
15 + "The capital of France is",
16 +]
17 +
18 +
19 +def config_buckets():
20 + """Configure context length and token gen buckets."""
21 + # creates XLA hlo graphs for all the context length buckets.
22 + os.environ['NEURON_CONTEXT_LENGTH_BUCKETS'] = "128,512,1024,2048"
23 + # creates XLA hlo graphs for all the token gen buckets.
24 + os.environ['NEURON_TOKEN_GEN_BUCKETS'] = "128,512,1024,2048"
25 +
26 +
27 +def initialize_model():
28 + """Create an LLM with speculative decoding."""
29 + return LLM(
30 + model="openlm-research/open_llama_7b",
31 + speculative_config={
32 + "model": "openlm-research/open_llama_3b",
33 + "num_speculative_tokens": 4,
34 + "max_model_len": 2048
35 + },
36 + max_num_seqs=4,
37 + max_model_len=2048,
38 + block_size=2048,
39 + use_v2_block_manager=True,
40 + device="neuron",
41 + tensor_parallel_size=32,
42 + )
43 +
44 +
45 +def process_requests(model: LLM, sampling_params: SamplingParams):
46 + """Generate texts from prompts and print them."""
47 + outputs = model.generate(prompts, sampling_params)
48 + for output in outputs:
49 + prompt = output.prompt
50 + generated_text = output.outputs[0].text
51 + print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
52 +
53 +
54 +def main():
55 + """Main function that sets up the model and processes prompts."""
56 + config_buckets()
57 + model = initialize_model()
58 + # Create a sampling params object.
59 + sampling_params = SamplingParams(max_tokens=100, top_k=1)
60 + process_requests(model, sampling_params)
61 +
62 +
63 +if __name__ == '__main__':
64 + main()
requirements/neuron.txt
+2 -1
@@ -5,4 +5,5 @@
5 packaging>=24.2
6 setuptools>=77.0.3,<80.0.0
7 torch-neuronx >= 2.5.0
8 -neuronx-cc
8 +neuronx-cc>=2.0.0a0
9 +torchvision # Required for Llama3.2 multimodal image preprocessing
tests/neuron/1_core/test_neuron_model_runner.py new
+126
@@ -0,0 +1,126 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +import os
3 +from unittest.mock import MagicMock
4 +
5 +from vllm.config import VllmConfig
6 +from vllm.engine.arg_utils import EngineArgs
7 +from vllm.platforms import current_platform
8 +from vllm.platforms.neuron import NeuronFramework
9 +from vllm.sampling_params import SamplingParams
10 +from vllm.sequence import SequenceData, SequenceGroupMetadata
11 +from vllm.worker.neuron_model_runner import NeuronModelRunner
12 +
13 +os.environ[
14 + 'VLLM_NEURON_FRAMEWORK'] = NeuronFramework.TRANSFORMERS_NEURONX.value
15 +
16 +
17 +def _create_neuron_model_runner(model: str, *args,
18 + **kwargs) -> NeuronModelRunner:
19 + engine_args = EngineArgs(model, *args, **kwargs)
20 + engine_config = engine_args.create_engine_config()
21 + vllm_config = VllmConfig(
22 + model_config=engine_config.model_config,
23 + parallel_config=engine_config.parallel_config,
24 + scheduler_config=engine_config.scheduler_config,
25 + device_config=engine_config.device_config,
26 + )
27 + neuron_model_runner = NeuronModelRunner(vllm_config=vllm_config)
28 + return neuron_model_runner
29 +
30 +
31 +def test_update_neuron_sampling_params_not_full_batch():
32 + os.environ["NEURON_ON_DEVICE_SAMPLING_DISABLED"] = "0"
33 + model_runner = _create_neuron_model_runner(
34 + "facebook/opt-125m",
35 + seed=0,
36 + dtype="float16",
37 + max_num_seqs=2,
38 + )
39 + assert not model_runner._on_device_sampling_disabled
40 + # Test sampling param updating only when TNx is framework
41 + # NxDI handles sampling parameter updating inside model
42 + if current_platform.use_transformers_neuronx():
43 + model_mock = MagicMock()
44 + model_runner.model = model_mock
45 +
46 + seq_group_metadata_list = [
47 + SequenceGroupMetadata(
48 + request_id="test_0",
49 + is_prompt=True,
50 + seq_data={0: SequenceData.from_seqs([1, 2, 3])},
51 + sampling_params=SamplingParams(temperature=0.5,
52 + top_k=1,
53 + top_p=0.5),
54 + block_tables={0: [1]},
55 + )
56 + ]
57 +
58 + model_runner.prepare_model_input(seq_group_metadata_list)
59 +
60 + # Index neuron sampling parameters based on block_tables indices.
61 + # The first block_id of the sequence 0 is 1, so its parameters are
62 + # placed at index 1. So the sampling parameters will be:
63 + # Index 0: default sampling parameters
64 + # Index 1: sequecne 0's sampling parameters.
65 + neuron_sampling_params = (
66 + model_runner.model_config.neuron_sampling_params)
67 + assert neuron_sampling_params.temperature == [1.0, 0.5]
68 + assert neuron_sampling_params.top_k == [
69 + model_runner._MAX_NEURON_SAMPLING_TOP_K, 1
70 + ]
71 + assert neuron_sampling_params.top_p == [1.0, 0.5]
72 + model_mock.model.update_generation_config.assert_called_once_with(
73 + neuron_sampling_params)
74 +
75 +
76 +def test_update_neuron_sampling_params_full_batch():
77 + os.environ["NEURON_ON_DEVICE_SAMPLING_DISABLED"] = "0"
78 + model_runner = _create_neuron_model_runner(
79 + "facebook/opt-125m",
80 + seed=0,
81 + dtype="float16",
82 + max_num_seqs=2,
83 + )
84 + assert not model_runner._on_device_sampling_disabled
85 +
86 + # Test sampling param updating only when TNx is framework
87 + # NxDI handles sampling parameter updating inside model
88 + if current_platform.use_transformers_neuronx():
89 + model_mock = MagicMock()
90 + model_runner.model = model_mock
91 +
92 + seq_group_metadata_list = [
93 + SequenceGroupMetadata(
94 + request_id="test_0",
95 + is_prompt=True,
96 + seq_data={0: SequenceData.from_seqs([1, 2, 3])},
97 + sampling_params=SamplingParams(temperature=0.5,
98 + top_k=1,
99 + top_p=0.5),
100 + block_tables={0: [1]},
101 + ),
102 + SequenceGroupMetadata(
103 + request_id="test_0",
104 + is_prompt=True,
105 + seq_data={1: SequenceData.from_seqs([4, 5, 6])},
106 + sampling_params=SamplingParams(temperature=0.2,
107 + top_k=2,
108 + top_p=0.2),
109 + block_tables={1: [0]},
110 + )
111 + ]
112 +
113 + model_runner.prepare_model_input(seq_group_metadata_list)
114 +
115 + # Index neuron sampling parameters based on block_tables indices.
116 + # The first block_id of the sequence 0 is 1, so its parameters are
117 + # placed at index 1. So the sampling parameters will be:
118 + # Index 0: sequence 1's sampling parameters
119 + # Index 1: sequecne 0's sampling parameters.
120 + neuron_sampling_params = (
121 + model_runner.model_config.neuron_sampling_params)
122 + assert neuron_sampling_params.temperature == [0.2, 0.5]
123 + assert neuron_sampling_params.top_k == [2, 1]
124 + assert neuron_sampling_params.top_p == [0.2, 0.5]
125 + model_mock.model.update_generation_config.assert_called_once_with(
126 + neuron_sampling_params)
vllm/config.py
+5 -1
@@ -2273,6 +2273,9 @@ class SpeculativeConfig:
2273 """Scaling factor for entropy-based threshold, applied when using
2274 `TypicalAcceptanceSampler`."""
2275
2276 + speculative_token_tree: Optional[str] = None
2277 + """Specifies the tree structure for speculative token generation.
2278 + """
2279 # required configuration params passed from engine
2280 target_model_config: ModelConfig = field(default=None,
2281 init=True) # type: ignore
@@ -2447,10 +2450,11 @@ class SpeculativeConfig:
2450 "Chunked prefill and EAGLE are not compatible "
2451 "when using V0.")
2452
2453 + from vllm.platforms import current_platform
2454 from vllm.transformers_utils.configs.eagle import (
2455 EAGLEConfig)
2456 if isinstance(self.draft_model_config.hf_config,
2453 - EAGLEConfig):
2457 + EAGLEConfig) or current_platform.is_neuron():
2458 pass
2459 else:
2460 eagle_config = EAGLEConfig(
vllm/engine/llm_engine.py
+2 -4
@@ -399,10 +399,8 @@ class LLMEngine:
399 self.scheduler,
400 self.seq_counter,
401 get_tokenizer_for_seq,
402 - stop_checker=StopChecker(
403 - self.scheduler_config.max_model_len,
404 - get_tokenizer_for_seq,
405 - ),
402 + stop_checker=StopChecker(self.scheduler_config.max_model_len,
403 + get_tokenizer_for_seq),
404 ))
405
406 self.seq_id_to_seq_group: Dict[str, SequenceGroupBase] = {}
vllm/model_executor/model_loader/neuron.py
+237 -4
@@ -1,5 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 -"""Utilities for selecting and loading neuron models."""
2 +"""Utilities for selecting and loading Neuron models in transformers-neuronx
3 +framework."""
4 +import ast
5 import copy
6 import importlib
7 import os
@@ -9,7 +11,8 @@ import torch
11 import torch.nn as nn
12 from transformers import PretrainedConfig
13
12 -from vllm.config import ModelConfig, ParallelConfig, SchedulerConfig
14 +from vllm.config import (ModelConfig, ParallelConfig, SchedulerConfig,
15 + SpeculativeConfig)
16 from vllm.model_executor.layers.logits_processor import LogitsProcessor
17 from vllm.model_executor.layers.quantization import get_quantization_config
18 from vllm.model_executor.layers.sampler import Sampler, SamplerOutput
@@ -113,6 +116,67 @@ class NeuronCausalLM(nn.Module):
116 self.model.to_neuron()
117
118
119 +class NeuronSpeculationCausalLM(nn.Module):
120 + """A Neuron-optimized causal language model with speculative decoding."""
121 +
122 + SPECULATION_TERMINATION_ID = -1
123 +
124 + def __init__(self, speculation_model) -> None:
125 + super().__init__()
126 + self.model = speculation_model
127 +
128 + def forward(
129 + self,
130 + input_ids: torch.Tensor,
131 + positions: torch.Tensor,
132 + input_block_ids: torch.Tensor,
133 + ) -> torch.Tensor:
134 + tokens, counts = self.model.speculative_iteration(
135 + input_ids, positions, input_block_ids)
136 +
137 + # Mark the end of accepted speculative tokens for each sequence with the
138 + # speculation termination id.
139 + batch_size, steps = tokens.shape
140 + mask = torch.arange(steps).expand(batch_size, -1) >= counts
141 + tokens[mask] = self.SPECULATION_TERMINATION_ID
142 +
143 + return tokens
144 +
145 + def sample(
146 + self,
147 + logits: torch.Tensor,
148 + sampling_metadata: SamplingMetadata,
149 + ) -> Optional[List[SamplerOutput]]:
150 + batch_size, num_steps = logits.shape
151 + seq_ids = [
152 + seq_id for sg in sampling_metadata.seq_groups
153 + for seq_id in sg.seq_ids
154 + ]
155 + # Organize input tensors by step instead of by sequence.
156 + accepted_token_ids_by_step = logits.transpose(0, 1)
157 + accepted_token_ids_by_step = accepted_token_ids_by_step.tolist()
158 +
159 + sampler_output_list = []
160 + for step_index in range(num_steps):
161 + if all(token_id == self.SPECULATION_TERMINATION_ID
162 + for token_id in accepted_token_ids_by_step[step_index]):
163 + break
164 + step_output_token_ids = []
165 + for sequence_index in range(batch_size):
166 + token_id = accepted_token_ids_by_step[step_index][
167 + sequence_index]
168 + step_output_token_ids.append(
169 + CompletionSequenceGroupOutput(samples=[
170 + SequenceOutput(parent_seq_id=seq_ids[sequence_index],
171 + output_token=token_id,
172 + logprobs={token_id: Logprob(token_id)})
173 + ],
174 + prompt_logprobs=None))
175 + sampler_output_list.append(
176 + SamplerOutput(outputs=step_output_token_ids))
177 + return sampler_output_list
178 +
179 +
180 def _get_model_architecture(config: PretrainedConfig) -> str:
181 architectures = getattr(config, "architectures", [])
182 for arch in architectures:
@@ -138,6 +202,7 @@ def _get_buckets(env: str, default_value: List[int]) -> List[int]:
202 def _get_default_neuron_config(model_config: ModelConfig,
203 parallel_config: ParallelConfig,
204 scheduler_config: SchedulerConfig):
205 + """Generate a neuron config based on vllm config args."""
206 from transformers_neuronx.config import ContinuousBatchingConfig
207 from transformers_neuronx.constants import LAYOUT_BSH
208
@@ -162,6 +227,27 @@ def _get_default_neuron_config(model_config: ModelConfig,
227 return default_neuron_args
228
229
230 +def _get_default_neuron_config_for_speculation(
231 + model_config: ModelConfig, parallel_config: ParallelConfig,
232 + scheduler_config: SchedulerConfig):
233 + """Generate a neuron config for speculative decoding based on
234 + vllm config args."""
235 + from transformers_neuronx.config import ContinuousBatchingConfig
236 + from transformers_neuronx.constants import LAYOUT_BSH
237 +
238 + continuous_batching_config = ContinuousBatchingConfig(
239 + batch_size_for_shared_caches=scheduler_config.max_num_seqs)
240 +
241 + default_neuron_args = dict(collectives_layout=LAYOUT_BSH,
242 + attention_layout=LAYOUT_BSH,
243 + fuse_qkv=True,
244 + on_device_embedding=True,
245 + continuous_batching=continuous_batching_config,
246 + on_device_generation=copy.deepcopy(
247 + model_config.neuron_sampling_params))
248 + return default_neuron_args
249 +
250 +
251 def _get_neuron_on_device_generation_config(model_config: ModelConfig):
252 if not _is_neuron_on_device_sampling_disabled(model_config):
253 return copy.deepcopy(model_config.neuron_sampling_params)
@@ -213,7 +299,7 @@ def _get_neuron_config_after_override(default_neuron_config,
299 def get_neuron_model(model_config: ModelConfig,
300 parallel_config: ParallelConfig,
301 scheduler_config: SchedulerConfig) -> nn.Module:
216 -
302 + """Initializes a neuron-optimized model for inference."""
303 # Create a model instance.
304 model = NeuronCausalLM(
305 model_config.hf_config,
@@ -230,7 +316,6 @@ def get_neuron_model(model_config: ModelConfig,
316 n_positions = _get_buckets("NEURON_TOKEN_GEN_BUCKETS",
317 [scheduler_config.max_model_len])
318
233 - # Load the weights from the cached or downloaded files.
319 model.load_weights(model_config.model,
320 tp_degree=parallel_config.tensor_parallel_size,
321 amp=TORCH_DTYPE_TO_NEURON_AMP[model_config.dtype],
@@ -240,3 +325,151 @@ def get_neuron_model(model_config: ModelConfig,
325 batch_size=scheduler_config.max_num_seqs)
326
327 return model.eval()
328 +
329 +
330 +def get_neuron_speculation_model(model_config: ModelConfig,
331 + parallel_config: ParallelConfig,
332 + scheduler_config: SchedulerConfig,
333 + speculation_config: SpeculativeConfig):
334 + """Initializes a neuron-optimized speculation model for inference.
335 +
336 + This method is only applicable for speculation with a standalone draft model
337 + """
338 + from transformers_neuronx.fused_speculation import FusedSpeculativeDecoder
339 +
340 + # For Eagle SD, we need to pass in additional parameters in neuron config.
341 + is_eagle = getattr(speculation_config.draft_model_config.hf_config,
342 + "is_eagle", False)
343 +
344 + # Create target model instance.
345 + target_model = NeuronCausalLM(model_config.hf_config)
346 +
347 + default_neuron_config_args = _get_default_neuron_config_for_speculation(
348 + model_config, parallel_config, scheduler_config)
349 + if is_eagle:
350 + default_neuron_config_args['is_eagle_target'] = True
351 +
352 + neuron_config = _get_neuron_config_after_override(
353 + default_neuron_config_args, model_config.override_neuron_config)
354 +
355 + context_length_estimates = _get_buckets("NEURON_CONTEXT_LENGTH_BUCKETS",
356 + [scheduler_config.max_model_len])
357 + n_positions = _get_buckets("NEURON_TOKEN_GEN_BUCKETS",
358 + [scheduler_config.max_model_len])
359 +
360 + target_model.load_weights(
361 + model_config.model,
362 + tp_degree=parallel_config.tensor_parallel_size,
363 + amp=TORCH_DTYPE_TO_NEURON_AMP[model_config.dtype],
364 + neuron_config=neuron_config,
365 + context_length_estimate=context_length_estimates,
366 + n_positions=n_positions,
367 + batch_size=scheduler_config.max_num_seqs)
368 +
369 + target_model.eval()
370 +
371 + # Create draft model instance.
372 + draft_model = NeuronCausalLM(
373 + speculation_config.draft_model_config.hf_config)
374 +
375 + default_draft_neuron_config_args = (
376 + _get_default_neuron_config_for_speculation(
377 + speculation_config.draft_model_config, parallel_config,
378 + scheduler_config))
379 + if is_eagle:
380 + default_draft_neuron_config_args['is_eagle_draft'] = True
381 + default_draft_neuron_config_args['has_pre_attention_norm'] = False
382 +
383 + draft_neuron_config = _get_neuron_config_after_override(
384 + default_draft_neuron_config_args,
385 + speculation_config.draft_model_config.override_neuron_config)
386 +
387 + draft_model.load_weights(speculation_config.draft_model_config.model,
388 + tp_degree=speculation_config.
389 + draft_parallel_config.tensor_parallel_size,
390 + amp=TORCH_DTYPE_TO_NEURON_AMP[
391 + speculation_config.draft_model_config.dtype],
392 + neuron_config=draft_neuron_config,
393 + context_length_estimate=context_length_estimates,
394 + n_positions=n_positions,
395 + batch_size=scheduler_config.max_num_seqs)
396 +
397 + draft_model.eval()
398 +
399 + num_speculative_tokens = speculation_config.num_speculative_tokens
400 + # Create speculation model instance.
401 + speculation_model = FusedSpeculativeDecoder(draft_model.model,
402 + target_model.model,
403 + num_speculative_tokens)
404 + speculation_model.to_neuron()
405 +
406 + return NeuronSpeculationCausalLM(speculation_model)
407 +
408 +
409 +def get_neuron_eagle_speculation_model(model_config: ModelConfig,
410 + parallel_config: ParallelConfig,
411 + scheduler_config: SchedulerConfig,
412 + speculation_config: SpeculativeConfig):
413 + """Initializes a neuron-optimized EAGLE speculation model for inference."""
414 + from transformers_neuronx.eagle_speculation import EagleSpeculativeDecoder
415 +
416 + # Create target model instance.
417 + target_model = NeuronCausalLM(model_config.hf_config)
418 +
419 + default_neuron_config_args = _get_default_neuron_config_for_speculation(
420 + model_config, parallel_config, scheduler_config)
421 + default_neuron_config_args['is_eagle_target'] = True
422 + neuron_config = _get_neuron_config_after_override(
423 + default_neuron_config_args, model_config.override_neuron_config)
424 +
425 + context_length_estimates = _get_buckets("NEURON_CONTEXT_LENGTH_BUCKETS",
426 + [scheduler_config.max_model_len])
427 + n_positions = _get_buckets("NEURON_TOKEN_GEN_BUCKETS",
428 + [scheduler_config.max_model_len])
429 +
430 + target_model.load_weights(
431 + model_config.model,
432 + tp_degree=parallel_config.tensor_parallel_size,
433 + amp=TORCH_DTYPE_TO_NEURON_AMP[model_config.dtype],
434 + neuron_config=neuron_config,
435 + context_length_estimate=context_length_estimates,
436 + n_positions=n_positions,
437 + batch_size=scheduler_config.max_num_seqs)
438 +
439 + target_model.eval()
440 +
441 + # Create draft model instance.
442 + draft_model = NeuronCausalLM(
443 + speculation_config.draft_model_config.hf_config)
444 +
445 + default_draft_neuron_config_args = (
446 + _get_default_neuron_config_for_speculation(
447 + speculation_config.draft_model_config, parallel_config,
448 + scheduler_config))
449 + default_draft_neuron_config_args['is_eagle_draft'] = True
450 + default_draft_neuron_config_args['has_pre_attention_norm'] = False
451 + draft_neuron_config = _get_neuron_config_after_override(
452 + default_draft_neuron_config_args,
453 + speculation_config.draft_model_config.override_neuron_config)
454 +
455 + draft_model.load_weights(speculation_config.draft_model_config.model,
456 + tp_degree=speculation_config.
457 + draft_parallel_config.tensor_parallel_size,
458 + amp=TORCH_DTYPE_TO_NEURON_AMP[
459 + speculation_config.draft_model_config.dtype],
460 + neuron_config=draft_neuron_config,
461 + context_length_estimate=context_length_estimates,
462 + n_positions=n_positions,
463 + batch_size=scheduler_config.max_num_seqs)
464 +
465 + draft_model.eval()
466 +
467 + token_tree: Dict[int, List[int]] = ast.literal_eval(
468 + speculation_config.speculative_token_tree)
469 +
470 + speculation_model = EagleSpeculativeDecoder(draft_model.model,
471 + target_model.model,
472 + token_tree=token_tree)
473 + speculation_model.to_neuron()
474 +
475 + return NeuronSpeculationCausalLM(speculation_model)
vllm/model_executor/model_loader/neuronx_distributed.py new
+584
@@ -0,0 +1,584 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +"""Utilities for selecting and loading Neuron models in
3 +neuronx-distributed-inference framework."""
4 +# Disabling yapf because yapf and isort have conflicts for the below imports
5 +# yapf: disable
6 +import copy
7 +import hashlib
8 +import importlib
9 +import multiprocessing
10 +import os
11 +import shutil
12 +from typing import Dict, List, Optional, Tuple
13 +
14 +import torch
15 +import torch.nn as nn
16 +from neuronx_distributed_inference.models.config import (
17 + FusedSpecNeuronConfig, OnDeviceSamplingConfig)
18 +from neuronx_distributed_inference.models.mllama.utils import (
19 + create_vision_mask)
20 +from neuronx_distributed_inference.utils.hf_adapter import (
21 + load_pretrained_config)
22 +from transformers import AutoModelForCausalLM, AutoTokenizer, PretrainedConfig
23 +
24 +from vllm.config import (ModelConfig, ParallelConfig, SchedulerConfig,
25 + SpeculativeConfig)
26 +from vllm.logger import init_logger
27 +from vllm.model_executor.layers.logits_processor import LogitsProcessor
28 +from vllm.model_executor.layers.sampler import Sampler, SamplerOutput
29 +from vllm.model_executor.sampling_metadata import SamplingMetadata
30 +from vllm.sequence import (CompletionSequenceGroupOutput, Logprob,
31 + SequenceOutput)
32 +
33 +# yapf: enable
34 +logger = init_logger(__name__)
35 +
36 +TORCH_DTYPE_TO_NEURON_AMP = {
37 + "auto": "float32",
38 + "half": "float16",
39 + "float16": "float16",
40 + "bfloat16": "bfloat16",
41 + "float": "float32",
42 + "float32": "float32",
43 + torch.float16: "float16",
44 + torch.bfloat16: "bfloat16",
45 + torch.float32: "float32",
46 +}
47 +
48 +# Models supported by Neuronx distributed for inference.
49 +_NEURON_SUPPORTED_MODELS: Dict[str, Tuple[str, str]] = {
50 + "LlamaForCausalLM":
51 + ("neuronx_distributed_inference.models.llama.modeling_llama",
52 + "NeuronLlamaForCausalLM"),
53 + "DbrxForCausalLM":
54 + ("neuronx_distributed_inference.models.dbrx.modeling_dbrx",
55 + "NeuronDbrxForCausalLM"),
56 + "MixtralForCausalLM":
57 + ("neuronx_distributed_inference.models.mixtral.modeling_mixtral",
58 + "NeuronMixtralForCausalLM"),
59 + "MllamaForConditionalGeneration":
60 + ("neuronx_distributed_inference.models.mllama.modeling_mllama",
61 + "NeuronMllamaForCausalLM"),
62 +}
63 +
64 +
65 +class NeuronCausalLM(nn.Module):
66 +
67 + def __init__(
68 + self,
69 + config: PretrainedConfig,
70 + ) -> None:
71 + super().__init__()
72 + self.config = config
73 + self.logits_processor = LogitsProcessor(config.vocab_size,
74 + logits_as_input=True)
75 + self.sampler = Sampler()
76 +
77 + # Lazy initialized
78 + self.model: nn.Module
79 +
80 + def forward(
81 + self,
82 + input_ids: torch.Tensor,
83 + positions: torch.Tensor,
84 + input_block_ids: torch.Tensor,
85 + sampling_params: torch.Tensor,
86 + ) -> torch.Tensor:
87 + output = self.model(input_ids,
88 + attention_mask=None,
89 + position_ids=positions,
90 + seq_ids=input_block_ids,
91 + sampling_params=sampling_params)
92 + # on-device sampling
93 + if self.config.neuron_config.on_device_sampling_config:
94 + return output.hidden_states
95 + else:
96 + return output.logits[:, -1, :]
97 +
98 + def compute_logits(self, hidden_states: torch.Tensor,
99 + sampling_metadata: SamplingMetadata) -> torch.Tensor:
100 + logits = self.logits_processor(None, hidden_states, sampling_metadata)
101 + return logits
102 +
103 + def sample(
104 + self,
105 + logits: torch.Tensor,
106 + sampling_metadata: SamplingMetadata,
107 + ) -> Optional[SamplerOutput]:
108 + # on-device sampling
109 + if self.config.neuron_config.on_device_sampling_config:
110 + batch_size = logits.shape
111 + seq_ids = [
112 + seq_id for sg in sampling_metadata.seq_groups
113 + for seq_id in sg.seq_ids
114 + ]
115 + assert len(seq_ids) == list(batch_size)[0], "batch size mismatch"
116 + # Organize input tensors by step instead of by sequence.
117 + accepted_token_ids_by_step = logits.flatten()
118 + accepted_token_ids_by_step = accepted_token_ids_by_step.tolist()
119 +
120 + step_output_token_ids = []
121 + for i, seq_id in enumerate(seq_ids):
122 + token_id = accepted_token_ids_by_step[i]
123 + step_output_token_ids.append(
124 + CompletionSequenceGroupOutput(samples=[
125 + SequenceOutput(parent_seq_id=seq_id,
126 + output_token=token_id,
127 + logprobs={token_id: Logprob(token_id)})
128 + ],
129 + prompt_logprobs=None))
130 + return SamplerOutput(outputs=step_output_token_ids)
131 + else:
132 + return self.sampler(logits, sampling_metadata)
133 +
134 + def load_weights(self, model_name_or_path: str, **kwargs):
135 + arch = _get_model_architecture(self.config)
136 + neuronx_module_path, neuronx_model_cls_name = (
137 + _NEURON_SUPPORTED_MODELS[arch])
138 + neuronx_module = importlib.import_module(neuronx_module_path)
139 + neuronx_model_cls = getattr(neuronx_module, neuronx_model_cls_name)
140 + neuron_config = neuronx_model_cls.get_neuron_config_cls()(
141 + **kwargs['neuron_config'])
142 + self.config.neuron_config = neuron_config
143 + config = neuronx_model_cls.get_config_cls()(
144 + neuron_config,
145 + load_config=load_pretrained_config(model_name_or_path))
146 + hashed_config = hashlib.md5(
147 + config.to_json_string().encode('utf-8')).hexdigest()
148 + if os.getenv("NEURON_COMPILED_ARTIFACTS") is not None:
149 + compiled_model_path = os.getenv("NEURON_COMPILED_ARTIFACTS")
150 + elif os.path.exists(model_name_or_path):
151 + compiled_model_path = os.path.join(model_name_or_path,
152 + "neuron-compiled-artifacts",
153 + hashed_config)
154 + shutil.rmtree(compiled_model_path, ignore_errors=True)
155 + else:
156 + compiled_model_path = os.path.join("local-models",
157 + model_name_or_path,
158 + "neuron-compiled-artifacts",
159 + hashed_config)
160 + shutil.rmtree(compiled_model_path, ignore_errors=True)
161 + try:
162 + self.model = neuronx_model_cls(compiled_model_path)
163 + override_neuron_config = kwargs["override_neuron_config"]
164 + for k, v in override_neuron_config.items():
165 + setattr(self.model.config.neuron_config, k, v)
166 + self.model.load(compiled_model_path)
167 + return
168 + except (FileNotFoundError, ValueError) as e:
169 + logger.warning("Exception: %s", e)
170 + logger.warning("Failed to load the model from %s, Recompiling...",
171 + compiled_model_path)
172 + if not os.path.exists(model_name_or_path):
173 + hf_model = AutoModelForCausalLM.from_pretrained(model_name_or_path)
174 + saved_path = os.path.join("local-models", model_name_or_path)
175 + hf_model.save_pretrained(saved_path)
176 + model_name_or_path = saved_path
177 + self.model = neuronx_model_cls(model_name_or_path, config)
178 + self.model.compile(compiled_model_path)
179 + self.model.load(compiled_model_path)
180 +
181 +
182 +class NeuronMllamaForCausalLM(nn.Module):
183 +
184 + def __init__(self,
185 + config: PretrainedConfig,
186 + on_device_sampling_disabled: bool = False) -> None:
187 + super().__init__()
188 + self.config = config
189 + self.logits_processor = LogitsProcessor(
190 + config.get_text_config().vocab_size, logits_as_input=True)
191 +
192 + self.on_device_sampling_disabled = on_device_sampling_disabled
193 + if self.on_device_sampling_disabled:
194 + # Use default sampler
195 + self.sampler = Sampler()
196 +
197 + # Lazy initialized
198 + self.model: nn.Module
199 +
200 + def forward(self, input_ids: torch.Tensor, positions: torch.Tensor,
201 + seq_ids: torch.Tensor, pixel_values: torch.Tensor,
202 + aspect_ratios: torch.Tensor, num_chunks: torch.Tensor,
203 + has_image: torch.Tensor, sampling_params) -> torch.Tensor:
204 + self.vision_mask = create_vision_mask(input_ids, self.vision_token_id)
205 + output = self.model(
206 + input_ids.to(torch.int32),
207 + attention_mask=None,
208 + position_ids=positions.to(torch.int32),
209 + seq_ids=seq_ids.flatten().to(torch.int32),
210 + pixel_values=pixel_values.to(
211 + self.config.vision_config.torch_dtype),
212 + aspect_ratios=aspect_ratios.to(torch.int32),
213 + vision_mask=self.vision_mask.to(torch.int32),
214 + sampling_params=sampling_params,
215 + num_chunks=num_chunks.to(torch.int32),
216 + has_image=has_image.to(torch.int32),
217 + )
218 + if self.config.neuron_config.on_device_sampling_config:
219 + return output.hidden_states
220 + return output.logits[:, -1, :]
221 +
222 + def compute_logits(self, hidden_states: torch.Tensor,
223 + sampling_metadata: SamplingMetadata) -> torch.Tensor:
224 + logits = self.logits_processor(None, hidden_states, sampling_metadata)
225 + return logits
226 +
227 + def sample(self, hidden_states, sampling_metadata):
228 + if not self.on_device_sampling_disabled:
229 + with torch.profiler.record_function("sample"):
230 + hidden_states = hidden_states.flatten()
231 + res = []
232 + sample_idx = 0
233 + for seq_group in sampling_metadata.seq_groups:
234 + seq_ids = seq_group.seq_ids
235 + samples = []
236 + for seq_id in seq_ids:
237 + token_id = hidden_states[sample_idx].item()
238 + samples.append(
239 + SequenceOutput(
240 + parent_seq_id=seq_id,
241 + output_token=token_id,
242 + logprobs={token_id: Logprob(token_id)}))
243 + sample_idx += 1
244 + res.append(
245 + CompletionSequenceGroupOutput(samples=samples,
246 + prompt_logprobs=None))
247 + next_tokens = SamplerOutput(outputs=res)
248 + else:
249 + next_tokens = self.sampler(None, hidden_states, sampling_metadata)
250 + return next_tokens
251 +
252 + def load_weights(self, model_name_or_path: str, **kwargs):
253 + arch = _get_model_architecture(self.config)
254 + neuronx_module_path, neuronx_model_cls_name = (
255 + _NEURON_SUPPORTED_MODELS[arch])
256 + neuronx_module = importlib.import_module(neuronx_module_path)
257 + neuronx_model_cls = getattr(neuronx_module, neuronx_model_cls_name)
258 + neuron_config = neuronx_model_cls.get_neuron_config_cls()(
259 + **kwargs['neuron_config'])
260 + self.config.neuron_config = neuron_config
261 + logger.info("neuron_config buckets: %s",
262 + self.config.neuron_config.buckets)
263 + config = neuronx_model_cls.get_config_cls()(
264 + neuron_config,
265 + load_config=load_pretrained_config(model_name_or_path))
266 + hashed_config = hashlib.md5(
267 + config.to_json_string().encode('utf-8')).hexdigest()
268 + if os.getenv("NEURON_COMPILED_ARTIFACTS") is not None:
269 + compiled_model_path = os.getenv("NEURON_COMPILED_ARTIFACTS")
270 + elif os.path.exists(model_name_or_path):
271 + compiled_model_path = os.path.join(model_name_or_path,
272 + "neuron-compiled-artifacts",
273 + hashed_config)
274 + else:
275 + compiled_model_path = os.path.join("local-models",
276 + model_name_or_path,
277 + "neuron-compiled-artifacts",
278 + hashed_config)
279 + try:
280 + self.model = neuronx_model_cls(compiled_model_path)
281 + tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
282 + self.vision_token_id = tokenizer(
283 + "<|image|>", add_special_tokens=False).input_ids
284 + self.model.load(compiled_model_path)
285 + return
286 + except (FileNotFoundError, ValueError):
287 + logger.warning("Failed to load the model from %s, Recompiling...",
288 + compiled_model_path)
289 + if not os.path.exists(model_name_or_path):
290 + hf_model = AutoModelForCausalLM.from_pretrained(model_name_or_path)
291 + saved_path = os.path.join("local-models", model_name_or_path)
292 + hf_model.save_pretrained(saved_path)
293 + model_name_or_path = saved_path
294 + self.model = neuronx_model_cls(model_name_or_path, config)
295 +
296 + logger.info("\nCompiling and saving model to %s", model_name_or_path)
297 +
298 + p = multiprocessing.Process(target=compile_model,
299 + args=(self, compiled_model_path))
300 + p.start()
301 + p.join()
302 +
303 + tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
304 + tokenizer.save_pretrained(compiled_model_path)
305 + logger.info("Successfully compiled and saved the model in %s",
306 + compiled_model_path)
307 +
308 + # Read "<|image|>" token_id from the tokenizer
309 + self.vision_token_id = tokenizer("<|image|>",
310 + add_special_tokens=False).input_ids
311 + logger.info("\nLoading model from compiled checkpoint...")
312 + self.model.load(compiled_model_path)
313 +
314 +
315 +def compile_model(neuron_model, traced_model_path):
316 + neuron_model.model.compile(traced_model_path)
317 +
318 +
319 +class NeuronSpeculationCausalLM(nn.Module):
320 + """A Neuron-optimized causal language model with speculative decoding."""
321 +
322 + def __init__(
323 + self,
324 + config: PretrainedConfig,
325 + ) -> None:
326 + super().__init__()
327 + self.config = config
328 + self.logits_processor = LogitsProcessor(config.vocab_size,
329 + logits_as_input=True)
330 + # Lazy initialized
331 + self.model: nn.Module
332 +
333 + def forward(
334 + self,
335 + input_ids: torch.Tensor,
336 + positions: torch.Tensor,
337 + input_block_ids: torch.Tensor,
338 + sampling_params: torch.Tensor,
339 + ) -> torch.Tensor:
340 + output = self.model(input_ids,
341 + attention_mask=None,
342 + position_ids=positions,
343 + seq_ids=input_block_ids,
344 + sampling_params=sampling_params)
345 + # CTX encoding
346 + if (positions[:, 0]).sum().item() == 0:
347 + return output.fused_outputs[0][:, 0:1]
348 +
349 + # Fused Spec (Generation)
350 + accepted_tokens_with_padding = output.fused_outputs[0]
351 + next_pos_ids = output.fused_outputs[-1]
352 + generated_token_counts = next_pos_ids - positions
353 +
354 + assert torch.any(generated_token_counts == 0).item() is False, \
355 + "NxDI model generated no output for one or more sequences."
356 +
357 + batch_size, steps = accepted_tokens_with_padding.shape
358 + mask = torch.arange(steps).expand(batch_size,
359 + -1) >= generated_token_counts
360 + accepted_tokens_with_padding[mask] = -1
361 +
362 + return accepted_tokens_with_padding
363 +
364 + def sample(
365 + self,
366 + logits: torch.Tensor,
367 + sampling_metadata: SamplingMetadata,
368 + ) -> Optional[List[SamplerOutput]]:
369 + batch_size, num_steps = logits.shape
370 + seq_ids = [
371 + seq_id for sg in sampling_metadata.seq_groups
372 + for seq_id in sg.seq_ids
373 + ]
374 + # Organize input tensors by step instead of by sequence.
375 + accepted_token_ids_by_step = logits.transpose(0, 1)
376 + accepted_token_ids_by_step = accepted_token_ids_by_step.tolist()
377 +
378 + sampler_output_list = []
379 + for step_index in range(num_steps):
380 + if all(token_id == -1
381 + for token_id in accepted_token_ids_by_step[step_index]):
382 + break
383 + step_output_token_ids = []
384 + for sequence_index in range(batch_size):
385 + token_id = accepted_token_ids_by_step[step_index][
386 + sequence_index]
387 + step_output_token_ids.append(
388 + CompletionSequenceGroupOutput(samples=[
389 + SequenceOutput(parent_seq_id=seq_ids[sequence_index],
390 + output_token=token_id,
391 + logprobs={token_id: Logprob(token_id)})
392 + ],
393 + prompt_logprobs=None))
394 + sampler_output_list.append(
395 + SamplerOutput(outputs=step_output_token_ids))
396 + return sampler_output_list
397 +
398 + def load_weights(self, model_name_or_path: str,
399 + draft_model_name_or_path: str, **kwargs):
400 + arch = _get_model_architecture(self.config)
401 + neuronx_module_path, neuronx_model_cls_name = (
402 + _NEURON_SUPPORTED_MODELS[arch])
403 + neuronx_module = importlib.import_module(neuronx_module_path)
404 + neuronx_model_cls = getattr(neuronx_module, neuronx_model_cls_name)
405 + neuron_config = neuronx_model_cls.get_neuron_config_cls()(
406 + **kwargs['neuron_config'])
407 + config = neuronx_model_cls.get_config_cls()(
408 + neuron_config,
409 + load_config=load_pretrained_config(model_name_or_path))
410 +
411 + draft_neuron_config = copy.deepcopy(config.neuron_config)
412 + if not config.neuron_config.enable_eagle_speculation:
413 + draft_neuron_config.speculation_length = 0
414 + draft_neuron_config.trace_tokengen_model = True
415 + draft_neuron_config.enable_fused_speculation = False
416 + if config.neuron_config.enable_eagle_speculation:
417 + draft_neuron_config.is_eagle_draft = True
418 + draft_neuron_config.sequence_parallel_enabled = False
419 + draft_config = neuronx_model_cls.get_config_cls()(
420 + draft_neuron_config,
421 + load_config=load_pretrained_config(draft_model_name_or_path))
422 + fused_spec_config = (FusedSpecNeuronConfig(
423 + neuronx_model_cls._model_cls,
424 + draft_config=draft_config,
425 + draft_model_path=draft_model_name_or_path))
426 + config.fused_spec_config = fused_spec_config
427 + self.config.neuron_config = neuron_config
428 +
429 + hashed_config = hashlib.md5(
430 + config.to_json_string().encode('utf-8')).hexdigest()
431 + if os.getenv("NEURON_COMPILED_ARTIFACTS") is not None:
432 + compiled_model_path = os.getenv("NEURON_COMPILED_ARTIFACTS")
433 + elif os.path.exists(model_name_or_path):
434 + compiled_model_path = os.path.join(model_name_or_path,
435 + "neuron-compiled-artifacts",
436 + hashed_config)
437 + shutil.rmtree(compiled_model_path, ignore_errors=True)
438 + else:
439 + compiled_model_path = os.path.join("local-models",
440 + model_name_or_path,
441 + "neuron-compiled-artifacts",
442 + hashed_config)
443 + shutil.rmtree(compiled_model_path, ignore_errors=True)
444 + try:
445 + self.model = neuronx_model_cls(compiled_model_path)
446 + override_neuron_config = kwargs["override_neuron_config"]
447 + for k, v in override_neuron_config.items():
448 + setattr(self.model.config.neuron_config, k, v)
449 + self.model.load(compiled_model_path)
450 + return
451 + except (FileNotFoundError, ValueError) as e:
452 + logger.warning("Exception: %s", e)
453 + logger.warning("Failed to load the model from %s Recompiling...",
454 + compiled_model_path)
455 + if not os.path.exists(model_name_or_path):
456 + hf_model = AutoModelForCausalLM.from_pretrained(model_name_or_path)
457 + saved_path = os.path.join("local-models", model_name_or_path)
458 + hf_model.save_pretrained(saved_path)
459 + model_name_or_path = saved_path
460 + if not os.path.exists(draft_model_name_or_path):
461 + if draft_model_name_or_path != model_name_or_path:
462 + hf_model = AutoModelForCausalLM.from_pretrained(
463 + draft_model_name_or_path)
464 + saved_path = os.path.join("local-models",
465 + draft_model_name_or_path)
466 + hf_model.save_pretrained(saved_path)
467 + draft_model_name_or_path = saved_path
468 + else:
469 + draft_model_name_or_path = model_name_or_path
470 + config.fused_spec_config.draft_model_path = draft_model_name_or_path
471 + self.model = neuronx_model_cls(model_name_or_path, config)
472 + self.model.compile(compiled_model_path)
473 + self.model.load(compiled_model_path)
474 +
475 +
476 +def _get_model_architecture(config: PretrainedConfig) -> str:
477 + architectures = getattr(config, "architectures", [])
478 + for arch in architectures:
479 + if arch in _NEURON_SUPPORTED_MODELS:
480 + return arch
481 + raise ValueError(
482 + f"Model architectures {architectures} are not supported on Neuron "
483 + f"for now. Supported architectures: "
484 + f"{list(_NEURON_SUPPORTED_MODELS.keys())}")
485 +
486 +
487 +def _get_default_neuron_config(model_config: ModelConfig,
488 + parallel_config: ParallelConfig,
489 + scheduler_config: SchedulerConfig):
490 + """Generate a neuron config based on vllm config args."""
491 + on_device_sampling_config = OnDeviceSamplingConfig(dynamic=True,
492 + deterministic=False)
493 + batch_size = scheduler_config.max_num_seqs
494 +
495 + neuron_config = dict(
496 + tp_degree=parallel_config.tensor_parallel_size,
497 + ctx_batch_size=1,
498 + batch_size=batch_size,
499 + max_context_length=scheduler_config.max_model_len,
500 + seq_len=scheduler_config.max_model_len,
501 + enable_bucketing=True,
502 + is_continuous_batching=(batch_size > 1),
503 + quantized=False,
504 + torch_dtype=TORCH_DTYPE_TO_NEURON_AMP[model_config.dtype],
505 + padding_side="right",
506 + on_device_sampling_config=on_device_sampling_config,
507 + sequence_parallel_enabled=True,
508 + )
509 + return neuron_config
510 +
511 +
512 +def _get_default_speculation_config(model_config: ModelConfig,
513 + parallel_config: ParallelConfig,
514 + scheduler_config: SchedulerConfig,
515 + speculation_config: SpeculativeConfig):
516 + """Generate a neuron config for speculative decoding based on vllm config
517 + args."""
518 + neuron_config = dict(
519 + tp_degree=parallel_config.tensor_parallel_size,
520 + batch_size=scheduler_config.max_num_seqs,
521 + max_context_length=scheduler_config.max_model_len,
522 + seq_len=scheduler_config.max_model_len,
523 + speculation_length=speculation_config.num_speculative_tokens,
524 + trace_tokengen_model=False,
525 + enable_fused_speculation=True,
526 + enable_bucketing=True,
527 + quantized=False,
528 + torch_dtype=TORCH_DTYPE_TO_NEURON_AMP[model_config.dtype],
529 + on_device_sampling_config=dict(
530 + top_k=1,
531 + do_sample=False,
532 + ))
533 + return neuron_config
534 +
535 +
536 +def _get_neuron_config_after_override(default_neuron_config,
537 + overridden_neuron_config):
538 + """Update default neuron config values with override args"""
539 + overridden_neuron_config = overridden_neuron_config or {}
540 + default_neuron_config.update(overridden_neuron_config)
541 + return default_neuron_config
542 +
543 +
544 +def get_neuron_model(model_config: ModelConfig,
545 + parallel_config: ParallelConfig,
546 + scheduler_config: SchedulerConfig) -> nn.Module:
547 + """Initializes a neuron-optimized model for inference."""
548 + model_arch = _get_model_architecture(model_config.hf_config)
549 + if model_arch == "MllamaForConditionalGeneration":
550 + model = NeuronMllamaForCausalLM(model_config.hf_config)
551 + else:
552 + model = NeuronCausalLM(model_config.hf_config)
553 + default_neuron_config_args = _get_default_neuron_config(
554 + model_config, parallel_config, scheduler_config)
555 + neuron_config = _get_neuron_config_after_override(
556 + default_neuron_config_args, model_config.override_neuron_config)
557 +
558 + override_neuron_config = model_config.override_neuron_config
559 + model.load_weights(model_config.model,
560 + neuron_config=neuron_config,
561 + override_neuron_config=override_neuron_config)
562 + return model.eval()
563 +
564 +
565 +def get_neuron_speculation_model(model_config: ModelConfig,
566 + parallel_config: ParallelConfig,
567 + scheduler_config: SchedulerConfig,
568 + speculation_config: SpeculativeConfig):
569 + """Initializes a neuron-optimized speculation model for inference.
570 +
571 + This model handles speculation using both a draft model and an EAGLE draft.
572 + """
573 + model = NeuronSpeculationCausalLM(model_config.hf_config)
574 + default_neuron_config_args = _get_default_speculation_config(
575 + model_config, parallel_config, scheduler_config, speculation_config)
576 + neuron_config = _get_neuron_config_after_override(
577 + default_neuron_config_args, model_config.override_neuron_config)
578 +
579 + override_neuron_config = model_config.override_neuron_config
580 + model.load_weights(model_config.model,
581 + speculation_config.draft_model_config.model,
582 + neuron_config=neuron_config,
583 + override_neuron_config=override_neuron_config)
584 + return model.eval()
vllm/platforms/__init__.py
+13 -4
@@ -176,17 +176,26 @@ def cpu_platform_plugin() -> Optional[str]:
176
177
178 def neuron_platform_plugin() -> Optional[str]:
179 - is_neuron = False
179 + tnx_installed = False
180 + nxd_installed = False
181 logger.debug("Checking if Neuron platform is available.")
182 try:
183 import transformers_neuronx # noqa: F401
183 - is_neuron = True
184 + tnx_installed = True
185 logger.debug("Confirmed Neuron platform is available because"
186 " transformers_neuronx is found.")
186 - except ImportError as e:
187 - logger.debug("Neuron platform is not available because: %s", str(e))
187 + except ImportError:
188 pass
189
190 + try:
191 + import neuronx_distributed_inference # noqa: F401
192 + nxd_installed = True
193 + logger.debug("Confirmed Neuron platform is available because"
194 + " neuronx_distributed_inference is found.")
195 + except ImportError:
196 + pass
197 +
198 + is_neuron = tnx_installed or nxd_installed
199 return "vllm.platforms.neuron.NeuronPlatform" if is_neuron else None
200
201
vllm/platforms/neuron.py
+76 -3
@@ -1,5 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 -
2 +import enum
3 +import os
4 +from functools import lru_cache
5 from typing import TYPE_CHECKING, Optional
6
7 from vllm import envs
@@ -15,6 +17,11 @@ else:
17 logger = init_logger(__name__)
18
19
20 +class NeuronFramework(enum.Enum):
21 + TRANSFORMERS_NEURONX = "transformers-neuronx"
22 + NEURONX_DISTRIBUTED_INFERENCE = "neuronx-distributed-inference"
23 +
24 +
25 class NeuronPlatform(Platform):
26 _enum = PlatformEnum.NEURON
27 device_name: str = "neuron"
@@ -43,8 +50,6 @@ class NeuronPlatform(Platform):
50
51 assert (vllm_config.lora_config
52 is None), "LoRA is not supported for Neuron backend."
46 - assert (not vllm_config.speculative_config
47 - ), "Speculative decoding not yet supported for Neuron backend."
53
54 cache_config = vllm_config.cache_config
55 if cache_config:
@@ -67,3 +72,71 @@ class NeuronPlatform(Platform):
72 @classmethod
73 def use_all_gather(cls) -> bool:
74 return True
75 +
76 + @classmethod
77 + @lru_cache
78 + def is_neuronx_distributed_inference(cls) -> bool:
79 + try:
80 + import neuronx_distributed_inference
81 + except ImportError:
82 + neuronx_distributed_inference = None
83 + return neuronx_distributed_inference is not None
84 +
85 + @classmethod
86 + @lru_cache
87 + def is_transformers_neuronx(cls) -> bool:
88 + try:
89 + import transformers_neuronx
90 + except ImportError:
91 + transformers_neuronx = None
92 + return transformers_neuronx is not None
93 +
94 + def get_neuron_framework_to_use(self):
95 + """Return the specified framework if corresponding installations are
96 + available.
97 +
98 + If no framework is specified, use neuronx-distributed-inference by
99 + default.
100 + If that's unavailable, check and switch to transformers-neuronx.
101 + """
102 + if not self.is_neuron():
103 + raise AssertionError(
104 + f"Neuron Framework unavailable for platform: {self}")
105 +
106 + tnx_installed = self.is_transformers_neuronx()
107 + nxd_installed = self.is_neuronx_distributed_inference()
108 +
109 + specified_framework = os.environ.get("VLLM_NEURON_FRAMEWORK")
110 + tnx_framework = NeuronFramework.TRANSFORMERS_NEURONX.value
111 + nxd_framework = NeuronFramework.NEURONX_DISTRIBUTED_INFERENCE.value
112 + if specified_framework == tnx_framework and tnx_installed:
113 + return self.TRANSFORMERS_NEURONX
114 +
115 + if ((specified_framework == nxd_framework and nxd_installed)
116 + or (specified_framework is None and nxd_installed)):
117 + return NeuronFramework.NEURONX_DISTRIBUTED_INFERENCE
118 +
119 + if specified_framework is None and tnx_installed:
120 + return NeuronFramework.TRANSFORMERS_NEURONX
121 +
122 + return None
123 +
124 + def use_neuronx_distributed(self):
125 + """
126 + Return True if the framework determined in get_neuron_framework_to_use()
127 + is NeuronFramework.NEURONX_DISTRIBUTED_INFERENCE, False otherwise. This
128 + is used to select the Neuron model framework and framework-specific
129 + configuration to apply during model compilation.
130 + """
131 + nxd_framework = NeuronFramework.NEURONX_DISTRIBUTED_INFERENCE
132 + return self.get_neuron_framework_to_use() == nxd_framework
133 +
134 + def use_transformers_neuronx(self):
135 + """
136 + Return True if the framework determined in get_neuron_framework_to_use()
137 + is NeuronFramework.TRANSFORMERS_NEURONX, False otherwise. This is used
138 + to select the Neuron model framework and framework-specific
139 + configuration to apply during model compilation.
140 + """
141 + return self.get_neuron_framework_to_use(
142 + ) == NeuronFramework.TRANSFORMERS_NEURONX
vllm/worker/multi_step_neuron_model_runner.py new
+81
@@ -0,0 +1,81 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +
3 +from importlib.util import find_spec
4 +from typing import List, Optional
5 +
6 +import torch
7 +
8 +from vllm.config import VllmConfig
9 +from vllm.model_executor.layers.sampler import SamplerOutput
10 +from vllm.multimodal import MultiModalKwargs
11 +from vllm.sequence import IntermediateTensors
12 +from vllm.worker.neuron_model_runner import (ModelInputForNeuron,
13 + NeuronModelRunner)
14 +
15 +
16 +class MultiStepNeuronModelRunner(NeuronModelRunner):
17 + """A model runner for multi step decoding using the transformers_neuronx
18 + framework"""
19 +
20 + def __init__(
21 + self,
22 + vllm_config: VllmConfig,
23 + ):
24 + super().__init__(vllm_config)
25 + self.speculation_config = self.speculative_config
26 + from transformers_neuronx.config import GenerationConfig
27 + self.speculation_config.draft_model_config.neuron_sampling_params = (
28 + GenerationConfig(
29 + max_length=self.scheduler_config.max_model_len,
30 + do_sample=True,
31 + per_batch_line=True,
32 + top_k=[self._MAX_NEURON_SAMPLING_TOP_K] \
33 + * self.scheduler_config.max_num_seqs,
34 + top_p=[1.0] * self.scheduler_config.max_num_seqs,
35 + temperature=[1.0] * self.scheduler_config.max_num_seqs,
36 + dynamic=True,
37 + global_top_k=self._MAX_NEURON_SAMPLING_TOP_K
38 + ))
39 +
40 + def load_model(self) -> None:
41 + if find_spec("transformers_neuronx") is not None:
42 + from vllm.model_executor.model_loader.neuron import (
43 + get_neuron_eagle_speculation_model,
44 + get_neuron_speculation_model)
45 + if self.speculation_config.speculative_token_tree is not None:
46 + self.model = get_neuron_eagle_speculation_model(
47 + self.model_config,
48 + parallel_config=self.parallel_config,
49 + scheduler_config=self.scheduler_config,
50 + speculation_config=self.speculation_config)
51 + else:
52 + self.model = get_neuron_speculation_model(
53 + self.model_config,
54 + parallel_config=self.parallel_config,
55 + scheduler_config=self.scheduler_config,
56 + speculation_config=self.speculation_config)
57 + else:
58 + raise NotImplementedError(
59 + "Supports only Transformer-NeuronX based models.")
60 +
61 + @torch.inference_mode()
62 + def execute_model(
63 + self,
64 + model_input: ModelInputForNeuron,
65 + kv_caches: Optional[List[torch.Tensor]] = None,
66 + intermediate_tensors: Optional[IntermediateTensors] = None,
67 + num_steps: int = 1,
68 + ) -> Optional[List[SamplerOutput]]:
69 + logits = self.model(
70 + input_ids=model_input.input_tokens,
71 + positions=model_input.input_positions,
72 + input_block_ids=model_input.input_block_ids,
73 + **MultiModalKwargs.as_kwargs(model_input.multi_modal_kwargs or {},
74 + device=self.device),
75 + )
76 +
77 + output = self.model.sample(
78 + logits=logits,
79 + sampling_metadata=model_input.sampling_metadata,
80 + )
81 + return output
vllm/worker/multi_step_neuronx_distributed_model_runner.py new
+60
@@ -0,0 +1,60 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +from typing import List, Optional
3 +
4 +import torch
5 +
6 +from vllm.config import VllmConfig
7 +from vllm.model_executor.layers.sampler import SamplerOutput
8 +from vllm.multimodal import MultiModalKwargs
9 +from vllm.sequence import IntermediateTensors
10 +from vllm.worker.neuronx_distributed_model_runner import (
11 + NeuronxDistributedModelRunner)
12 +
13 +
14 +class MultiStepNeuronxDistributedModelRunner(NeuronxDistributedModelRunner):
15 + """A model runner for multi-step decoding using the
16 + neuronx-distributed-inference framework"""
17 +
18 + def __init__(
19 + self,
20 + vllm_config: VllmConfig,
21 + ):
22 + super().__init__(vllm_config)
23 +
24 + def load_model(self) -> None:
25 + from vllm.model_executor.model_loader.neuronx_distributed import (
26 + get_neuron_speculation_model)
27 + self.model = get_neuron_speculation_model(
28 + self.model_config,
29 + parallel_config=self.parallel_config,
30 + scheduler_config=self.scheduler_config,
31 + speculation_config=self.speculative_config)
32 +
33 + @torch.inference_mode()
34 + def execute_model(
35 + self,
36 + model_input,
37 + kv_caches: Optional[List[torch.Tensor]] = None,
38 + intermediate_tensors: Optional[IntermediateTensors] = None,
39 + num_steps: int = 1,
40 + ) -> Optional[List[SamplerOutput]]:
41 + sampling_params = torch.tensor([[
42 + seq_group.sampling_params.top_k,
43 + seq_group.sampling_params.top_p,
44 + seq_group.sampling_params.temperature,
45 + ] for seq_group in model_input.sampling_metadata.seq_groups])
46 +
47 + logits = self.model(
48 + input_ids=model_input.input_tokens,
49 + positions=model_input.input_positions,
50 + input_block_ids=model_input.input_block_ids,
51 + sampling_params=sampling_params,
52 + **MultiModalKwargs.as_kwargs(model_input.multi_modal_kwargs or {},
53 + device=self.device),
54 + )
55 +
56 + output = self.model.sample(
57 + logits=logits,
58 + sampling_metadata=model_input.sampling_metadata,
59 + )
60 + return output
vllm/worker/neuron_model_runner.py
+130 -52
@@ -2,20 +2,20 @@
2
3 import os
4 from dataclasses import dataclass
5 -from importlib.util import find_spec
5 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
6
7 import torch
8 from torch import nn
10 -from transformers_neuronx.config import GenerationConfig
9
12 -from vllm.config import VllmConfig
13 -from vllm.forward_context import set_forward_context
10 +from vllm.config import DeviceConfig, VllmConfig
11 from vllm.logger import init_logger
12 from vllm.model_executor import SamplingMetadata
13 from vllm.model_executor.layers.sampler import SamplerOutput
14 from vllm.model_executor.model_loader.neuron import get_neuron_model
18 -from vllm.multimodal import BatchedTensorInputs, MultiModalKwargs
15 +from vllm.multimodal import (MULTIMODAL_REGISTRY, BatchedTensorInputs,
16 + MultiModalKwargs)
17 +from vllm.platforms import current_platform
18 +from vllm.sampling_params import SamplingParams
19 from vllm.sequence import IntermediateTensors, SequenceGroupMetadata
20 from vllm.utils import is_pin_memory_available, make_tensor_with_pad
21 from vllm.worker.model_runner_base import ModelRunnerBase, ModelRunnerInputBase
@@ -34,12 +34,18 @@ class ModelInputForNeuron(ModelRunnerInputBase):
34 input_tokens: Optional[torch.Tensor] = None
35 input_positions: Optional[torch.Tensor] = None
36 input_block_ids: Optional[torch.Tensor] = None
37 - sampling_metadata: Optional["SamplingMetadata"] = None
38 - multi_modal_kwargs: Optional[BatchedTensorInputs] = None
37 + sampling_metadata: SamplingMetadata = None
38 + multi_modal_kwargs: BatchedTensorInputs = None
39
40 def as_broadcastable_tensor_dict(
41 self) -> Dict[str, Union[int, torch.Tensor]]:
42 - raise NotImplementedError("ModelInputForNeuron cannot be broadcast.")
42 + return {
43 + "input_tokens": self.input_tokens,
44 + "input_positions": self.input_positions,
45 + "input_block_ids": self.input_block_ids,
46 + "sampling_metadata": self.sampling_metadata,
47 + "multi_modal_kwargs": self.multi_modal_kwargs,
48 + }
49
50 @classmethod
51 def from_broadcasted_tensor_dict(
@@ -47,11 +53,17 @@ class ModelInputForNeuron(ModelRunnerInputBase):
53 tensor_dict: Dict[str, Any],
54 attn_backend: Optional["AttentionBackend"] = None,
55 ) -> "ModelInputForNeuron":
50 - assert attn_backend is None
51 - return cls.from_broadcasted_tensor_dict(tensor_dict)
56 + return ModelInputForNeuron(
57 + input_tokens=tensor_dict["input_tokens"],
58 + input_positions=tensor_dict["input_positions"],
59 + input_block_ids=tensor_dict["input_block_ids"],
60 + sampling_metadata=tensor_dict["sampling_metadata"],
61 + multi_modal_kwargs=tensor_dict["multi_modal_kwargs"],
62 + )
63
64
65 class NeuronModelRunner(ModelRunnerBase[ModelInputForNeuron]):
66 + """A model runner for AWS Neuron hardware"""
67
68 # NEURON has an upper limit on the top_k
69 _MAX_NEURON_SAMPLING_TOP_K = 256
@@ -61,13 +73,20 @@ class NeuronModelRunner(ModelRunnerBase[ModelInputForNeuron]):
73 vllm_config: VllmConfig,
74 ):
75 ModelRunnerBase.__init__(self, vllm_config)
64 - model_config = self.model_config
65 - if model_config is not None and model_config.get_sliding_window():
76 +
77 + if (self.model_config is not None
78 + and self.model_config.get_sliding_window()):
79 logger.warning("Sliding window is not supported on Neuron. "
80 "The model will run without sliding window.")
81 + self.device_config = (self.device_config if self.device_config
82 + is not None else DeviceConfig())
83 self.device = self.device_config.device
84 self.pin_memory = is_pin_memory_available()
85
86 + # Multi-modal data support
87 + self.multi_modal_input_mapper = MULTIMODAL_REGISTRY \
88 + .create_input_mapper(self.model_config)
89 +
90 # Lazy initialization.
91 self.model: nn.Module # initialize after load_model.
92
@@ -82,32 +101,33 @@ class NeuronModelRunner(ModelRunnerBase[ModelInputForNeuron]):
101 self._previous_batch_request_ids: List[str] = []
102
103 if not self._on_device_sampling_disabled:
85 - logger.warning(
86 - "On-device sampling is turned on in Neuron by default, only "
87 - "top_k, top_p, and temperature are current supported sampling "
88 - "parameters. To turn off the on-device sampling, please set "
89 - "the environment variable NEURON_ON_DEVICE_SAMPLING_DISABLED=1."
90 - )
91 - self.model_config.neuron_sampling_params = GenerationConfig(
92 - max_length=self.scheduler_config.max_model_len,
93 - do_sample=True,
94 - per_batch_line=True,
95 - top_k=[self._MAX_NEURON_SAMPLING_TOP_K] \
96 - * self.scheduler_config.max_num_seqs,
97 - top_p=[1.0] * self.scheduler_config.max_num_seqs,
98 - temperature=[1.0] * self.scheduler_config.max_num_seqs,
99 - dynamic=True,
100 - global_top_k=self._MAX_NEURON_SAMPLING_TOP_K)
104 + self._init_neuron_sampling()
105
102 - def load_model(self) -> None:
103 - if find_spec("transformers_neuronx") is not None:
104 - self.model = get_neuron_model(
105 - self.model_config,
106 - parallel_config=self.parallel_config,
107 - scheduler_config=self.scheduler_config)
106 + def _init_neuron_sampling(self) -> None:
107 + if current_platform.use_transformers_neuronx():
108 + from transformers_neuronx.config import GenerationConfig
109 else:
109 - raise NotImplementedError(
110 - "Supports only Transformer-NeuronX based models.")
110 + from transformers import GenerationConfig
111 + logger.warning(
112 + "On-device sampling is turned on in Neuron by default, only "
113 + "top_k, top_p, and temperature are current supported sampling "
114 + "parameters. To turn off the on-device sampling, please set "
115 + "the environment variable NEURON_ON_DEVICE_SAMPLING_DISABLED=1.")
116 + self.model_config.neuron_sampling_params = GenerationConfig(
117 + max_length=self.scheduler_config.max_model_len,
118 + do_sample=True,
119 + per_batch_line=True,
120 + top_k=[self._MAX_NEURON_SAMPLING_TOP_K] \
121 + * self.scheduler_config.max_num_seqs,
122 + top_p=[1.0] * self.scheduler_config.max_num_seqs,
123 + temperature=[1.0] * self.scheduler_config.max_num_seqs,
124 + dynamic=True,
125 + global_top_k=self._MAX_NEURON_SAMPLING_TOP_K)
126 +
127 + def load_model(self) -> None:
128 + self.model = get_neuron_model(self.model_config,
129 + parallel_config=self.parallel_config,
130 + scheduler_config=self.scheduler_config)
131
132 def get_model(self) -> nn.Module:
133 return self.model
@@ -240,6 +260,16 @@ class NeuronModelRunner(ModelRunnerBase[ModelInputForNeuron]):
260 (input_tokens, input_positions,
261 input_block_ids) = self._prepare_decode(seq_group_metadata_list)
262 seq_lens = None
263 +
264 + if not self._on_device_sampling_disabled:
265 + for seq_group_metadata in seq_group_metadata_list:
266 + sampling_params = seq_group_metadata.sampling_params
267 + top_k, top_p, temperature = (
268 + self._convert_to_neuron_sampling_params(sampling_params))
269 + sampling_params.top_k = top_k
270 + sampling_params.top_p = top_p
271 + sampling_params.temperature = temperature
272 +
273 sampling_metadata = SamplingMetadata.prepare(
274 seq_group_metadata_list,
275 seq_lens,
@@ -251,7 +281,8 @@ class NeuronModelRunner(ModelRunnerBase[ModelInputForNeuron]):
281 self.pin_memory,
282 generators=self.get_generators(finished_requests_ids))
283
254 - if not self._on_device_sampling_disabled:
284 + if current_platform.use_transformers_neuronx(
285 + ) and not self._on_device_sampling_disabled:
286 # Once the request IDs are changed in current iteration, we will
287 # update the on-device sampling parameters.
288 current_batch_request_ids = [
@@ -259,7 +290,7 @@ class NeuronModelRunner(ModelRunnerBase[ModelInputForNeuron]):
290 for seq_group_meta_data in seq_group_metadata_list
291 ]
292 if current_batch_request_ids != self._previous_batch_request_ids:
262 - self._update_neuron_sampling_params(sampling_metadata)
293 + self._update_neuron_sampling_params(seq_group_metadata_list)
294 self._previous_batch_request_ids = current_batch_request_ids
295
296 return ModelInputForNeuron(input_tokens=input_tokens,
@@ -268,31 +299,59 @@ class NeuronModelRunner(ModelRunnerBase[ModelInputForNeuron]):
299 sampling_metadata=sampling_metadata,
300 multi_modal_kwargs=multi_modal_kwargs)
301
271 - def _update_neuron_sampling_params(self,
272 - sampling_metadata: SamplingMetadata):
302 + def _update_neuron_sampling_params(
303 + self, seq_group_metadata_list: List[SequenceGroupMetadata]):
304 # Update Neuron sampling parameters (GenerationConfig in Neuron)
305 current_sampling_params = self.model_config.neuron_sampling_params
306 assert current_sampling_params is not None, (
307 f"Failed to update sampling_params, "
308 f"current sampling params is {current_sampling_params}")
309
310 + is_update_needed = False
311 +
312 top_k = current_sampling_params.top_k
313 top_p = current_sampling_params.top_p
314 temperature = current_sampling_params.temperature
282 - for index, sequence_group_to_sample in enumerate(
283 - sampling_metadata.seq_groups):
284 - top_k[index] = self._convert_to_neuron_top_k(
285 - sequence_group_to_sample.sampling_params.top_k)
286 - top_p[index] = sequence_group_to_sample.sampling_params.top_p
287 - temperature[index] = \
288 - sequence_group_to_sample.sampling_params.temperature
315
290 - self.model.model.update_generation_config(current_sampling_params)
316 + # The index of a sequence's sampling parameters in neuron is equal to
317 + # its index in `input_block_ids`.
318 + for seq_group_metadata in seq_group_metadata_list:
319 + seq_ids = list(seq_group_metadata.seq_data.keys())
320 + sampling_params = seq_group_metadata.sampling_params
321 +
322 + seq_group_top_k = sampling_params.top_k
323 + seq_group_top_p = sampling_params.top_p
324 + seq_group_temperature = sampling_params.temperature
325
292 - def _convert_to_neuron_top_k(self, top_k: int) -> int:
326 + for seq_id in seq_ids:
327 + index = seq_group_metadata.block_tables[seq_id][0]
328 + if (top_k[index] != seq_group_top_k
329 + or top_p[index] != seq_group_top_p
330 + or temperature[index] != seq_group_temperature):
331 + is_update_needed = True
332 +
333 + top_k[index] = seq_group_top_k
334 + top_p[index] = seq_group_top_p
335 + temperature[index] = seq_group_temperature
336 +
337 + # update_generation_config is only available in transformers-neuronx
338 + if is_update_needed and current_platform.use_transformers_neuronx():
339 + self.model.model.update_generation_config(current_sampling_params)
340 +
341 + def _convert_to_neuron_sampling_params(
342 + self, sampling_params: SamplingParams) -> Tuple[int, float, float]:
343 + # Returns the top_k, top_p and temperature parameters for neuron.
344 + top_k = sampling_params.top_k
345 + top_p = sampling_params.top_p
346 + temperature = sampling_params.temperature
347 +
348 + if temperature == 0.0:
349 + # Enable greedy sampling on zero temperature
350 + return (1, 1.0, 1.0)
351 if top_k < 0 or top_k > self._MAX_NEURON_SAMPLING_TOP_K:
294 - return self._MAX_NEURON_SAMPLING_TOP_K
295 - return top_k
352 + top_k = self._MAX_NEURON_SAMPLING_TOP_K
353 +
354 + return (top_k, top_p, temperature)
355
356 @torch.inference_mode()
357 def execute_model(
@@ -306,7 +365,26 @@ class NeuronModelRunner(ModelRunnerBase[ModelInputForNeuron]):
365 raise ValueError(
366 "NeuronModelRunner does not support multi-step execution.")
367
309 - with set_forward_context(None, self.vllm_config, 0):
368 + # extract top_k, top_p and temperature from model_input for neuron
369 + # forward call
370 + sampling_params = (torch.tensor([[
371 + seq_group.sampling_params.top_k, seq_group.sampling_params.top_p,
372 + seq_group.sampling_params.temperature
373 + ] for seq_group in model_input.sampling_metadata.seq_groups]))
374 +
375 + if current_platform.use_neuronx_distributed():
376 + hidden_states = self.model(
377 + input_ids=model_input.input_tokens,
378 + positions=model_input.input_positions,
379 + input_block_ids=model_input.input_block_ids,
380 + sampling_params=sampling_params,
381 + **MultiModalKwargs.as_kwargs(model_input.multi_modal_kwargs
382 + or {},
383 + device=self.device),
384 + )
385 + elif current_platform.use_transformers_neuronx():
386 + # [TODO] validate on-device sampling
387 + # The model signature may need change for on-device sampling
388 hidden_states = self.model(
389 input_ids=model_input.input_tokens,
390 positions=model_input.input_positions,
vllm/worker/neuron_worker.py
+53 -33
@@ -1,61 +1,81 @@
1 # SPDX-License-Identifier: Apache-2.0
2 """A Neuron worker class."""
3 +import os
4 from typing import List, Optional, Tuple
5
5 -import torch
6 import torch.distributed
7
8 from vllm.config import VllmConfig
9 from vllm.distributed import (ensure_model_parallel_initialized,
10 init_distributed_environment)
11 +from vllm.logger import init_logger
12 from vllm.model_executor import set_random_seed
12 -from vllm.model_executor.layers.sampler import SamplerOutput
13 +from vllm.platforms import current_platform
14 +from vllm.platforms.neuron import NeuronFramework
15 from vllm.sequence import ExecuteModelRequest
16 from vllm.worker.neuron_model_runner import NeuronModelRunner
17 from vllm.worker.worker_base import (LocalOrDistributedWorkerBase,
18 LoRANotSupportedWorkerBase, WorkerBase,
19 WorkerInput)
20
21 +logger = init_logger(__name__)
22 +
23
24 class NeuronWorker(LoRANotSupportedWorkerBase, LocalOrDistributedWorkerBase):
25 """A worker class that executes the model on a group of neuron cores.
26 """
27
24 - def __init__(
25 - self,
26 - vllm_config: VllmConfig,
27 - local_rank: int,
28 - rank: int,
29 - distributed_init_method: str,
30 - is_driver_worker: bool = True,
31 - ) -> None:
28 + model_runner: NeuronModelRunner
29 +
30 + def __init__(self,
31 + vllm_config: VllmConfig,
32 + local_rank: int,
33 + rank: int,
34 + distributed_init_method: str,
35 + is_driver_worker: bool = False) -> None:
36 WorkerBase.__init__(self, vllm_config=vllm_config)
37 self.local_rank = local_rank
38 self.rank = rank
39 self.distributed_init_method = distributed_init_method
40 + self.is_driver_worker = is_driver_worker
41 +
42 if self.model_config.trust_remote_code:
43 # note: lazy import to avoid importing torch before initializing
44 from vllm.utils import init_cached_hf_modules
45 init_cached_hf_modules()
46
41 - self.model_runner: NeuronModelRunner = NeuronModelRunner(
42 - vllm_config=vllm_config)
43 - self.is_driver_worker = is_driver_worker
44 -
45 - def execute_model(
46 - self,
47 - execute_model_req: Optional[ExecuteModelRequest] = None,
48 - ) -> Optional[List[SamplerOutput]]:
49 - assert execute_model_req is not None
50 - assert (not execute_model_req.blocks_to_swap_in
51 - and not execute_model_req.blocks_to_swap_out
52 - and not execute_model_req.blocks_to_copy), (
53 - "Cache operations are not supported for Neuron backend.")
54 - assert execute_model_req.num_lookahead_slots == 0, (
55 - "lookahead not supported for Neuron backend.")
56 - output = LocalOrDistributedWorkerBase.execute_model(
57 - self, execute_model_req)
58 - return output
47 + neuron_framework = current_platform.get_neuron_framework_to_use()
48 + if neuron_framework == NeuronFramework.TRANSFORMERS_NEURONX:
49 + self.model_runner = self.get_tnx_model_runner(vllm_config)
50 + elif neuron_framework == NeuronFramework.NEURONX_DISTRIBUTED_INFERENCE:
51 + self.model_runner = self.get_neuronx_distributed_model_runner(
52 + vllm_config)
53 + else:
54 + raise NotImplementedError(
55 + "Specified framework" +
56 + f" {os.environ.get('VLLM_NEURON_FRAMEWORK')}" +
57 + " is either not installed or not supported." +
58 + " Supported frameworks: " +
59 + "[transformers-neuronx, neuronx-distributed-inference]")
60 +
61 + def get_tnx_model_runner(self, vllm_config):
62 + from vllm.worker.multi_step_neuron_model_runner import (
63 + MultiStepNeuronModelRunner)
64 + if self.speculative_config is not None:
65 + return MultiStepNeuronModelRunner(vllm_config=vllm_config)
66 + else:
67 + return NeuronModelRunner(vllm_config=vllm_config)
68 +
69 + def get_neuronx_distributed_model_runner(self, vllm_config):
70 + from vllm.worker.multi_step_neuronx_distributed_model_runner import (
71 + MultiStepNeuronxDistributedModelRunner)
72 + from vllm.worker.neuronx_distributed_model_runner import (
73 + NeuronxDistributedModelRunner)
74 + if self.speculative_config is not None:
75 + return MultiStepNeuronxDistributedModelRunner(
76 + vllm_config=vllm_config)
77 + else:
78 + return NeuronxDistributedModelRunner(vllm_config=vllm_config)
79
80 def init_device(self) -> None:
81 self.init_distributed_environment()
@@ -121,17 +141,17 @@ class NeuronWorker(LoRANotSupportedWorkerBase, LocalOrDistributedWorkerBase):
141
142 def init_distributed_environment(self):
143 """Neuron uses transformers-neuronx for tensor parallelism.
124 - It has only one process to control multiple devices.
125 - vLLM still needs the environment initialized when TP/PP > 1,
126 - so we initialize a distributed environment with one process.
144 +
145 + vLLM still needs the environment initialized when TP/PP > 1
146 """
147 init_distributed_environment(
148 world_size=1,
130 - rank=0,
131 - local_rank=0,
149 + rank=self.rank,
150 + local_rank=self.local_rank,
151 distributed_init_method=self.distributed_init_method,
152 backend="gloo",
153 )
154 +
155 ensure_model_parallel_initialized(
156 1,
157 1,
vllm/worker/neuronx_distributed_model_runner.py new
+136
@@ -0,0 +1,136 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +
3 +from typing import List, Optional
4 +
5 +import torch
6 +from neuronx_distributed_inference.modules.generation.sampling import (
7 + prepare_sampling_params)
8 +
9 +from vllm.config import VllmConfig
10 +from vllm.logger import init_logger
11 +from vllm.model_executor.layers.sampler import SamplerOutput
12 +from vllm.model_executor.model_loader.neuronx_distributed import (
13 + _get_model_architecture, get_neuron_model)
14 +from vllm.sequence import IntermediateTensors
15 +from vllm.worker.neuron_model_runner import (ModelInputForNeuron,
16 + NeuronModelRunner)
17 +
18 +logger = init_logger(__name__)
19 +
20 +
21 +class NeuronxDistributedModelRunner(NeuronModelRunner):
22 +
23 + def __init__(
24 + self,
25 + vllm_config: VllmConfig,
26 + ):
27 + super().__init__(vllm_config)
28 +
29 + def load_model(self) -> None:
30 + self.model = get_neuron_model(self.model_config,
31 + parallel_config=self.parallel_config,
32 + scheduler_config=self.scheduler_config)
33 +
34 + def get_nxd_sampling_params(self, sampling_metadata):
35 + if self.model.config.neuron_config.on_device_sampling_config:
36 + max_topk = (self.model.config.neuron_config.
37 + on_device_sampling_config.global_topk)
38 + else:
39 + max_topk = self.model.config.vocab_size
40 +
41 + top_k = [1] * self.scheduler_config.max_num_seqs
42 + top_p = [1.0] * self.scheduler_config.max_num_seqs
43 + temperature = [1.0] * self.scheduler_config.max_num_seqs
44 +
45 + for index, sequenceGroupToSample in enumerate(
46 + sampling_metadata.seq_groups):
47 + top_k[index] = (sequenceGroupToSample.sampling_params.top_k
48 + if sequenceGroupToSample.sampling_params.top_k > 0
49 + else max_topk)
50 + top_p[index] = sequenceGroupToSample.sampling_params.top_p
51 + temperature[index] = (
52 + sequenceGroupToSample.sampling_params.temperature)
53 +
54 + sampling_params = prepare_sampling_params(
55 + batch_size=self.scheduler_config.max_num_seqs,
56 + top_k=top_k,
57 + top_p=top_p,
58 + temperature=temperature)
59 + return sampling_params
60 +
61 + def get_multi_modal_data_neuron(self, input_images):
62 + raise NotImplementedError("need to restore multi-modal support")
63 +
64 + @torch.inference_mode()
65 + def execute_model(
66 + self,
67 + model_input: ModelInputForNeuron,
68 + kv_caches: Optional[List[torch.Tensor]] = None,
69 + intermediate_tensors: Optional[IntermediateTensors] = None,
70 + num_steps: int = 1,
71 + ) -> Optional[List[SamplerOutput]]:
72 + if num_steps > 1:
73 + raise ValueError(
74 + "NeuronModelRunner does not support multi-step execution.")
75 +
76 + if _get_model_architecture(
77 + self.model.config) != "MllamaForConditionalGeneration":
78 + return super().execute_model(model_input, kv_caches,
79 + intermediate_tensors, num_steps)
80 +
81 + sampling_params = self.get_nxd_sampling_params(
82 + model_input.sampling_metadata)
83 +
84 + if model_input.multi_modal_kwargs.get('image') is not None:
85 + pixel_values = []
86 + aspect_ratios = []
87 + num_chunks = []
88 + has_image = []
89 + for multi_modal_input in model_input.multi_modal_kwargs.get(
90 + 'image'):
91 + image_tensors = self.get_multi_modal_data_neuron(
92 + multi_modal_input.squeeze(0))
93 + pixel_values.append(image_tensors[0])
94 + aspect_ratios.append(image_tensors[1])
95 + num_chunks.append(image_tensors[2])
96 + has_image.append(image_tensors[3])
97 +
98 + pixel_values = torch.cat(pixel_values, dim=0)
99 + aspect_ratios = torch.cat(aspect_ratios, dim=0)
100 + num_chunks = torch.cat(num_chunks, dim=0)
101 + has_image = torch.cat(has_image, dim=0)
102 +
103 + hidden_states = self.model(
104 + input_ids=model_input.input_tokens,
105 + positions=model_input.input_positions,
106 + seq_ids=model_input.input_block_ids,
107 + pixel_values=pixel_values,
108 + aspect_ratios=aspect_ratios,
109 + sampling_params=sampling_params,
110 + num_chunks=num_chunks,
111 + has_image=has_image,
112 + )
113 + else:
114 + empty_pixel_values = torch.zeros([1, 1, 4, 3, 560, 560],
115 + dtype=torch.bfloat16)
116 + empty_aspect_ratios = torch.ones([1, 1, 2], dtype=torch.int64)
117 + num_chunks = torch.tensor([[1]
118 + ]) # dummy num_chunks, will not be used
119 + has_image = torch.tensor([0])
120 + hidden_states = self.model(
121 + input_ids=model_input.input_tokens,
122 + positions=model_input.input_positions,
123 + seq_ids=model_input.input_block_ids,
124 + pixel_values=empty_pixel_values,
125 + aspect_ratios=empty_aspect_ratios,
126 + sampling_params=sampling_params,
127 + num_chunks=num_chunks,
128 + has_image=has_image,
129 + )
130 +
131 + output = self.model.sample(
132 + hidden_states=hidden_states,
133 + sampling_metadata=model_input.sampling_metadata,
134 + )
135 +
136 + return [output]