@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / f23fb5a7c1
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 9, 21:48 UTC

[Bugfix] Support HF sharded weights for Mistral3/Pixtral models (#32673)

Signed-off-by: ricky-chaoju <ricky.chen@infinirc.com> Signed-off-by: vllm-dev <ricky.chen@infinirc.com>

RickyChen / 陳昭儒 committed Jan 21, 2026 at 15:27 UTC f23fb5a7c1b61350c5c40ca1115d3bf8cf2b8cc9
1 file changed +14 -8
vllm/model_executor/models/pixtral.py
+14 -8
@@ -502,10 +502,12 @@ class PixtralForConditionalGeneration(
502
503 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
504 def is_vision_encoder_weights(weight: tuple[str, torch.Tensor]):
505 - return weight[0].startswith("vision_encoder")
505 + return weight[0].startswith(("vision_encoder", "vision_tower"))
506
507 def is_vision_lang_adapter_weights(weight: tuple[str, torch.Tensor]):
508 - return weight[0].startswith("vision_language_adapter")
508 + return weight[0].startswith(
509 + ("vision_language_adapter", "multi_modal_projector")
510 + )
511
512 def is_patch_merger(weight: tuple[str, torch.Tensor]):
513 return weight[0].startswith("patch_merger")
@@ -543,9 +545,10 @@ class PixtralForConditionalGeneration(
545 continue
546 # Load vision encoder weights directly
547 trimmed_name = ".".join(name.split(".")[1:])
546 - param = vision_encoder_dict[trimmed_name]
547 - with torch.no_grad():
548 - default_weight_loader(param, w)
548 + param = vision_encoder_dict.get(trimmed_name)
549 + if param is not None:
550 + with torch.no_grad():
551 + default_weight_loader(param, w)
552 elif is_patch_merger((name, w)):
553 if self.patch_merger is None:
554 continue
@@ -567,12 +570,15 @@ class PixtralForConditionalGeneration(
570 continue
571 # Load vision-language adapter weights directly
572 trimmed_name = ".".join(name.split(".")[1:])
570 - param = vision_lang_adapter_dict[trimmed_name]
571 - with torch.no_grad():
572 - default_weight_loader(param, w)
573 + param = vision_lang_adapter_dict.get(trimmed_name)
574 + if param is not None:
575 + with torch.no_grad():
576 + default_weight_loader(param, w)
577 else:
578 # LLM weights: yield them to be loaded
579 # by language_model.load_weights
580 + # Strip "language_model." prefix if present (HF sharded format)
581 + name = name.removeprefix("language_model.")
582 yield (name, w)
583
584 # Now we call the language model load with the generator