@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 1adeb3b84c
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 11, 11:15 UTC

[New Model] BAGEL support (AR only) (#28439)

Signed-off-by: princepride <wangzhipeng628@gmail.com> Signed-off-by: 汪志鹏 <wangzhipeng628@gmail.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>

汪志鹏 committed Dec 15, 2025 at 14:58 UTC 1adeb3b84c2dcf776b13a9933904c6214c3fe745
11 files changed +777
docs/models/supported_models.md
+1
@@ -661,6 +661,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
661 | `AriaForConditionalGeneration` | Aria | T + I<sup>+</sup> | `rhymes-ai/Aria` | | |
662 | `AudioFlamingo3ForConditionalGeneration` | AudioFlamingo3 | T + A<sup>+</sup> | `nvidia/audio-flamingo-3-hf`, `nvidia/music-flamingo-hf` | ✅︎ | ✅︎ |
663 | `AyaVisionForConditionalGeneration` | Aya Vision | T + I<sup>+</sup> | `CohereLabs/aya-vision-8b`, `CohereLabs/aya-vision-32b`, etc. | | ✅︎ |
664 +| `BagelForConditionalGeneration` | BAGEL | T + I<sup>+</sup> | `ByteDance-Seed/BAGEL-7B-MoT` | ✅︎ | ✅︎ |
665 | `BeeForConditionalGeneration` | Bee-8B | T + I<sup>E+</sup> | `Open-Bee/Bee-8B-RL`, `Open-Bee/Bee-8B-SFT` | | ✅︎ |
666 | `Blip2ForConditionalGeneration` | BLIP-2 | T + I<sup>E</sup> | `Salesforce/blip2-opt-2.7b`, `Salesforce/blip2-opt-6.7b`, etc. | | ✅︎ |
667 | `ChameleonForConditionalGeneration` | Chameleon | T + I | `facebook/chameleon-7b`, etc. | | ✅︎ |
examples/offline_inference/vision_language.py
+27
@@ -118,6 +118,32 @@ def run_bee(questions: list[str], modality: str) -> ModelRequestData:
118 )
119
120
121 +def run_bagel(questions: list[str], modality: str) -> ModelRequestData:
122 + assert modality == "image"
123 + model_name = "ByteDance-Seed/BAGEL-7B-MoT"
124 +
125 + engine_args = EngineArgs(
126 + model=model_name,
127 + trust_remote_code=True,
128 + max_model_len=8192,
129 + max_num_seqs=2,
130 + limit_mm_per_prompt={modality: 1},
131 + )
132 +
133 + prompts = [
134 + (
135 + f"<|im_start|>user\n<|image_pad|>\n{question}<|im_end|>\n"
136 + f"<|im_start|>assistant\n"
137 + )
138 + for question in questions
139 + ]
140 +
141 + return ModelRequestData(
142 + engine_args=engine_args,
143 + prompts=prompts,
144 + )
145 +
146 +
147 # BLIP-2
148 def run_blip2(questions: list[str], modality: str) -> ModelRequestData:
149 assert modality == "image"
@@ -1832,6 +1858,7 @@ def run_tarsier2(questions: list[str], modality: str) -> ModelRequestData:
1858 model_example_map = {
1859 "aria": run_aria,
1860 "aya_vision": run_aya_vision,
1861 + "bagel": run_bagel,
1862 "bee": run_bee,
1863 "blip-2": run_blip2,
1864 "chameleon": run_chameleon,
tests/models/registry.py
+1
@@ -582,6 +582,7 @@ _MULTIMODAL_EXAMPLE_MODELS = {
582 "nvidia/audio-flamingo-3-hf", min_transformers_version="5.0.0.dev"
583 ),
584 "AyaVisionForConditionalGeneration": _HfExamplesInfo("CohereLabs/aya-vision-8b"),
585 + "BagelForConditionalGeneration": _HfExamplesInfo("ByteDance-Seed/BAGEL-7B-MoT"),
586 "BeeForConditionalGeneration": _HfExamplesInfo(
587 "Open-Bee/Bee-8B-RL",
588 trust_remote_code=True,
vllm/model_executor/models/bagel.py new
+584
@@ -0,0 +1,584 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +# Copyright 2025 Bytedance Ltd. and/or its affiliates.
4 +"""Inference-only BAGEL model compatible with HuggingFace weights.
5 +
6 +BAGEL is a unified multimodal model for image understanding and generation.
7 +For vLLM, we focus on the image understanding (vision-to-text) capabilities.
8 +"""
9 +
10 +from collections.abc import Iterable, Mapping, Sequence
11 +from typing import Any, Literal, TypeAlias
12 +
13 +import torch
14 +import torch.nn as nn
15 +
16 +from vllm.config import VllmConfig
17 +from vllm.config.multimodal import BaseDummyOptions
18 +from vllm.logger import init_logger
19 +from vllm.model_executor.layers.activation import get_act_fn
20 +from vllm.model_executor.layers.linear import (
21 + ColumnParallelLinear,
22 + RowParallelLinear,
23 +)
24 +from vllm.model_executor.layers.quantization import QuantizationConfig
25 +from vllm.multimodal import MULTIMODAL_REGISTRY
26 +from vllm.multimodal.inputs import (
27 + MultiModalDataDict,
28 + MultiModalFieldConfig,
29 + MultiModalKwargsItems,
30 +)
31 +from vllm.multimodal.parse import MultiModalDataItems
32 +from vllm.multimodal.processing import (
33 + BaseMultiModalProcessor,
34 + BaseProcessingInfo,
35 + PromptReplacement,
36 +)
37 +from vllm.multimodal.profiling import BaseDummyInputsBuilder
38 +from vllm.sequence import IntermediateTensors
39 +from vllm.transformers_utils.processors.bagel import BagelProcessor
40 +from vllm.utils.tensor_schema import TensorSchema
41 +
42 +from .interfaces import (
43 + MultiModalEmbeddings,
44 + SupportsLoRA,
45 + SupportsMultiModal,
46 + SupportsPP,
47 +)
48 +from .siglip import SiglipVisionModel
49 +from .utils import (
50 + AutoWeightsLoader,
51 + WeightsMapper,
52 + init_vllm_registered_model,
53 + maybe_prefix,
54 +)
55 +
56 +logger = init_logger(__name__)
57 +
58 +
59 +class BagelImagePixelInputs(TensorSchema):
60 + """
61 + Dimensions:
62 + - bn: Batch size * number of images
63 + - c: Number of channels (3)
64 + - h: Height of each image
65 + - w: Width of each image
66 + """
67 +
68 + type: Literal["pixel_values"]
69 + pixel_values: torch.Tensor # Shape: (bn, 3, h, w)
70 +
71 +
72 +BagelImageInputs: TypeAlias = BagelImagePixelInputs
73 +
74 +
75 +class BagelVisionMLP(nn.Module):
76 + """MLP connector for vision features."""
77 +
78 + def __init__(
79 + self,
80 + in_features: int,
81 + hidden_features: int,
82 + out_features: int,
83 + act_layer: str = "gelu_pytorch_tanh",
84 + quant_config: QuantizationConfig | None = None,
85 + prefix: str = "",
86 + ):
87 + super().__init__()
88 + self.fc1 = ColumnParallelLinear(
89 + in_features,
90 + hidden_features,
91 + bias=True,
92 + quant_config=quant_config,
93 + prefix=f"{prefix}.fc1",
94 + )
95 + self.act = get_act_fn(act_layer)
96 + self.fc2 = RowParallelLinear(
97 + hidden_features,
98 + out_features,
99 + bias=True,
100 + quant_config=quant_config,
101 + prefix=f"{prefix}.fc2",
102 + )
103 +
104 + def forward(self, x: torch.Tensor) -> torch.Tensor:
105 + x, _ = self.fc1(x)
106 + x = self.act(x)
107 + x, _ = self.fc2(x)
108 + return x
109 +
110 +
111 +class PositionEmbedding(nn.Module):
112 + """2D position embedding for vision tokens using sin-cos embeddings."""
113 +
114 + def __init__(self, max_num_patch_per_side: int, hidden_size: int):
115 + super().__init__()
116 + self.max_num_patch_per_side = max_num_patch_per_side
117 + self.hidden_size = hidden_size
118 +
119 + # Create learnable 2D position embeddings (frozen sin-cos)
120 + pos_embed = self._get_2d_sincos_pos_embed(hidden_size, max_num_patch_per_side)
121 + self.register_buffer(
122 + "pos_embed",
123 + torch.from_numpy(pos_embed).float(),
124 + persistent=False,
125 + )
126 +
127 + @staticmethod
128 + def _get_2d_sincos_pos_embed(embed_dim: int, grid_size: int):
129 + """Generate 2D sin-cos position embeddings."""
130 + import numpy as np
131 +
132 + grid_h = np.arange(grid_size, dtype=np.float32)
133 + grid_w = np.arange(grid_size, dtype=np.float32)
134 + grid = np.meshgrid(grid_w, grid_h) # w goes first
135 + grid = np.stack(grid, axis=0)
136 + grid = grid.reshape([2, 1, grid_size, grid_size])
137 + pos_embed = PositionEmbedding._get_2d_sincos_pos_embed_from_grid(
138 + embed_dim, grid
139 + )
140 + return pos_embed
141 +
142 + @staticmethod
143 + def _get_2d_sincos_pos_embed_from_grid(embed_dim: int, grid):
144 + """Generate 2D sin-cos position embeddings from grid."""
145 + import numpy as np
146 +
147 + assert embed_dim % 2 == 0
148 + # use half of dimensions to encode grid_h
149 + emb_h = PositionEmbedding._get_1d_sincos_pos_embed_from_grid(
150 + embed_dim // 2, grid[0]
151 + )
152 + emb_w = PositionEmbedding._get_1d_sincos_pos_embed_from_grid(
153 + embed_dim // 2, grid[1]
154 + )
155 + emb = np.concatenate([emb_h, emb_w], axis=1)
156 + return emb
157 +
158 + @staticmethod
159 + def _get_1d_sincos_pos_embed_from_grid(embed_dim: int, pos):
160 + """Generate 1D sin-cos position embeddings."""
161 + import numpy as np
162 +
163 + assert embed_dim % 2 == 0
164 + omega = np.arange(embed_dim // 2, dtype=np.float64)
165 + omega /= embed_dim / 2.0
166 + omega = 1.0 / 10000**omega
167 +
168 + pos = pos.reshape(-1)
169 + out = np.einsum("m,d->md", pos, omega)
170 +
171 + emb_sin = np.sin(out)
172 + emb_cos = np.cos(out)
173 + emb = np.concatenate([emb_sin, emb_cos], axis=1)
174 + return emb
175 +
176 + def forward(self, position_ids: torch.Tensor) -> torch.Tensor:
177 + """
178 + Args:
179 + position_ids: Flattened position IDs, shape (N,) where each ID
180 + corresponds to a position in the flattened grid
181 + Returns:
182 + Position embeddings of shape (N, hidden_size)
183 + """
184 + # Ensure position_ids are on the same device as pos_embed
185 + position_ids = position_ids.to(self.pos_embed.device)
186 + return self.pos_embed[position_ids]
187 +
188 +
189 +class BagelProcessingInfo(BaseProcessingInfo):
190 + """Processing information for BAGEL model."""
191 +
192 + def get_hf_processor(self, **kwargs: object) -> BagelProcessor:
193 + from vllm.transformers_utils.processor import cached_get_image_processor
194 +
195 + image_processor = cached_get_image_processor(
196 + self.ctx.model_config.model,
197 + revision=self.ctx.model_config.revision,
198 + trust_remote_code=self.ctx.model_config.trust_remote_code,
199 + )
200 +
201 + tokenizer = self.get_tokenizer()
202 +
203 + return BagelProcessor(
204 + image_processor=image_processor,
205 + tokenizer=tokenizer,
206 + **kwargs,
207 + )
208 +
209 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:
210 + return {"image": None}
211 +
212 + def get_mm_max_tokens_per_item(
213 + self,
214 + seq_len: int,
215 + mm_counts: Mapping[str, int],
216 + ) -> Mapping[str, int]:
217 + hf_config = self.get_hf_config()
218 + # Calculate max tokens per image
219 + # For BAGEL: (vit_max_num_patch_per_side) ** 2
220 + max_num_patches = hf_config.vit_max_num_patch_per_side**2
221 + return {"image": max_num_patches}
222 +
223 + def get_num_image_tokens(
224 + self,
225 + *,
226 + image_width: int,
227 + image_height: int,
228 + ) -> int:
229 + hf_config = self.get_hf_config()
230 + vit_config = hf_config.vit_config
231 + patch_size = vit_config.patch_size
232 +
233 + # Calculate number of patches
234 + num_patches_h = image_height // patch_size
235 + num_patches_w = image_width // patch_size
236 + return num_patches_h * num_patches_w
237 +
238 +
239 +class BagelDummyInputsBuilder(BaseDummyInputsBuilder[BagelProcessingInfo]):
240 + """Build dummy inputs for BAGEL model profiling."""
241 +
242 + def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
243 + num_images = mm_counts.get("image", 0)
244 + # Use a simple placeholder for each image
245 + return "<|image_pad|>" * num_images
246 +
247 + def get_dummy_mm_data(
248 + self,
249 + seq_len: int,
250 + mm_counts: Mapping[str, int],
251 + mm_options: Mapping[str, BaseDummyOptions] | None = None,
252 + ) -> MultiModalDataDict:
253 + num_images = mm_counts.get("image", 0)
254 + hf_config = self.info.get_hf_config()
255 + vit_config = hf_config.vit_config
256 +
257 + # Use the configured image size
258 + image_size = vit_config.image_size
259 + image_overrides = mm_options.get("image") if mm_options else None
260 +
261 + return {
262 + "image": self._get_dummy_images(
263 + width=image_size,
264 + height=image_size,
265 + num_images=num_images,
266 + overrides=image_overrides,
267 + ),
268 + }
269 +
270 +
271 +class BagelMultiModalProcessor(BaseMultiModalProcessor[BagelProcessingInfo]):
272 + """Multimodal processor for BAGEL model."""
273 +
274 + def _hf_processor_applies_updates(
275 + self,
276 + prompt_text: str,
277 + mm_items: MultiModalDataItems,
278 + hf_processor_mm_kwargs: Mapping[str, object],
279 + tokenization_kwargs: Mapping[str, object],
280 + ) -> bool:
281 + return False
282 +
283 + def _get_prompt_updates(
284 + self,
285 + mm_items: MultiModalDataItems,
286 + hf_processor_mm_kwargs: Mapping[str, Any],
287 + out_mm_kwargs: MultiModalKwargsItems,
288 + ) -> Sequence[PromptReplacement]:
289 + """Replace image placeholders with the correct number of tokens."""
290 + hf_config = self.info.get_hf_config()
291 +
292 + # Get the tokenizer to look up the image token ID
293 + tokenizer = self.info.get_tokenizer()
294 + image_token_id = tokenizer.get_vocab().get("<|image_pad|>")
295 + if image_token_id is None:
296 + raise ValueError(
297 + "Image token '<|image_pad|>' not found in tokenizer vocabulary"
298 + )
299 +
300 + def get_replacement_bagel(item_idx: int):
301 + # For BAGEL, calculate number of tokens based on max patch size
302 + num_tokens = hf_config.vit_max_num_patch_per_side**2
303 + # Use the image token ID from tokenizer
304 + return [image_token_id] * num_tokens
305 +
306 + return [
307 + PromptReplacement(
308 + modality="image",
309 + target=[image_token_id],
310 + replacement=get_replacement_bagel,
311 + )
312 + ]
313 +
314 + def _get_mm_fields_config(
315 + self,
316 + hf_inputs: Any,
317 + hf_processor_mm_kwargs: Mapping[str, object],
318 + ) -> Mapping[str, MultiModalFieldConfig]:
319 + return {
320 + "pixel_values": MultiModalFieldConfig.batched("image"),
321 + }
322 +
323 +
324 +@MULTIMODAL_REGISTRY.register_processor(
325 + BagelMultiModalProcessor,
326 + info=BagelProcessingInfo,
327 + dummy_inputs=BagelDummyInputsBuilder,
328 +)
329 +class BagelForConditionalGeneration(
330 + nn.Module, SupportsMultiModal, SupportsLoRA, SupportsPP
331 +):
332 + """
333 + BAGEL: A unified multimodal model for image understanding and generation.
334 +
335 + For vLLM, we focus on the image understanding (vision-to-text) capabilities.
336 + The image generation part is not supported in vLLM.
337 + """
338 +
339 + # Weight mapping from HF to vLLM
340 + hf_to_vllm_mapper = WeightsMapper(
341 + orig_to_new_prefix={
342 + "language_model.": "language_model.",
343 + "vit_model.": "vit_model.",
344 + "connector.": "connector.",
345 + "vit_pos_embed.": "vit_pos_embed.",
346 + }
347 + )
348 +
349 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
350 + super().__init__()
351 +
352 + config = vllm_config.model_config.hf_config
353 + quant_config = vllm_config.quant_config
354 + multimodal_config = vllm_config.model_config.multimodal_config
355 +
356 + # Ensure we have a BagelConfig (check by name to handle trust_remote_code)
357 + # When trust_remote_code=True, the config comes from transformers_modules
358 + if type(config).__name__ != "BagelConfig":
359 + raise ValueError(
360 + f"Expected BagelConfig, got {type(config).__name__}. "
361 + "Make sure the model config is properly loaded."
362 + )
363 +
364 + self.config = config
365 + self.multimodal_config = multimodal_config
366 +
367 + # Initialize language model (Qwen2)
368 + # Pass the llm_config from BagelConfig to initialize Qwen2 properly
369 + self.language_model = init_vllm_registered_model(
370 + vllm_config=vllm_config,
371 + hf_config=config.llm_config,
372 + prefix=maybe_prefix(prefix, "language_model"),
373 + architectures=["Qwen2ForCausalLM"],
374 + )
375 +
376 + # Initialize vision model (SigLIP) if visual understanding is enabled
377 + if config.visual_und:
378 + # Fix vit_config: checkpoint has 26 layers (0-25) but config says 27
379 + # Also disable head as it's not in checkpoint
380 + vit_config = config.vit_config
381 + if vit_config.num_hidden_layers == 27:
382 + logger.warning(
383 + "Overriding vit_config.num_hidden_layers from 27 to 26 "
384 + "to match the Bagel model checkpoint."
385 + )
386 + vit_config.num_hidden_layers = 26
387 + if not hasattr(vit_config, "vision_use_head"):
388 + logger.warning(
389 + "Setting vit_config.vision_use_head to False as it is not "
390 + "present in the Bagel model checkpoint."
391 + )
392 + vit_config.vision_use_head = False
393 +
394 + self.vit_model = SiglipVisionModel(
395 + config=vit_config,
396 + quant_config=quant_config,
397 + prefix=maybe_prefix(prefix, "vit_model"),
398 + )
399 +
400 + # Initialize connector (MLP)
401 + vit_hidden_size = config.vit_config.hidden_size
402 + llm_hidden_size = config.llm_config.hidden_size
403 +
404 + self.connector = BagelVisionMLP(
405 + in_features=vit_hidden_size,
406 + hidden_features=llm_hidden_size,
407 + out_features=llm_hidden_size,
408 + act_layer=config.connector_act,
409 + quant_config=quant_config,
410 + prefix=maybe_prefix(prefix, "connector"),
411 + )
412 +
413 + # Position embedding for vision tokens
414 + self.vit_pos_embed = PositionEmbedding(
415 + max_num_patch_per_side=config.vit_max_num_patch_per_side,
416 + hidden_size=llm_hidden_size,
417 + )
418 + else:
419 + self.vit_model = None
420 + self.connector = None
421 + self.vit_pos_embed = None
422 +
423 + self.make_empty_intermediate_tensors = (
424 + self.language_model.make_empty_intermediate_tensors
425 + )
426 +
427 + def _parse_and_validate_image_input(
428 + self, **kwargs: object
429 + ) -> BagelImageInputs | None:
430 + pixel_values = kwargs.pop("pixel_values", None)
431 +
432 + if pixel_values is None:
433 + return None
434 +
435 + return BagelImagePixelInputs(
436 + type="pixel_values",
437 + pixel_values=pixel_values,
438 + )
439 +
440 + def _process_image_input(
441 + self, image_input: BagelImageInputs
442 + ) -> tuple[torch.Tensor, ...]:
443 + """Process image inputs through vision encoder and connector."""
444 + pixel_values = image_input["pixel_values"]
445 +
446 + # Handle potential extra batch dimension
447 + # Expected shape: (batch_size * num_images, 3, H, W)
448 + # But might receive: (batch_size, num_images, 3, H, W)
449 + if pixel_values.ndim == 5:
450 + # Flatten batch and num_images dimensions
451 + batch_size, num_images, channels, height, width = pixel_values.shape
452 + pixel_values = pixel_values.reshape(
453 + batch_size * num_images, channels, height, width
454 + )
455 +
456 + # Get vision features from SigLIP
457 + # pixel_values shape: (batch_size * num_images, 3, H, W)
458 + vision_features = self.vit_model(pixel_values)
459 +
460 + # Pass through connector
461 + vision_embeds = self.connector(vision_features)
462 +
463 + # Add position embeddings
464 + batch_size, num_patches, hidden_size = vision_embeds.shape
465 + patch_size = self.config.vit_config.patch_size
466 + image_size = self.config.vit_config.image_size
467 +
468 + # Calculate grid dimensions
469 + num_patches_per_side = image_size // patch_size
470 +
471 + # Create flattened position IDs (0 to num_patches-1)
472 + # For BAGEL, we use extrapolate mode by default
473 + h_coords = torch.arange(num_patches_per_side, device=vision_embeds.device)
474 + w_coords = torch.arange(num_patches_per_side, device=vision_embeds.device)
475 + position_ids = (
476 + h_coords[:, None] * self.config.vit_max_num_patch_per_side + w_coords
477 + ).flatten()
478 + position_ids = position_ids.unsqueeze(0).expand(batch_size, -1).flatten()
479 +
480 + # Add position embeddings
481 + pos_embeds = self.vit_pos_embed(position_ids)
482 + pos_embeds = pos_embeds.reshape(batch_size, num_patches, hidden_size)
483 + # Ensure pos_embeds are on the same device as vision_embeds
484 + pos_embeds = pos_embeds.to(vision_embeds.device)
485 + vision_embeds = vision_embeds + pos_embeds
486 +
487 + # Split by image
488 + return tuple(vision_embeds)
489 +
490 + def get_multimodal_embeddings(self, **kwargs: object) -> MultiModalEmbeddings:
491 + """Get multimodal embeddings from input."""
492 + image_input = self._parse_and_validate_image_input(**kwargs)
493 + if image_input is None:
494 + return []
495 +
496 + return self._process_image_input(image_input)
497 +
498 + def get_language_model(self) -> nn.Module:
499 + return self.language_model
500 +
501 + def forward(
502 + self,
503 + input_ids: torch.Tensor,
504 + positions: torch.Tensor,
505 + intermediate_tensors: IntermediateTensors | None = None,
506 + inputs_embeds: torch.Tensor | None = None,
507 + **kwargs: object,
508 + ) -> torch.Tensor | IntermediateTensors:
509 + """Run forward pass for BAGEL.
510 +
511 + Args:
512 + input_ids: Flattened (concatenated) input_ids corresponding to a batch.
513 + positions: Flattened (concatenated) position ids corresponding to a batch.
514 + intermediate_tensors: Intermediate tensors from prior forward pass.
515 + inputs_embeds: Optional tensor of input embeddings.
516 + """
517 + if intermediate_tensors is not None:
518 + inputs_embeds = None
519 +
520 + hidden_states = self.language_model.model(
521 + input_ids=input_ids,
522 + positions=positions,
523 + intermediate_tensors=intermediate_tensors,
524 + inputs_embeds=inputs_embeds,
525 + )
526 + return hidden_states
527 +
528 + def compute_logits(
529 + self,
530 + hidden_states: torch.Tensor,
531 + ) -> torch.Tensor | None:
532 + return self.language_model.compute_logits(hidden_states)
533 +
534 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
535 + """Load weights from checkpoint."""
536 + skip_prefixes = []
537 + # Skip vit_pos_embed.pos_embed as it's handled by PositionEmbedding module
538 + skip_prefixes.append("vit_pos_embed.pos_embed")
539 +
540 + # If visual understanding is disabled, skip vision-related weights
541 + if self.vit_model is None:
542 + skip_prefixes.extend(["vit_model.", "connector.", "vit_pos_embed"])
543 +
544 + # Skip generation-related weights since we only support text2text and image2text
545 + # Filter out all image generation components:
546 + # - 'moe_gen': MoE generation weights
547 + # - 'latent_pos_embed': Latent position embeddings for VAE
548 + # - 'llm2vae', 'vae2llm': LLM-VAE projections
549 + # - 'time_embedder': Timestep embeddings for diffusion
550 + # - VAE encoder/decoder: Use specific prefixes to avoid matching vision encoder
551 + generation_keywords = [
552 + "moe_gen",
553 + "latent_pos_embed",
554 + "llm2vae",
555 + "vae2llm",
556 + "time_embedder",
557 + ]
558 + vae_prefixes = [
559 + "decoder.",
560 + "encoder.",
561 + ] # VAE encoder/decoder, not vision encoder
562 + filtered_weights = []
563 + for name, tensor in weights:
564 + # Skip generation-related keywords
565 + if any(skip in name for skip in generation_keywords):
566 + continue
567 + if any(name.startswith(prefix) for prefix in vae_prefixes):
568 + continue
569 +
570 + if "patch_embedding.weight" in name and tensor.ndim == 2:
571 + out_channels = tensor.shape[0]
572 + in_features = tensor.shape[1]
573 + patch_size = self.config.vit_config.patch_size
574 + in_channels = self.config.vit_config.num_channels
575 + if in_features == in_channels * patch_size * patch_size:
576 + tensor = tensor.reshape(
577 + out_channels, patch_size, patch_size, in_channels
578 + )
579 + tensor = tensor.permute(0, 3, 1, 2).contiguous()
580 +
581 + filtered_weights.append((name, tensor))
582 +
583 + loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes)
584 + return loader.load_weights(filtered_weights, mapper=self.hf_to_vllm_mapper)
vllm/model_executor/models/qwen2.py
+32
@@ -122,6 +122,8 @@ class Qwen2Attention(nn.Module):
122 prefix: str = "",
123 attn_type: str = AttentionType.DECODER,
124 dual_chunk_attention_config: dict[str, Any] | None = None,
125 + qk_norm: bool = False,
126 + rms_norm_eps: float = 1e-6,
127 ) -> None:
128 super().__init__()
129 self.hidden_size = hidden_size
@@ -144,6 +146,7 @@ class Qwen2Attention(nn.Module):
146 self.kv_size = self.num_kv_heads * self.head_dim
147 self.scaling = self.head_dim**-0.5
148 self.dual_chunk_attention_config = dual_chunk_attention_config
149 + self.qk_norm = qk_norm
150
151 self.qkv_proj = QKVParallelLinear(
152 hidden_size,
@@ -162,6 +165,11 @@ class Qwen2Attention(nn.Module):
165 prefix=f"{prefix}.o_proj",
166 )
167
168 + # QK Normalization support (used in BAGEL and some other models)
169 + if self.qk_norm:
170 + self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)
171 + self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)
172 +
173 self.rotary_emb = get_rope(
174 self.head_dim,
175 max_position=max_position,
@@ -197,6 +205,23 @@ class Qwen2Attention(nn.Module):
205 ) -> torch.Tensor:
206 qkv, _ = self.qkv_proj(hidden_states)
207 q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
208 +
209 + # Apply QK normalization if enabled (before RoPE)
210 + if self.qk_norm:
211 + # Reshape to apply per-head normalization
212 + # q shape: (total_tokens, q_size) -> (total_tokens, num_heads, head_dim)
213 + total_tokens = q.shape[0]
214 + q = q.view(total_tokens, self.num_heads, self.head_dim)
215 + k = k.view(total_tokens, self.num_kv_heads, self.head_dim)
216 +
217 + # Apply normalization
218 + q = self.q_norm(q)
219 + k = self.k_norm(k)
220 +
221 + # Reshape back
222 + q = q.view(total_tokens, self.q_size)
223 + k = k.view(total_tokens, self.kv_size)
224 +
225 q, k = self.rotary_emb(positions, q, k)
226 attn_output = self.attn(q, k, v)
227 output, _ = self.o_proj(attn_output)
@@ -227,6 +252,9 @@ class Qwen2DecoderLayer(nn.Module):
252 else:
253 attn_type = AttentionType.ENCODER_ONLY
254
255 + # Check if QK normalization is enabled (used in BAGEL and some other models)
256 + qk_norm = getattr(config, "qk_norm", False)
257 +
258 self.self_attn = Qwen2Attention(
259 hidden_size=self.hidden_size,
260 num_heads=config.num_attention_heads,
@@ -238,6 +266,8 @@ class Qwen2DecoderLayer(nn.Module):
266 prefix=f"{prefix}.self_attn",
267 attn_type=attn_type,
268 dual_chunk_attention_config=dual_chunk_attention_config,
269 + qk_norm=qk_norm,
270 + rms_norm_eps=config.rms_norm_eps,
271 )
272 self.mlp = Qwen2MLP(
273 hidden_size=self.hidden_size,
@@ -480,6 +510,8 @@ class Qwen2Model(nn.Module):
510 continue
511 if is_pp_missing_parameter(name, self):
512 continue
513 + if name not in params_dict:
514 + continue
515 param = params_dict[name]
516 weight_loader = getattr(param, "weight_loader", default_weight_loader)
517 weight_loader(param, loaded_weight)
vllm/model_executor/models/registry.py
+1
@@ -272,6 +272,7 @@ _MULTIMODAL_MODELS = {
272 "aya_vision",
273 "AyaVisionForConditionalGeneration",
274 ),
275 + "BagelForConditionalGeneration": ("bagel", "BagelForConditionalGeneration"),
276 "BeeForConditionalGeneration": ("bee", "BeeForConditionalGeneration"),
277 "Blip2ForConditionalGeneration": ("blip2", "Blip2ForConditionalGeneration"),
278 "ChameleonForConditionalGeneration": (
vllm/transformers_utils/config.py
+1
@@ -66,6 +66,7 @@ class LazyConfigDict(dict):
66
67 _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict(
68 afmoe="AfmoeConfig",
69 + bagel="BagelConfig",
70 chatglm="ChatGLMConfig",
71 deepseek_vl_v2="DeepseekVLV2Config",
72 deepseek_v32="DeepseekV3Config",
vllm/transformers_utils/configs/__init__.py
+2
@@ -16,6 +16,7 @@ import importlib
16
17 _CLASS_TO_MODULE: dict[str, str] = {
18 "AfmoeConfig": "vllm.transformers_utils.configs.afmoe",
19 + "BagelConfig": "vllm.transformers_utils.configs.bagel",
20 "ChatGLMConfig": "vllm.transformers_utils.configs.chatglm",
21 "DeepseekVLV2Config": "vllm.transformers_utils.configs.deepseek_vl2",
22 "DotsOCRConfig": "vllm.transformers_utils.configs.dotsocr",
@@ -54,6 +55,7 @@ _CLASS_TO_MODULE: dict[str, str] = {
55
56 __all__ = [
57 "AfmoeConfig",
58 + "BagelConfig",
59 "ChatGLMConfig",
60 "DeepseekVLV2Config",
61 "DeepseekV3Config",
vllm/transformers_utils/configs/bagel.py new
+53
@@ -0,0 +1,53 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from transformers import PretrainedConfig, SiglipVisionConfig
4 +from transformers.models.qwen2 import Qwen2Config
5 +
6 +
7 +class BagelConfig(PretrainedConfig):
8 + """Configuration class for BAGEL model."""
9 +
10 + model_type = "bagel"
11 +
12 + def __init__(
13 + self,
14 + visual_gen: bool = True,
15 + visual_und: bool = True,
16 + llm_config: dict | Qwen2Config | None = None,
17 + vit_config: dict | SiglipVisionConfig | None = None,
18 + vae_config: dict | None = None,
19 + latent_patch_size: int = 2,
20 + max_latent_size: int = 32,
21 + vit_max_num_patch_per_side: int = 70,
22 + connector_act: str = "gelu_pytorch_tanh",
23 + interpolate_pos: bool = False,
24 + timestep_shift: float = 1.0,
25 + **kwargs,
26 + ):
27 + super().__init__(**kwargs)
28 + self.visual_gen = visual_gen
29 + self.visual_und = visual_und
30 +
31 + # Convert dict configs to proper config objects
32 + if isinstance(llm_config, dict):
33 + self.llm_config = Qwen2Config(**llm_config)
34 + else:
35 + self.llm_config = llm_config or Qwen2Config()
36 +
37 + if isinstance(vit_config, dict):
38 + self.vit_config = SiglipVisionConfig(**vit_config)
39 + else:
40 + self.vit_config = vit_config or SiglipVisionConfig()
41 +
42 + self.vae_config = vae_config or {"z_channels": 16, "downsample": 8}
43 + self.latent_patch_size = latent_patch_size
44 + self.max_latent_size = max_latent_size
45 + self.vit_max_num_patch_per_side = vit_max_num_patch_per_side
46 + self.connector_act = connector_act
47 + self.interpolate_pos = interpolate_pos
48 + self.timestep_shift = timestep_shift
49 +
50 + @property
51 + def hidden_size(self) -> int:
52 + """Return the hidden size of the language model."""
53 + return self.llm_config.hidden_size
vllm/transformers_utils/processors/__init__.py
+2
@@ -8,6 +8,7 @@ reasons:
8 - There is a need to override the existing processor to support vLLM.
9 """
10
11 +from vllm.transformers_utils.processors.bagel import BagelProcessor
12 from vllm.transformers_utils.processors.deepseek_vl2 import DeepseekVLV2Processor
13 from vllm.transformers_utils.processors.hunyuan_vl import HunYuanVLProcessor
14 from vllm.transformers_utils.processors.hunyuan_vl_image import HunYuanVLImageProcessor
@@ -15,6 +16,7 @@ from vllm.transformers_utils.processors.ovis import OvisProcessor
16 from vllm.transformers_utils.processors.ovis2_5 import Ovis2_5Processor
17
18 __all__ = [
19 + "BagelProcessor",
20 "DeepseekVLV2Processor",
21 "HunYuanVLProcessor",
22 "HunYuanVLImageProcessor",
vllm/transformers_utils/processors/bagel.py new
+73
@@ -0,0 +1,73 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +# Copyright 2025 Bytedance Ltd. and/or its affiliates.
4 +"""BAGEL processor for image and text inputs."""
5 +
6 +from transformers import AutoProcessor
7 +from transformers.image_utils import ImageInput
8 +from transformers.processing_utils import ProcessorMixin
9 +from transformers.tokenization_utils_base import PreTokenizedInput, TextInput
10 +
11 +
12 +class BagelProcessor(ProcessorMixin):
13 + """
14 + Constructs a BAGEL processor which wraps a
15 + SigLIP image processor and a Qwen2 tokenizer.
16 + """
17 +
18 + attributes = ["image_processor", "tokenizer"]
19 + image_processor_class = "SiglipImageProcessor"
20 + tokenizer_class = "AutoTokenizer"
21 +
22 + def __call__(
23 + self,
24 + text: TextInput
25 + | PreTokenizedInput
26 + | list[TextInput]
27 + | list[PreTokenizedInput] = None,
28 + images: ImageInput = None,
29 + **kwargs,
30 + ):
31 + """
32 + Main method to prepare for the model one or several sequences(s) and image(s).
33 + """
34 + if images is not None:
35 + # Process images with the image processor
36 + # Ensure return_tensors is set to "pt" for PyTorch tensors
37 + image_kwargs = {**kwargs}
38 + if "return_tensors" not in image_kwargs:
39 + image_kwargs["return_tensors"] = "pt"
40 + pixel_values = self.image_processor(images, **image_kwargs)
41 + else:
42 + pixel_values = None
43 +
44 + text_inputs = self.tokenizer(text, **kwargs) if text is not None else None
45 +
46 + if pixel_values is not None and text_inputs is not None:
47 + text_inputs["pixel_values"] = pixel_values["pixel_values"]
48 + return text_inputs
49 + elif pixel_values is not None:
50 + return pixel_values
51 + else:
52 + return text_inputs
53 +
54 + def batch_decode(self, *args, **kwargs):
55 + """
56 + This method forwards all its arguments to Qwen2TokenizerFast's batch_decode.
57 + """
58 + return self.tokenizer.batch_decode(*args, **kwargs)
59 +
60 + def decode(self, *args, **kwargs):
61 + """
62 + This method forwards all its arguments to Qwen2TokenizerFast's decode.
63 + """
64 + return self.tokenizer.decode(*args, **kwargs)
65 +
66 + @property
67 + def model_input_names(self):
68 + tokenizer_input_names = self.tokenizer.model_input_names
69 + image_processor_input_names = self.image_processor.model_input_names
70 + return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names))
71 +
72 +
73 +AutoProcessor.register("BagelProcessor", BagelProcessor)