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 12, 13:30 UTC
[Misc] Allow LM only loading for Pixtral (#29451)
Signed-off-by: Roger Wang <hey@rogerw.io>
Roger Wang committed
Nov 26, 2025 at 05:00 UTC
452a7c9f7c949cd20c3c0c81cd4352b2a0045076
1 file changed
+51
-22
vllm/model_executor/models/pixtral.py
+51
-22
@@ -400,21 +400,30 @@ class PixtralForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP)
400
prefix=maybe_prefix(prefix, "language_model"),
401
)
402
403
- self.vision_encoder = VisionTransformer(self.vision_args)
404
-
405
- if self.vision_args.add_pre_mm_projector_layer_norm:
406
- self.pre_mm_projector_norm = RMSNorm(self.vision_args.hidden_size, eps=1e-5)
407
-
408
- if self.vision_args.mm_projector_id == PATCH_MERGE:
409
- self.patch_merger = PatchMerger(
410
- vision_encoder_dim=self.vision_args.hidden_size,
411
- spatial_merge_size=self.vision_args.spatial_merge_size,
412
- use_mlp_bias=False,
403
+ if multimodal_config.get_limit_per_prompt("image"):
404
+ self.vision_encoder = VisionTransformer(self.vision_args)
405
+ self.pre_mm_projector_norm = (
406
+ RMSNorm(self.vision_args.hidden_size, eps=1e-5)
407
+ if self.vision_args.add_pre_mm_projector_layer_norm
408
+ else None
409
)
414
-
415
- self.vision_language_adapter = VisionLanguageAdapter(
416
- self.vision_args, dim=config.text_config.hidden_size
417
- )
410
+ self.patch_merger = (
411
+ PatchMerger(
412
+ vision_encoder_dim=self.vision_args.hidden_size,
413
+ spatial_merge_size=self.vision_args.spatial_merge_size,
414
+ use_mlp_bias=False,
415
+ )
416
+ if self.vision_args.mm_projector_id == PATCH_MERGE
417
+ else None
418
+ )
419
+ self.vision_language_adapter = VisionLanguageAdapter(
420
+ self.vision_args, dim=config.text_config.hidden_size
421
+ )
422
+ else:
423
+ self.vision_encoder = None
424
+ self.pre_mm_projector_norm = None
425
+ self.patch_merger = None
426
+ self.vision_language_adapter = None
427
428
self.make_empty_intermediate_tensors = (
429
self.language_model.make_empty_intermediate_tensors
@@ -436,13 +445,17 @@ class PixtralForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP)
445
self,
446
image_input: PixtralImagePixelInputs,
447
) -> tuple[torch.Tensor, ...]:
448
+ assert (
449
+ self.vision_encoder is not None and self.vision_language_adapter is not None
450
+ )
451
+
452
images = image_input["images"]
453
image_features = self.vision_encoder(images)
454
feature_sizes = [image_feature.shape[0] for image_feature in image_features]
455
image_features = torch.cat(image_features)
443
- if self.vision_args.add_pre_mm_projector_layer_norm:
456
+ if self.pre_mm_projector_norm is not None:
457
image_features = self.pre_mm_projector_norm(image_features)
445
- if self.vision_args.mm_projector_id == PATCH_MERGE:
458
+ if self.patch_merger is not None:
459
patch_size = self.vision_args.patch_size
460
spatial_merge_size_square = self.vision_args.spatial_merge_size**2
461
img_patch_dims = [
@@ -508,41 +521,57 @@ class PixtralForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP)
521
return weight[0].startswith("pre_mm_projector_norm")
522
523
# Get references to parameters for direct loading
511
- vision_encoder_dict = dict(self.vision_encoder.named_parameters())
524
+ vision_encoder_dict = (
525
+ dict(self.vision_encoder.named_parameters())
526
+ if self.vision_encoder is not None
527
+ else {}
528
+ )
529
patch_merger_dict = (
530
dict(self.patch_merger.named_parameters())
514
- if self.vision_args.mm_projector_id == PATCH_MERGE
515
- else dict()
531
+ if self.patch_merger is not None
532
+ else {}
533
)
534
pre_mm_projector_norm_dict = (
535
dict(self.pre_mm_projector_norm.named_parameters())
519
- if self.vision_args.add_pre_mm_projector_layer_norm
520
- else dict()
536
+ if self.pre_mm_projector_norm is not None
537
+ else {}
538
+ )
539
+ vision_lang_adapter_dict = (
540
+ dict(self.vision_language_adapter.named_parameters())
541
+ if self.vision_language_adapter is not None
542
+ else {}
543
)
522
- vision_lang_adapter_dict = dict(self.vision_language_adapter.named_parameters())
544
545
def llm_weights_generator():
546
# Single pass over weights
547
for name, w in weights:
548
if is_vision_encoder_weights((name, w)):
549
+ if self.vision_encoder is None:
550
+ continue
551
# Load vision encoder weights directly
552
trimmed_name = ".".join(name.split(".")[1:])
553
param = vision_encoder_dict[trimmed_name]
554
with torch.no_grad():
555
default_weight_loader(param, w)
556
elif is_patch_merger((name, w)):
557
+ if self.patch_merger is None:
558
+ continue
559
# Load vision patch merger weights directly
560
trimmed_name = ".".join(name.split(".")[1:])
561
param = patch_merger_dict[trimmed_name]
562
with torch.no_grad():
563
default_weight_loader(param, w)
564
elif is_pre_mm_projector_norm((name, w)):
565
+ if self.pre_mm_projector_norm is None:
566
+ continue
567
# Load vision pre_mm_projector_norm weights directly
568
trimmed_name = ".".join(name.split(".")[1:])
569
param = pre_mm_projector_norm_dict[trimmed_name]
570
with torch.no_grad():
571
default_weight_loader(param, w)
572
elif is_vision_lang_adapter_weights((name, w)):
573
+ if self.vision_language_adapter is None:
574
+ continue
575
# Load vision-language adapter weights directly
576
trimmed_name = ".".join(name.split(".")[1:])
577
param = vision_lang_adapter_dict[trimmed_name]