ac30a8311e · tduncan/Nvidia-Optimization-Stack-VLLM

@@ -8,7 +8,6 @@ from typing import Annotated, Literal

8

9 import torch

10 import torch.nn as nn

11 -import torch.nn.functional as F

11 from mistral_common.protocol.instruct.chunk import ImageChunk, TextChunk

12 from mistral_common.protocol.instruct.messages import UserMessage

13 from mistral_common.protocol.instruct.request import ChatCompletionRequest

@@ -26,16 +25,18 @@ from vllm.config import VllmConfig

25 from vllm.config.multimodal import BaseDummyOptions

26 from vllm.distributed import divide, get_tensor_model_parallel_world_size

27 from vllm.inputs import MultiModalDataDict

29 -from vllm.model_executor.layers.activation import get_act_and_mul_fn

28 +from vllm.model_executor.layers.activation import SiluAndMul, get_act_and_mul_fn

29 from vllm.model_executor.layers.conv import Conv2dLayer

30 from vllm.model_executor.layers.layernorm import RMSNorm

31 from vllm.model_executor.layers.linear import (

32 MergedColumnParallelLinear,

33 QKVParallelLinear,

34 + ReplicatedLinear,

35 RowParallelLinear,

36 )

37 from vllm.model_executor.layers.quantization import QuantizationConfig

38 from vllm.model_executor.model_loader.weight_utils import default_weight_loader

39 +from vllm.model_executor.models.utils import WeightsMapper

40 from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalKwargsItems

41 from vllm.multimodal.inputs import (

42 MultiModalFieldConfig,

@@ -293,6 +294,23 @@ class PixtralMultiModalProcessor(BaseMultiModalProcessor[PixtralProcessingInfo])

294 class PixtralForConditionalGeneration(

295 nn.Module, SupportsLoRA, SupportsEagle3, SupportsMultiModal, SupportsPP

296 ):

297 + hf_to_vllm_mapper = WeightsMapper(

298 + orig_to_new_prefix={

299 + "model.language_model.": "language_model.model.",

300 + "model.vision_tower.": "vision_encoder.",

301 + "model.multi_modal_projector.": "vision_language_adapter.",

302 + },

303 + orig_to_new_substr={

304 + ".linear_1.": ".w_in.",

305 + ".linear_2.": ".w_out.",

306 + },

307 + )

308 +

309 + packed_modules_mapping = {

310 + "qkv_proj": ["q_proj", "k_proj", "v_proj"],

311 + "gate_up_proj": ["gate_proj", "up_proj"],

312 + }

313 +

314 @classmethod

315 def get_placeholder_str(cls, modality: str, i: int) -> str | None:

316 if modality.startswith("image"):

@@ -325,7 +343,10 @@ class PixtralForConditionalGeneration(

343 )

344

345 with self._mark_tower_model(vllm_config, "image"):

328 - self.vision_encoder = VisionTransformer(self.vision_args)

346 + self.vision_encoder = VisionTransformer(

347 + self.vision_args,

348 + prefix=maybe_prefix(prefix, "vision_encoder"),

349 + )

350 self.pre_mm_projector_norm = (

351 RMSNorm(self.vision_args.hidden_size, eps=1e-5)

352 if self.vision_args.add_pre_mm_projector_layer_norm

@@ -435,6 +456,15 @@ class PixtralForConditionalGeneration(

456 return self.language_model.get_eagle3_aux_hidden_state_layers()

457

458 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):

459 + _vision_encoder_stacked_params = [

460 + # (param_name, shard_name, shard_id)

461 + (".qkv_proj", ".q_proj", "q"),

462 + (".qkv_proj", ".k_proj", "k"),

463 + (".qkv_proj", ".v_proj", "v"),

464 + (".gate_up_proj", ".gate_proj", 0),

465 + (".gate_up_proj", ".up_proj", 1),

466 + ]

467 +

468 def is_vision_encoder_weights(weight: tuple[str, torch.Tensor]):

469 return weight[0].startswith(("vision_encoder", "vision_tower"))

470

@@ -449,7 +479,6 @@ class PixtralForConditionalGeneration(

479 def is_pre_mm_projector_norm(weight: tuple[str, torch.Tensor]):

480 return weight[0].startswith("pre_mm_projector_norm")

481

452 - # Get references to parameters for direct loading

482 vision_encoder_dict = (

483 dict(self.vision_encoder.named_parameters())

484 if self.vision_encoder is not None

@@ -472,29 +501,41 @@ class PixtralForConditionalGeneration(

501 )

502

503 def llm_weights_generator():

475 - # Single pass over weights

504 for name, w in weights:

505 if is_vision_encoder_weights((name, w)):

506 if _is_layer_none_or_staged(self.vision_encoder):

507 continue

480 - # Load vision encoder weights directly

508 trimmed_name = ".".join(name.split(".")[1:])

482 - param = vision_encoder_dict.get(trimmed_name)

483 - if param is not None:

484 - with torch.no_grad():

485 - default_weight_loader(param, w)

509 + for (

510 + param_name,

511 + weight_name,

512 + shard_id,

513 + ) in _vision_encoder_stacked_params:

514 + if weight_name in trimmed_name:

515 + trimmed_name = trimmed_name.replace(weight_name, param_name)

516 + param = vision_encoder_dict[trimmed_name]

517 + weight_loader = param.weight_loader

518 + weight_loader(param, w, shard_id)

519 + break

520 + else:

521 + param = vision_encoder_dict.get(trimmed_name)

522 + if param is not None:

523 + weight_loader = getattr(

524 + param, "weight_loader", default_weight_loader

525 + )

526 + weight_loader(param, w)

527 elif is_patch_merger((name, w)):

528 if _is_layer_none_or_staged(self.patch_merger):

529 continue

489 - # Load vision patch merger weights directly

530 trimmed_name = ".".join(name.split(".")[1:])

531 param = patch_merger_dict[trimmed_name]

492 - with torch.no_grad():

493 - default_weight_loader(param, w)

532 + weight_loader = getattr(

533 + param, "weight_loader", default_weight_loader

534 + )

535 + weight_loader(param, w)

536 elif is_pre_mm_projector_norm((name, w)):

537 if _is_layer_none_or_staged(self.pre_mm_projector_norm):

538 continue

497 - # Load vision pre_mm_projector_norm weights directly

539 trimmed_name = ".".join(name.split(".")[1:])

540 param = pre_mm_projector_norm_dict[trimmed_name]

541 with torch.no_grad():

@@ -502,26 +543,23 @@ class PixtralForConditionalGeneration(

543 elif is_vision_lang_adapter_weights((name, w)):

544 if _is_layer_none_or_staged(self.vision_language_adapter):

545 continue

505 - # Load vision-language adapter weights directly

546 trimmed_name = ".".join(name.split(".")[1:])

547 param = vision_lang_adapter_dict.get(trimmed_name)

548 if param is not None:

509 - with torch.no_grad():

510 - default_weight_loader(param, w)

549 + weight_loader = getattr(

550 + param, "weight_loader", default_weight_loader

551 + )

552 + weight_loader(param, w)

553 else:

512 - # LLM weights: yield them to be loaded

513 - # by language_model.load_weights

514 - # Strip "language_model." prefix if present (HF sharded format)

554 name = name.removeprefix("language_model.")

555 yield (name, w)

556

518 - # Now we call the language model load with the generator

557 self.language_model.load_weights(llm_weights_generator())

558

559 def get_mm_mapping(self) -> MultiModelKeys:

560 return MultiModelKeys.from_string_field(

523 - language_model="language_model",

524 - connector="vision_language_adapter",

561 + language_model="language_model.",

562 + connector="vision_language_adapter.",

563 tower_model="vision_encoder",

564 )

565

@@ -614,29 +652,78 @@ def apply_rotary_emb_vit(

652

653

654 class FeedForward(nn.Module):

617 - def __init__(self, args: VisionEncoderArgs):

655 + def __init__(

656 + self,

657 + hidden_size: int,

658 + intermediate_size: int,

659 + quant_config: QuantizationConfig | None = None,

660 + bias: bool = False,

661 + prefix: str = "",

662 + reduce_results: bool = True,

663 + disable_tp: bool = False,

664 + ) -> None:

665 super().__init__()

619 - assert args.intermediate_size is not None

620 - self.w1 = nn.Linear(args.hidden_size, args.intermediate_size, bias=False)

621 - self.w2 = nn.Linear(args.intermediate_size, args.hidden_size, bias=False)

622 - self.w3 = nn.Linear(args.hidden_size, args.intermediate_size, bias=False)

666

624 - def forward(self, x: torch.Tensor) -> torch.Tensor:

625 - return self.w2(F.silu(self.w1(x)) * self.w3(x))

667 + self.gate_up_proj = MergedColumnParallelLinear(

668 + input_size=hidden_size,

669 + output_sizes=[intermediate_size] * 2,

670 + bias=bias,

671 + quant_config=quant_config,

672 + disable_tp=disable_tp,

673 + prefix=f"{prefix}.w13",

674 + )

675 + self.down_proj = RowParallelLinear(

676 + input_size=intermediate_size,

677 + output_size=hidden_size,

678 + bias=bias,

679 + quant_config=quant_config,

680 + reduce_results=reduce_results,

681 + disable_tp=disable_tp,

682 + prefix=f"{prefix}.w2",

683 + )

684 +

685 + self.act_fn = SiluAndMul()

686 +

687 + def forward(self, x):

688 + x, _ = self.gate_up_proj(x)

689 + x = self.act_fn(x)

690 + x, _ = self.down_proj(x)

691 + return x

692

693

694 class Attention(nn.Module):

629 - def __init__(self, args: VisionEncoderArgs):

695 + def __init__(

696 + self,

697 + args: VisionEncoderArgs,

698 + quant_config: QuantizationConfig | None = None,

699 + prefix: str = "",

700 + disable_tp: bool = False,

701 + ):

702 super().__init__()

703 self.args = args

704 assert not args.hidden_size % args.num_attention_heads

633 - self.n_heads = args.num_attention_heads

705 self.head_dim = args.hidden_size // args.num_attention_heads

706

636 - self.wq = nn.Linear(args.hidden_size, args.hidden_size, bias=False)

637 - self.wk = nn.Linear(args.hidden_size, args.hidden_size, bias=False)

638 - self.wv = nn.Linear(args.hidden_size, args.hidden_size, bias=False)

639 - self.wo = nn.Linear(args.hidden_size, args.hidden_size, bias=False)

707 + self.qkv_proj = QKVParallelLinear(

708 + hidden_size=args.hidden_size,

709 + head_size=self.head_dim,

710 + total_num_heads=args.num_attention_heads,

711 + bias=False,

712 + quant_config=quant_config,

713 + prefix=f"{prefix}.wqkv",

714 + disable_tp=disable_tp,

715 + )

716 + self.o_proj = RowParallelLinear(

717 + input_size=args.hidden_size,

718 + output_size=args.hidden_size,

719 + bias=False,

720 + quant_config=quant_config,

721 + prefix=f"{prefix}.wo",

722 + disable_tp=disable_tp,

723 + )

724 +

725 + tp_size = 1 if disable_tp else get_tensor_model_parallel_world_size()

726 + self.n_heads = divide(args.num_attention_heads, tp_size)

727

728 def forward(

729 self,

@@ -646,7 +733,8 @@ class Attention(nn.Module):

733 ) -> torch.Tensor:

734 batch, patches, _ = x.shape

735

649 - q, k, v = self.wq(x), self.wk(x), self.wv(x)

736 + qkv, _ = self.qkv_proj(x)

737 + q, k, v = qkv.chunk(3, dim=-1)

738 q = q.reshape(batch, patches, self.n_heads, self.head_dim)

739 k = k.reshape(batch, patches, self.n_heads, self.head_dim)

740 v = v.reshape(batch, patches, self.n_heads, self.head_dim)

@@ -663,14 +751,32 @@ class Attention(nn.Module):

751 out = out.transpose(1, 2)

752

753 out = out.reshape(batch, patches, self.n_heads * self.head_dim)

666 - return self.wo(out)

754 + out, _ = self.o_proj(out)

755 + return out

756

757

758 class TransformerBlock(nn.Module):

670 - def __init__(self, args: VisionEncoderArgs):

759 + def __init__(

760 + self,

761 + args: VisionEncoderArgs,

762 + quant_config: QuantizationConfig | None = None,

763 + prefix: str = "",

764 + disable_tp: bool = False,

765 + ):

766 super().__init__()

672 - self.attention = Attention(args)

673 - self.feed_forward = FeedForward(args)

767 + self.attention = Attention(

768 + args,

769 + quant_config=quant_config,

770 + prefix=f"{prefix}.attention",

771 + disable_tp=disable_tp,

772 + )

773 + self.feed_forward = FeedForward(

774 + args.hidden_size,

775 + args.intermediate_size,

776 + quant_config=quant_config,

777 + prefix=f"{prefix}.feed_forward",

778 + disable_tp=disable_tp,

779 + )

780 self.attention_norm = RMSNorm(args.hidden_size, eps=1e-5)

781 self.ffn_norm = RMSNorm(args.hidden_size, eps=1e-5)

782

@@ -690,11 +796,24 @@ class TransformerBlock(nn.Module):

796

797

798 class Transformer(nn.Module):

693 - def __init__(self, args: VisionEncoderArgs):

799 + def __init__(

800 + self,

801 + args: VisionEncoderArgs,

802 + quant_config: QuantizationConfig | None = None,

803 + prefix: str = "",

804 + disable_tp: bool = False,

805 + ):

806 super().__init__()

807 self.layers = torch.nn.ModuleList()

696 - for _ in range(args.num_hidden_layers):

697 - self.layers.append(TransformerBlock(args))

808 + for idx in range(args.num_hidden_layers):

809 + self.layers.append(

810 + TransformerBlock(

811 + args,

812 + quant_config=quant_config,

813 + prefix=f"{prefix}.layers.{idx}",

814 + disable_tp=disable_tp,

815 + )

816 + )

817

818 def forward(

819 self,

@@ -727,9 +846,15 @@ def position_meshgrid(

846

847

848 class VisionTransformer(nn.Module):

730 - def __init__(self, args: VisionEncoderArgs):

849 + def __init__(

850 + self,

851 + args: VisionEncoderArgs,

852 + quant_config: QuantizationConfig | None = None,

853 + prefix: str = "",

854 + ):

855 super().__init__()

856 self.args = args

857 + disable_tp = is_vit_use_data_parallel()

858 self.patch_conv = Conv2dLayer(

859 in_channels=args.num_channels,

860 out_channels=args.hidden_size,

@@ -738,7 +863,12 @@ class VisionTransformer(nn.Module):

863 bias=False,

864 )

865 self.ln_pre = RMSNorm(args.hidden_size, eps=1e-5)

741 - self.transformer = Transformer(args)

866 + self.transformer = Transformer(

867 + args,

868 + quant_config=quant_config,

869 + prefix=f"{prefix}.transformer",

870 + disable_tp=disable_tp,

871 + )

872

873 head_dim = self.args.hidden_size // self.args.num_attention_heads

874 assert head_dim % 2 == 0, "ROPE requires even head_dim"

@@ -822,13 +952,16 @@ class VisionLanguageAdapter(nn.Module):

952 def __init__(self, args: VisionEncoderArgs, dim: int):

953 super().__init__()

954 assert isinstance(args, VisionEncoderArgs)

825 - self.w_in = nn.Linear(

955 + self.w_in = ReplicatedLinear(

956 args.hidden_size,

957 dim,

958 bias=args.adapter_bias,

959 + return_bias=False,

960 )

961 self.gelu = nn.GELU()

831 - self.w_out = nn.Linear(dim, dim, bias=args.adapter_bias)

962 + self.w_out = ReplicatedLinear(

963 + dim, dim, bias=args.adapter_bias, return_bias=False

964 + )

965

966 def forward(self, x: torch.Tensor) -> torch.Tensor:

967 return self.w_out(self.gelu(self.w_in(x)))

@@ -852,10 +985,8 @@ class PatchMerger(nn.Module):

985 self.spatial_merge_size = spatial_merge_size

986 self.mlp_input_dim = mlp_input_dim

987

855 - self.merging_layer = nn.Linear(

856 - mlp_input_dim,

857 - vision_encoder_dim,

858 - bias=use_mlp_bias,

988 + self.merging_layer = ReplicatedLinear(

989 + mlp_input_dim, vision_encoder_dim, bias=use_mlp_bias, return_bias=False

990 )

991

992 def forward(

[Bugfix][Model] Fix PixtralForConditionalGeneration LoRA (#36963)