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, 07:15 UTC
[Neuron] Add Multi-Modal model support for Neuron (#18921)
Signed-off-by: Satyajith Chilappagari <satchill@amazon.com> Co-authored-by: Ashraf Mahgoub <ashymahg@amazon.com> Co-authored-by: Rohith Nallamaddi <nalrohit@amazon.com> Co-authored-by: FeliciaLuo <luof@amazon.com> Co-authored-by: Elaine Zhao <elaineyz@amazon.com>
Satyajith Chilappagari committed
May 31, 2025 at 03:39 UTC
2a50ef57605ffe332a73e50597276b71e9d52676
5 files changed
+235
-46
examples/offline_inference/neuron_multimodal.py
new
+105
@@ -0,0 +1,105 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+import requests
3
+import torch
4
+from neuronx_distributed_inference.models.mllama.utils import add_instruct
5
+from PIL import Image
6
+
7
+from vllm import LLM, SamplingParams, TextPrompt
8
+
9
+
10
+def get_image(image_url):
11
+ image = Image.open(requests.get(image_url, stream=True).raw)
12
+ return image
13
+
14
+
15
+# Model Inputs
16
+PROMPTS = [
17
+ "What is in this image? Tell me a story",
18
+ "What is the recipe of mayonnaise in two sentences?",
19
+ "Describe this image",
20
+ "What is the capital of Italy famous for?",
21
+]
22
+IMAGES = [
23
+ get_image(
24
+ "https://images.pexels.com/photos/1108099/pexels-photo-1108099.jpeg?auto=compress&cs=tinysrgb&dpr=1&w=500"
25
+ ),
26
+ None,
27
+ get_image(
28
+ "https://images.pexels.com/photos/1108099/pexels-photo-1108099.jpeg?auto=compress&cs=tinysrgb&dpr=1&w=500"
29
+ ),
30
+ None,
31
+]
32
+SAMPLING_PARAMS = [
33
+ dict(top_k=1, temperature=1.0, top_p=1.0, max_tokens=16)
34
+ for _ in range(len(PROMPTS))
35
+]
36
+
37
+
38
+def get_VLLM_mllama_model_inputs(prompt, single_image, sampling_params):
39
+ # Prepare all inputs for mllama generation, including:
40
+ # 1. put text prompt into instruct chat template
41
+ # 2. compose single text and single image prompt into Vllm's prompt class
42
+ # 3. prepare sampling parameters
43
+ input_image = single_image
44
+ has_image = torch.tensor([1])
45
+ if isinstance(single_image, torch.Tensor) and single_image.numel() == 0:
46
+ has_image = torch.tensor([0])
47
+
48
+ instruct_prompt = add_instruct(prompt, has_image)
49
+ inputs = TextPrompt(prompt=instruct_prompt)
50
+
51
+ if input_image is not None:
52
+ inputs["multi_modal_data"] = {"image": input_image}
53
+
54
+ sampling_params = SamplingParams(**sampling_params)
55
+ return inputs, sampling_params
56
+
57
+
58
+def print_outputs(outputs):
59
+ # Print the outputs.
60
+ for output in outputs:
61
+ prompt = output.prompt
62
+ generated_text = output.outputs[0].text
63
+ print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
64
+
65
+
66
+if __name__ == "__main__":
67
+ assert (
68
+ len(PROMPTS) == len(IMAGES) == len(SAMPLING_PARAMS)
69
+ ), f"""Text, image prompts and sampling parameters should have the
70
+ same batch size; but got {len(PROMPTS)}, {len(IMAGES)},
71
+ and {len(SAMPLING_PARAMS)}"""
72
+
73
+ # Create an LLM.
74
+ llm = LLM(
75
+ model="meta-llama/Llama-3.2-11B-Vision-Instruct",
76
+ max_num_seqs=1,
77
+ max_model_len=4096,
78
+ block_size=4096,
79
+ device="neuron",
80
+ tensor_parallel_size=32,
81
+ override_neuron_config={
82
+ "sequence_parallel_enabled": False,
83
+ "skip_warmup": True,
84
+ "save_sharded_checkpoint": True,
85
+ "on_device_sampling_config": {
86
+ "global_topk": 1,
87
+ "dynamic": False,
88
+ "deterministic": False,
89
+ },
90
+ },
91
+ )
92
+
93
+ batched_inputs = []
94
+ batched_sample_params = []
95
+ for pmpt, img, params in zip(PROMPTS, IMAGES, SAMPLING_PARAMS):
96
+ inputs, sampling_params = get_VLLM_mllama_model_inputs(pmpt, img, params)
97
+ # test batch-size = 1
98
+ outputs = llm.generate(inputs, sampling_params)
99
+ print_outputs(outputs)
100
+ batched_inputs.append(inputs)
101
+ batched_sample_params.append(sampling_params)
102
+
103
+ # test batch-size = 4
104
+ outputs = llm.generate(batched_inputs, batched_sample_params)
105
+ print_outputs(outputs)
vllm/config.py
+10
@@ -1360,6 +1360,16 @@ class ModelConfig:
1360
@property
1361
def is_encoder_decoder(self) -> bool:
1362
"""Extract the HF encoder/decoder model flag."""
1363
+ """
1364
+ For Mllama, VLLM overrides HF's is_encoder_decoder flag and sets it to
1365
+ True to enable cross-attention
1366
+ Neuron needs all multimodal data to be in the decoder and does not
1367
+ need to explicitly enable cross-attention
1368
+ """
1369
+ if (current_platform.is_neuron()
1370
+ and self.hf_config.model_type == "mllama"):
1371
+ return False
1372
+
1373
return is_encoder_decoder(self.hf_config)
1374
1375
@property
vllm/model_executor/model_loader/neuronx_distributed.py
+61
-4
@@ -204,6 +204,11 @@ class NeuronMllamaForCausalLM(nn.Module):
204
config: PretrainedConfig,
205
on_device_sampling_disabled: bool = False) -> None:
206
super().__init__()
207
+ # has_image is the only multimodal input that is used in
208
+ # token-generation
209
+ # This is a cache (on CPU) that saves has_image data per sequence id
210
+ # The number of entries in this cache is <= Batch-Size
211
+ self.has_image_cache: dict[int, torch.Tensor] = {}
212
self.config = config
213
self.logits_processor = LogitsProcessor(
214
config.get_text_config().vocab_size, logits_as_input=True)
@@ -215,11 +220,57 @@ class NeuronMllamaForCausalLM(nn.Module):
220
221
# Lazy initialized
222
self.model: nn.Module
223
+ self.is_reorder_needed: bool = True
224
+
225
+ def read_from_has_image_cache(self, seq_ids: torch.Tensor):
226
+ has_image_list = []
227
+ for index in range(len(seq_ids)):
228
+ seq_id = seq_ids[index].item()
229
+ if seq_id in self.has_image_cache:
230
+ has_image_list.append(self.has_image_cache[seq_id])
231
+ else:
232
+ has_image_list.append(torch.tensor([0]))
233
+ return torch.tensor(has_image_list)
234
+
235
+ def write_to_has_image_cache(self, seq_ids: torch.Tensor,
236
+ has_image: torch.Tensor):
237
+ for index in range(len(seq_ids)):
238
+ seq_id = seq_ids[index].item()
239
+ if index < len(has_image):
240
+ self.has_image_cache[seq_id] = has_image[index]
241
+ else:
242
+ self.has_image_cache[seq_id] = torch.zeros(1)
243
244
def forward(self, input_ids: torch.Tensor, positions: torch.Tensor,
245
seq_ids: torch.Tensor, pixel_values: torch.Tensor,
246
aspect_ratios: torch.Tensor, num_chunks: torch.Tensor,
247
has_image: torch.Tensor, sampling_params) -> torch.Tensor:
248
+
249
+ # We update the has_image cache during prefill
250
+ # and read the has_image cache during decode
251
+ if input_ids.shape[-1] > 1: # prefill
252
+ self.write_to_has_image_cache(seq_ids, has_image)
253
+ else:
254
+ has_image = self.read_from_has_image_cache(seq_ids)
255
+ bs = input_ids.shape[0]
256
+ num_chunks = torch.zeros((bs, 1))
257
+ aspect_ratios = torch.zeros((bs, 1, 2))
258
+
259
+ input_block_ids = seq_ids
260
+ origin_input_block_ids = seq_ids
261
+ if self.is_reorder_needed:
262
+ # sort block ids sequentially for perf/neuron support reasons
263
+ input_block_ids, sorted_indices = torch.sort(input_block_ids)
264
+ input_ids = torch.index_select(input_ids, 0, sorted_indices)
265
+ positions = torch.index_select(positions, 0, sorted_indices)
266
+ sampling_params = torch.index_select(sampling_params, 0,
267
+ sorted_indices)
268
+ pixel_values = torch.index_select(pixel_values, 0, sorted_indices)
269
+ aspect_ratios = torch.index_select(aspect_ratios, 0,
270
+ sorted_indices)
271
+ num_chunks = torch.index_select(num_chunks, 0, sorted_indices)
272
+ has_image = torch.index_select(has_image, 0, sorted_indices)
273
+
274
self.vision_mask = create_vision_mask(input_ids, self.vision_token_id)
275
output = self.model(
276
input_ids.to(torch.int32),
@@ -235,8 +286,14 @@ class NeuronMllamaForCausalLM(nn.Module):
286
has_image=has_image.to(torch.int32),
287
)
288
if self.config.neuron_config.on_device_sampling_config:
238
- return output.hidden_states
239
- return output.logits[:, -1, :]
289
+ output = output.hidden_states
290
+ else:
291
+ output = output.logits[:, -1, :]
292
+
293
+ if self.is_reorder_needed and origin_input_block_ids.shape[0] != 1:
294
+ restored_indices = torch.argsort(sorted_indices)
295
+ output = torch.index_select(output, 0, restored_indices)
296
+ return output
297
298
def compute_logits(self, hidden_states: torch.Tensor,
299
sampling_metadata: SamplingMetadata) -> torch.Tensor:
@@ -299,7 +356,7 @@ class NeuronMllamaForCausalLM(nn.Module):
356
self.model = neuronx_model_cls(compiled_model_path)
357
tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
358
self.vision_token_id = tokenizer(
302
- "<|image|>", add_special_tokens=False).input_ids
359
+ "<|image|>", add_special_tokens=False).input_ids[0]
360
self.model.load(compiled_model_path)
361
return
362
except (FileNotFoundError, ValueError):
@@ -326,7 +383,7 @@ class NeuronMllamaForCausalLM(nn.Module):
383
384
# Read "<|image|>" token_id from the tokenizer
385
self.vision_token_id = tokenizer("<|image|>",
329
- add_special_tokens=False).input_ids
386
+ add_special_tokens=False).input_ids[0]
387
logger.info("\nLoading model from compiled checkpoint...")
388
self.model.load(compiled_model_path)
389
vllm/worker/neuron_model_runner.py
+13
@@ -169,6 +169,7 @@ class NeuronModelRunner(ModelRunnerBase[ModelInputForNeuron]):
169
170
mm_kwargs = seq_group_metadata.multi_modal_data
171
if mm_kwargs:
172
+ mm_kwargs = self.process_multi_modal_data_neuron(mm_kwargs)
173
multi_modal_kwargs_list.append(mm_kwargs)
174
175
max_seq_len = max(seq_lens)
@@ -274,6 +275,14 @@ class NeuronModelRunner(ModelRunnerBase[ModelInputForNeuron]):
275
sampling_params.top_p = top_p
276
sampling_params.temperature = temperature
277
278
+ # we need multi_modal_data for later tokens as well
279
+ multi_modal_kwargs_list: List[MultiModalKwargs] = []
280
+ for seq_group_metadata in seq_group_metadata_list:
281
+ mm_data = seq_group_metadata.multi_modal_data
282
+ if mm_data:
283
+ multi_modal_kwargs_list.append(mm_data)
284
+ multi_modal_kwargs = MultiModalKwargs.batch(multi_modal_kwargs_list)
285
+
286
sampling_metadata = SamplingMetadata.prepare(
287
seq_group_metadata_list,
288
seq_lens,
@@ -422,6 +431,10 @@ class NeuronModelRunner(ModelRunnerBase[ModelInputForNeuron]):
431
def vocab_size(self) -> int:
432
return self.model_config.get_vocab_size()
433
434
+ def process_multi_modal_data_neuron(self, mm_data):
435
+ # this is a no-op for NeuronModelRunner
436
+ return mm_data
437
+
438
def remove_all_loras(self):
439
raise NotImplementedError(
440
"LoRAs are not supported for Transformers NeuronX framework")
vllm/worker/neuronx_distributed_model_runner.py
+46
-42
@@ -3,6 +3,8 @@
3
from typing import List, Optional, Set
4
5
import torch
6
+from neuronx_distributed_inference.models.mllama.aspect_ratio_utils import (
7
+ get_all_supported_aspect_ratios)
8
from neuronx_distributed_inference.modules.generation.sampling import (
9
prepare_sampling_params)
10
from neuronx_distributed_inference.modules.lora_serving import (
@@ -17,7 +19,7 @@ from vllm.model_executor import SamplingMetadata
19
from vllm.model_executor.layers.sampler import SamplerOutput
20
from vllm.model_executor.model_loader.neuronx_distributed import (
21
_get_model_architecture, get_neuron_model)
20
-from vllm.platforms import current_platform
22
+from vllm.multimodal import MultiModalKwargs
23
from vllm.sequence import IntermediateTensors, SequenceGroupMetadata
24
from vllm.worker.neuron_model_runner import (ModelInputForNeuron,
25
NeuronModelRunner)
@@ -121,42 +123,28 @@ class NeuronxDistributedModelRunner(NeuronModelRunner):
123
sampling_params = self.get_nxd_sampling_params(
124
model_input.sampling_metadata)
125
124
- if model_input.multi_modal_kwargs.get('image') is not None:
125
- pixel_values = []
126
- aspect_ratios = []
127
- num_chunks = []
128
- has_image = []
129
- for multi_modal_input in model_input.multi_modal_kwargs.get(
130
- 'image'):
131
- image_tensors = self.get_multi_modal_data_neuron(
132
- multi_modal_input.squeeze(0))
133
- pixel_values.append(image_tensors[0])
134
- aspect_ratios.append(image_tensors[1])
135
- num_chunks.append(image_tensors[2])
136
- has_image.append(image_tensors[3])
137
-
138
- pixel_values = torch.cat(pixel_values, dim=0)
139
- aspect_ratios = torch.cat(aspect_ratios, dim=0)
140
- num_chunks = torch.cat(num_chunks, dim=0)
141
- has_image = torch.cat(has_image, dim=0)
142
-
126
+ if model_input.multi_modal_kwargs.get('pixel_values') is not None:
127
hidden_states = self.model(
128
input_ids=model_input.input_tokens,
129
positions=model_input.input_positions,
130
seq_ids=model_input.input_block_ids,
147
- pixel_values=pixel_values,
148
- aspect_ratios=aspect_ratios,
131
+ pixel_values=model_input.multi_modal_kwargs.get(
132
+ 'pixel_values'),
133
+ aspect_ratios=model_input.multi_modal_kwargs.get(
134
+ 'aspect_ratios'),
135
sampling_params=sampling_params,
150
- num_chunks=num_chunks,
151
- has_image=has_image,
136
+ num_chunks=model_input.multi_modal_kwargs.get('num_chunks'),
137
+ has_image=model_input.multi_modal_kwargs.get(
138
+ 'has_image').squeeze(1),
139
)
140
else:
154
- empty_pixel_values = torch.zeros([1, 1, 4, 3, 560, 560],
141
+ bs = model_input.input_tokens.shape[0] if (model_input.input_tokens
142
+ is not None) else 1
143
+ empty_pixel_values = torch.zeros([bs, 1, 4, 3, 560, 560],
144
dtype=torch.bfloat16)
156
- empty_aspect_ratios = torch.ones([1, 1, 2], dtype=torch.int64)
157
- num_chunks = torch.tensor([[1]
158
- ]) # dummy num_chunks, will not be used
159
- has_image = torch.tensor([0])
145
+ empty_aspect_ratios = torch.ones([bs, 1, 2], dtype=torch.int64)
146
+ num_chunks = torch.zeros((bs, 1), dtype=torch.int32)
147
+ has_image = torch.zeros([bs], dtype=torch.int32)
148
hidden_states = self.model(
149
input_ids=model_input.input_tokens,
150
positions=model_input.input_positions,
@@ -175,6 +163,27 @@ class NeuronxDistributedModelRunner(NeuronModelRunner):
163
164
return [output]
165
166
+ def process_multi_modal_data_neuron(self, mm_data):
167
+ # Neuron uses aspect_ratios instead of aspect_ratio_ids
168
+ all_supported_aspect_ratios = get_all_supported_aspect_ratios(
169
+ self.model.config.vision_config.max_num_tiles)
170
+ aspect_ratio_ids = mm_data.get("aspect_ratio_ids")
171
+ mm_data["aspect_ratios"] = torch.tensor(
172
+ all_supported_aspect_ratios[aspect_ratio_ids]).unsqueeze(0)
173
+
174
+ # Neuron's num_chunks is HF's num_tiles
175
+ mm_data["num_chunks"] = mm_data.get("num_tiles")
176
+
177
+ # Input has an image if it has pixel_values
178
+ bs = mm_data["num_chunks"].shape[0]
179
+ pixel_values = mm_data.get("pixel_values")
180
+ if pixel_values is not None and not torch.all(pixel_values == 0):
181
+ mm_data["has_image"] = torch.ones(bs)
182
+
183
+ else:
184
+ mm_data["has_image"] = torch.zeros(bs)
185
+ return mm_data
186
+
187
def _get_lora_adapter_ids(self, seq_group_metadata_list):
188
# set LoRA adapter IDs for multi-lora serving
189
batch_size = len(seq_group_metadata_list)
@@ -200,7 +209,6 @@ class NeuronxDistributedModelRunner(NeuronModelRunner):
209
virtual_engine: int = 0,
210
finished_requests_ids: Optional[List[str]] = None
211
) -> ModelInputForNeuron:
203
- multi_modal_kwargs = None
212
# NOTE: We assume that all sequences in the group are all prompts or
213
# all decodes.
214
is_prompt = seq_group_metadata_list[0].is_prompt
@@ -223,6 +231,14 @@ class NeuronxDistributedModelRunner(NeuronModelRunner):
231
sampling_params.top_p = top_p
232
sampling_params.temperature = temperature
233
234
+ # we need multi_modal_data for later tokens as well
235
+ multi_modal_kwargs_list: List[MultiModalKwargs] = []
236
+ for seq_group_metadata in seq_group_metadata_list:
237
+ mm_data = seq_group_metadata.multi_modal_data
238
+ if mm_data:
239
+ multi_modal_kwargs_list.append(mm_data)
240
+ multi_modal_kwargs = MultiModalKwargs.batch(multi_modal_kwargs_list)
241
+
242
lora_adapter_ids = self._get_lora_adapter_ids(seq_group_metadata_list)
243
244
sampling_metadata = SamplingMetadata.prepare(
@@ -236,18 +252,6 @@ class NeuronxDistributedModelRunner(NeuronModelRunner):
252
self.pin_memory,
253
generators=self.get_generators(finished_requests_ids))
254
239
- if current_platform.use_transformers_neuronx(
240
- ) and not self._on_device_sampling_disabled:
241
- # Once the request IDs are changed in current iteration, we will
242
- # update the on-device sampling parameters.
243
- current_batch_request_ids = [
244
- seq_group_meta_data.request_id
245
- for seq_group_meta_data in seq_group_metadata_list
246
- ]
247
- if current_batch_request_ids != self._previous_batch_request_ids:
248
- self._update_neuron_sampling_params(seq_group_metadata_list)
249
- self._previous_batch_request_ids = current_batch_request_ids
250
-
255
return ModelInputForNeuron(input_tokens=input_tokens,
256
input_positions=input_positions,
257
input_block_ids=input_block_ids,