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
[VLM] Add PP support and fix GPTQ inference for Ovis models (#18958)
Signed-off-by: isotr0py <2037008807@qq.com> Signed-off-by: Isotr0py <2037008807@qq.com>
Isotr0py committed
May 31, 2025 at 01:11 UTC
5a8641638a5ec6cf089490d6b7f474b135ff590e
5 files changed
+145
-91
docs/models/supported_models.md
+1
-1
@@ -538,7 +538,7 @@ Specified using `--task generate`.
538
| `MllamaForConditionalGeneration` | Llama 3.2 | T + I<sup>+</sup> | `meta-llama/Llama-3.2-90B-Vision-Instruct`, `meta-llama/Llama-3.2-11B-Vision`, etc. | | | |
539
| `MolmoForCausalLM` | Molmo | T + I<sup>+</sup> | `allenai/Molmo-7B-D-0924`, `allenai/Molmo-7B-O-0924`, etc. | ✅︎ | ✅︎ | ✅︎ |
540
| `NVLM_D_Model` | NVLM-D 1.0 | T + I<sup>+</sup> | `nvidia/NVLM-D-72B`, etc. | | ✅︎ | ✅︎ |
541
-| `Ovis` | Ovis2, Ovis1.6 | T + I<sup>+</sup> | `AIDC-AI/Ovis2-1B`, `AIDC-AI/Ovis1.6-Llama3.2-3B`, etc. | | | ✅︎ |
541
+| `Ovis` | Ovis2, Ovis1.6 | T + I<sup>+</sup> | `AIDC-AI/Ovis2-1B`, `AIDC-AI/Ovis1.6-Llama3.2-3B`, etc. | | ✅︎ | ✅︎ |
542
| `PaliGemmaForConditionalGeneration` | PaliGemma, PaliGemma 2 | T + I<sup>E</sup> | `google/paligemma-3b-pt-224`, `google/paligemma-3b-mix-224`, `google/paligemma2-3b-ft-docci-448`, etc. | | ✅︎ | ⚠️ |
543
| `Phi3VForCausalLM` | Phi-3-Vision, Phi-3.5-Vision | T + I<sup>E+</sup> | `microsoft/Phi-3-vision-128k-instruct`, `microsoft/Phi-3.5-vision-instruct`, etc. | | ✅︎ | ✅︎ |
544
| `Phi4MMForCausalLM` | Phi-4-multimodal | T + I<sup>+</sup> / T + A<sup>+</sup> / I<sup>+</sup> + A<sup>+</sup> | `microsoft/Phi-4-multimodal-instruct`, etc. | ✅︎ | ✅︎ | ✅︎ |
tests/distributed/test_pipeline_parallel.py
+1
@@ -227,6 +227,7 @@ MULTIMODAL_MODELS = {
227
"llava-hf/llava-onevision-qwen2-0.5b-ov-hf": PPTestSettings.fast(),
228
"openbmb/MiniCPM-Llama3-V-2_5": PPTestSettings.fast(),
229
"allenai/Molmo-7B-D-0924": PPTestSettings.fast(),
230
+ "AIDC-AI/Ovis2-1B": PPTestSettings.fast(),
231
"microsoft/Phi-3.5-vision-instruct": PPTestSettings.fast(),
232
"mistralai/Pixtral-12B-2409": PPTestSettings.fast(load_format="dummy"),
233
"Qwen/Qwen-VL-Chat": PPTestSettings.fast(),
vllm/model_executor/models/aimv2.py
+120
-74
@@ -2,16 +2,23 @@
2
3
# A modified implementation of the AIMv2 Transformer
4
# inserted here also the image tokenizer used by Ovis2
5
+from collections.abc import Iterable
6
from typing import Optional
7
8
import torch
9
import torch.nn as nn
9
-from torch.nn import functional as F
10
11
+from vllm.attention.layer import MultiHeadAttention
12
+from vllm.distributed import get_tensor_model_parallel_world_size
13
+from vllm.distributed.utils import divide
14
+from vllm.model_executor.layers.activation import SiluAndMul
15
from vllm.model_executor.layers.layernorm import RMSNorm
12
-from vllm.model_executor.layers.linear import ReplicatedLinear
16
+from vllm.model_executor.layers.linear import (MergedColumnParallelLinear,
17
+ QKVParallelLinear,
18
+ RowParallelLinear)
19
from vllm.model_executor.layers.quantization.base_config import (
20
QuantizationConfig)
21
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
22
from vllm.transformers_utils.configs.ovis import AIMv2Config
23
24
@@ -24,29 +31,27 @@ class AIMv2SwiGLUFFN(nn.Module):
31
in_features = config.hidden_size
32
bias = config.use_bias
33
27
- # TODO(Isotr0py): investigate if we can add TP to visual tokenizer
28
- self.fc1 = ReplicatedLinear(in_features,
29
- hidden_features,
30
- bias=bias,
31
- quant_config=quant_config,
32
- prefix=f"{prefix}.fc1")
33
- self.fc2 = ReplicatedLinear(hidden_features,
34
- in_features,
35
- bias=bias,
36
- quant_config=quant_config,
37
- prefix=f"{prefix}.fc2")
38
- self.fc3 = ReplicatedLinear(in_features,
39
- hidden_features,
40
- bias=bias,
41
- quant_config=quant_config,
42
- prefix=f"{prefix}.fc3")
34
+ self.fc13 = MergedColumnParallelLinear(
35
+ in_features,
36
+ [hidden_features] * 2,
37
+ bias=bias,
38
+ quant_config=quant_config,
39
+ prefix=f"{prefix}.fc13",
40
+ )
41
+ self.fc2 = RowParallelLinear(
42
+ input_size=hidden_features,
43
+ output_size=in_features,
44
+ bias=bias,
45
+ quant_config=quant_config,
46
+ prefix=f"{prefix}.fc2",
47
+ )
48
+ self.act_fn = SiluAndMul()
49
50
def forward(self, x: torch.Tensor) -> torch.Tensor:
45
- x_parallel, _ = self.fc1(x)
46
- gate, _ = self.fc3(x)
47
- x_parallel = F.silu(x_parallel) * gate
48
- out, _ = self.fc2(x_parallel)
49
- return out
51
+ x, _ = self.fc13(x)
52
+ x = self.act_fn(x)
53
+ x, _ = self.fc2(x)
54
+ return x
55
56
57
class AIMv2PatchEmbed(nn.Module):
@@ -90,39 +95,45 @@ class AIMv2Attention(nn.Module):
95
def __init__(self, config: AIMv2Config, quant_config: QuantizationConfig,
96
prefix: str):
97
super().__init__()
93
- dim = config.hidden_size
94
-
95
- # TODO(Isotr0py): investigate if we can add TP to visual tokenizer
98
+ self.config = config
99
+ self.embed_dim = config.hidden_size
100
self.num_heads = config.num_attention_heads
97
- self.qkv = ReplicatedLinear(dim, dim * 3, bias=config.qkv_bias)
98
- # self.qkv = QKVParallelLinear(
99
- # hidden_size=dim,
100
- # head_size=dim // config.num_attention_heads,
101
- # total_num_heads=config.num_attention_heads,
102
- # bias=config.qkv_bias,
103
- # quant_config=quant_config,
104
- # prefix=f"{prefix}.qkv")
105
- self.proj = ReplicatedLinear(dim, dim, bias=config.use_bias)
106
- # self.proj = RowParallelLinear(input_size=dim,
107
- # output_size=dim,
108
- # bias = config.use_bias,
109
- # quant_config=quant_config,
110
- # prefix=f"{prefix}.proj")
111
-
112
- def forward( # todo might implement multiple attn implementations
113
- self,
114
- x: torch.Tensor,
115
- mask: Optional[torch.Tensor] = None) -> torch.Tensor:
116
- B, N, C = x.shape
117
- qkv, _ = self.qkv(x)
101
+ self.head_dim = self.embed_dim // self.num_heads
102
+ if self.head_dim * self.num_heads != self.embed_dim:
103
+ raise ValueError(
104
+ "embed_dim must be divisible by num_heads "
105
+ f"(got `embed_dim`: {self.embed_dim} and `num_heads`:"
106
+ f" {self.num_heads}).")
107
+ self.scale = self.head_dim**-0.5
108
+
109
+ self.qkv = QKVParallelLinear(
110
+ hidden_size=self.embed_dim,
111
+ head_size=self.head_dim,
112
+ total_num_heads=self.num_heads,
113
+ bias=config.qkv_bias,
114
+ quant_config=quant_config,
115
+ prefix=f"{prefix}.qkv",
116
+ )
117
+
118
+ self.proj = RowParallelLinear(
119
+ input_size=self.embed_dim,
120
+ output_size=self.embed_dim,
121
+ bias=config.use_bias,
122
+ quant_config=quant_config,
123
+ prefix=f"{prefix}.proj",
124
+ )
125
+
126
+ self.tp_size = get_tensor_model_parallel_world_size()
127
+ self.num_heads_per_partition = divide(self.num_heads, self.tp_size)
128
119
- qkv = qkv.reshape(B, N, 3, self.num_heads,
120
- C // self.num_heads).permute(2, 0, 3, 1, 4)
129
+ self.attn = MultiHeadAttention(self.num_heads_per_partition,
130
+ self.head_dim, self.scale)
131
122
- q, k, v = qkv.unbind(0)
132
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
133
+ qkv, _ = self.qkv(x)
134
+ q, k, v = qkv.chunk(3, dim=-1)
135
124
- x = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
125
- x = x.transpose(1, 2).contiguous().reshape(B, N, C)
136
+ x = self.attn(q, k, v)
137
x, _ = self.proj(x)
138
return x
139
@@ -141,37 +152,40 @@ class AIMv2Block(nn.Module):
152
prefix=f"{prefix}.mlp")
153
self.norm_2 = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
154
144
- def forward(self,
145
- x: torch.Tensor,
146
- mask: Optional[torch.Tensor] = None) -> torch.Tensor:
147
- x = x + self.attn(self.norm_1.forward_native(x), mask)
155
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
156
+ x = x + self.attn(self.norm_1.forward_native(x))
157
x = x + self.mlp(self.norm_2.forward_native(x))
158
return x
159
160
161
class AIMv2Transformer(nn.Module):
162
154
- def __init__(self, config: AIMv2Config, quant_config: QuantizationConfig,
155
- prefix: str):
163
+ def __init__(
164
+ self,
165
+ config: AIMv2Config,
166
+ quant_config: QuantizationConfig,
167
+ *,
168
+ require_post_norm: Optional[bool] = None,
169
+ prefix: str = "",
170
+ ):
171
super().__init__()
172
173
self.blocks = nn.ModuleList([
174
AIMv2Block(config, quant_config, prefix=f"{prefix}.blocks.{i}")
175
for i in range(config.num_hidden_layers)
176
])
162
- self.post_trunk_norm = RMSNorm(config.hidden_size,
163
- eps=config.rms_norm_eps)
177
+ if require_post_norm:
178
+ self.post_trunk_norm = RMSNorm(config.hidden_size,
179
+ eps=config.rms_norm_eps)
180
+ else:
181
+ self.post_trunk_norm = None
182
165
- def forward(
166
- self,
167
- tokens: torch.Tensor,
168
- mask: Optional[torch.Tensor] = None,
169
- ) -> torch.Tensor:
183
+ def forward(self, tokens: torch.Tensor) -> torch.Tensor:
184
# they take the -1 as the ref embeddings, like a clip skip
185
for block in self.blocks:
172
- tokens = block(tokens, mask)
173
- # NO NORM IN THE OG IMPLEMENTATION
174
- # tokens = self.post_trunk_norm(tokens)
186
+ tokens = block(tokens)
187
+ if self.post_trunk_norm is not None:
188
+ tokens = self.post_trunk_norm(tokens)
189
return tokens
190
191
@@ -180,20 +194,52 @@ class AIMv2Model(torch.nn.Module):
194
def __init__(self,
195
config: AIMv2Config,
196
quant_config: QuantizationConfig,
197
+ *,
198
+ require_post_norm: Optional[bool] = None,
199
prefix: str = ""):
200
super().__init__()
201
self.preprocessor = AIMv2ViTPreprocessor(config)
202
self.trunk = AIMv2Transformer(config,
203
quant_config=quant_config,
204
+ require_post_norm=require_post_norm,
205
prefix=f"{prefix}.trunk")
206
190
- def forward(
191
- self,
192
- pixel_values: torch.Tensor,
193
- mask: Optional[torch.Tensor] = None,
194
- ) -> torch.Tensor:
207
+ def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
208
209
x = self.preprocessor(pixel_values)
197
- x = self.trunk(x, mask)
210
+ x = self.trunk(x)
211
212
return x
213
+
214
+ def load_weights(self, weights: Iterable[tuple[str,
215
+ torch.Tensor]]) -> set[str]:
216
+ stacked_params_mapping = [
217
+ # (param_name, shard_name, shard_id)
218
+ (".fc13", ".fc1", 0),
219
+ (".fc13", ".fc3", 1),
220
+ ]
221
+ params_dict = dict(self.named_parameters())
222
+ loaded_params: set[str] = set()
223
+
224
+ for name, loaded_weight in weights:
225
+ # post_layernorm is optional in SiglipVisionModel
226
+ if (name.startswith("trunk.post_trunk_norm")
227
+ and self.trunk.post_trunk_norm is None):
228
+ continue
229
+
230
+ for (param_name, weight_name, shard_id) in stacked_params_mapping:
231
+ if weight_name not in name:
232
+ continue
233
+ name = name.replace(weight_name, param_name)
234
+
235
+ param = params_dict[name]
236
+ weight_loader = param.weight_loader
237
+ weight_loader(param, loaded_weight, shard_id)
238
+ break
239
+ else:
240
+ param = params_dict[name]
241
+ weight_loader = getattr(param, "weight_loader",
242
+ default_weight_loader)
243
+ weight_loader(param, loaded_weight)
244
+ loaded_params.add(name)
245
+ return loaded_params
vllm/model_executor/models/clip.py
-5
@@ -106,7 +106,6 @@ class CLIPAttention(nn.Module):
106
f"(got `embed_dim`: {self.embed_dim} and `num_heads`:"
107
f" {self.num_heads}).")
108
self.scale = self.head_dim**-0.5
109
- self.dropout = config.attention_dropout
109
110
self.qkv_proj = QKVParallelLinear(
111
hidden_size=self.embed_dim,
@@ -129,10 +128,6 @@ class CLIPAttention(nn.Module):
128
self.attn = MultiHeadAttention(self.num_heads_per_partition,
129
self.head_dim, self.scale)
130
132
- def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
133
- return tensor.view(bsz, seq_len, self.num_heads,
134
- self.head_dim).transpose(1, 2).contiguous()
135
-
131
def forward(
132
self,
133
hidden_states: torch.Tensor,
vllm/model_executor/models/ovis.py
+23
-11
@@ -30,6 +30,9 @@ from vllm.config import VllmConfig
30
from vllm.model_executor.layers.linear import ReplicatedLinear
31
from vllm.model_executor.layers.quantization.base_config import (
32
QuantizationConfig)
33
+from vllm.model_executor.layers.quantization.gptq import GPTQConfig
34
+from vllm.model_executor.layers.quantization.gptq_marlin import (
35
+ GPTQMarlinConfig)
36
from vllm.model_executor.models.aimv2 import AIMv2Model
37
from vllm.model_executor.models.siglip import SiglipVisionModel
38
from vllm.model_executor.models.utils import (AutoWeightsLoader, flatten_bn,
@@ -48,7 +51,7 @@ from vllm.transformers_utils.configs.ovis import (BaseVisualTokenizerConfig,
51
OvisConfig)
52
from vllm.transformers_utils.processors.ovis import OvisProcessor
53
51
-from .interfaces import MultiModalEmbeddings, SupportsMultiModal
54
+from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
55
from .utils import merge_multimodal_embeddings
56
57
# Cannot find the following number from hf config.
@@ -106,12 +109,14 @@ class VisualTokenizer(torch.nn.Module):
109
config: BaseVisualTokenizerConfig,
110
quant_config: Optional[QuantizationConfig] = None,
111
prefix: str = "",
109
- ):
112
+ ) -> nn.Module:
113
model_type = config.backbone_config.model_type
114
if model_type == "aimv2":
115
+ # No post rms_norm in Ovis2's AIMv2 ViT.
116
return AIMv2Model(
117
config=config.backbone_config,
118
quant_config=quant_config,
119
+ require_post_norm=False,
120
prefix=prefix,
121
)
122
elif model_type == "siglip_vision_model":
@@ -124,14 +129,14 @@ class VisualTokenizer(torch.nn.Module):
129
f"Unsupported visual tokenizer model_type: {model_type}")
130
131
@property
127
- def dtype(self):
132
+ def dtype(self) -> torch.dtype:
133
return next(self.head.parameters()).dtype
134
135
@property
131
- def device(self):
136
+ def device(self) -> torch.device:
137
return next(self.head.parameters()).device
138
134
- def tokenize(self, logits):
139
+ def tokenize(self, logits: torch.Tensor) -> torch.Tensor:
140
if self.config.tokenize_function == 'softmax':
141
tokens = softmax(logits, dim=-1)
142
elif self.config.tokenize_function == 'gumbel_argmax':
@@ -144,7 +149,7 @@ class VisualTokenizer(torch.nn.Module):
149
f'or st_argmax, but got {self.config.tokenize_function}')
150
return tokens
151
147
- def encode(self, pixel_values):
152
+ def encode(self, pixel_values: torch.Tensor) -> torch.Tensor:
153
features = self.backbone(pixel_values)
154
if self.config.drop_cls_token:
155
features = features[:, 1:, :]
@@ -395,7 +400,7 @@ class OvisMultiModalProcessor(BaseMultiModalProcessor[OvisProcessingInfo]):
400
@MULTIMODAL_REGISTRY.register_processor(OvisMultiModalProcessor,
401
info=OvisProcessingInfo,
402
dummy_inputs=OvisDummyInputsBuilder)
398
-class Ovis(nn.Module, SupportsMultiModal):
403
+class Ovis(nn.Module, SupportsMultiModal, SupportsPP):
404
405
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
406
super().__init__()
@@ -410,7 +415,7 @@ class Ovis(nn.Module, SupportsMultiModal):
415
416
self.visual_tokenizer = VisualTokenizer(
417
config=config.visual_tokenizer_config,
413
- quant_config=quant_config,
418
+ quant_config=self._maybe_ignore_quant_config(quant_config),
419
prefix=f"{prefix}.visual_tokenizer",
420
)
421
@@ -421,9 +426,16 @@ class Ovis(nn.Module, SupportsMultiModal):
426
text_model_type = self.config.get_text_config().model_type
427
self.image_pad_token_id = IMAGE_PAD_TOKEN_ID_MAP[text_model_type]
428
424
- # TODO(Isotr0py): PP support
425
- # self.make_empty_intermediate_tensors = (
426
- # self.language_model.make_empty_intermediate_tensors)
429
+ self.make_empty_intermediate_tensors = (
430
+ self.get_language_model().make_empty_intermediate_tensors)
431
+
432
+ def _maybe_ignore_quant_config(self, quant_config: QuantizationConfig):
433
+ # GPTQ configs do not have a list of ignored modules, however AutoGPTQ
434
+ # seems to avoid vision encoder sections for some models.
435
+ # See: https://huggingface.co/AIDC-AI/Ovis2-2B-GPTQ-Int4
436
+ if isinstance(quant_config, (GPTQConfig, GPTQMarlinConfig)):
437
+ return None
438
+ return quant_config
439
440
def _parse_and_validate_image_input(
441
self, **kwargs: object) -> Optional[OvisImagePatchInputs]: