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 4, 07:30 UTC
[Model] Use AutoWeightsLoader for Plamo2 (#41699)
Signed-off-by: bittoby <218712309+bittoby@users.noreply.github.com>
BitToby committed
May 4, 2026 at 22:56 UTC
27cc676be3c08d41c6e4cf0df94aad94c8f9c616
1 file changed
+83
-85
vllm/model_executor/models/plamo2.py
+83
-85
@@ -797,6 +797,83 @@ class Plamo2Model(torch.nn.Module):
797
hidden_states, _ = self.norm(hidden_states, residual)
798
return hidden_states
799
800
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
801
+ params_dict = dict(self.named_parameters())
802
+ loaded_params: set[str] = set()
803
+ for name, loaded_weight in weights:
804
+ # Update the weight names to be compatible with the vllm version
805
+ # of the model.
806
+ # Do not change the order of the replacements.
807
+ replacements = {
808
+ # Rename incompatible weight names.
809
+ ".A_log": ".A",
810
+ ".B_norm_weight": ".B_norm.weight",
811
+ ".C_norm_weight": ".C_norm.weight",
812
+ ".dt_norm_weight": ".dt_norm.weight",
813
+ ".q_weight": ".q_norm.weight",
814
+ ".k_weight": ".k_norm.weight",
815
+ }
816
+ # Apply replacements based on the defined mappings
817
+ for old, new in replacements.items():
818
+ if old in name:
819
+ name = name.replace(old, new)
820
+
821
+ # Reshape the in_proj weights to match the shape expected
822
+ # by MergedColumnParallelLinear.
823
+ # This works both for unquantized weights and
824
+ # for quantized weights.
825
+ # In the quantized case, the weights are already transposed.
826
+ # Also, in addition to the quantized weights,
827
+ # the zero points and scales have to be reshaped as well.
828
+ # Packing should not be affected by this.
829
+ if (
830
+ ".mixer.in_proj.weight" in name
831
+ or "mixer.in_proj.qweight" in name
832
+ or "mixer.in_proj.scales" in name
833
+ or "mixer.in_proj.qzeros" in name
834
+ ):
835
+ if "mixer.in_proj.weight" in name:
836
+ loaded_weight = loaded_weight.transpose(0, 1)
837
+ # for weight:
838
+ # loaded_weight.shape[0] == self.config.hidden_size
839
+ # for qweight:
840
+ # loaded_weight.shape[0] == self.config.hidden_size // param.pack_factor # noqa
841
+ # for scales and qzeros:
842
+ # loaded_weight.shape[0] == self.config.hidden_size // self.vllm_config.quant_config.group_size # noqa
843
+ loaded_weight = loaded_weight.reshape(
844
+ loaded_weight.shape[0], self.config.mamba_num_heads, -1
845
+ )
846
+ gate_weight, hidden_states_weight = loaded_weight.chunk(2, dim=-1)
847
+ gate_weight = gate_weight.reshape(loaded_weight.shape[0], -1)
848
+ hidden_states_weight = hidden_states_weight.reshape(
849
+ loaded_weight.shape[0], -1
850
+ )
851
+ loaded_weight = torch.cat([gate_weight, hidden_states_weight], dim=-1)
852
+ if "mixer.in_proj.weight" in name:
853
+ loaded_weight = loaded_weight.transpose(0, 1)
854
+
855
+ # Offset parameter with vllm's RMSNorm haven't been supported yet.
856
+ if ".pre_mixer_norm" in name:
857
+ loaded_weight += 1.0
858
+ elif ".post_mixer_norm" in name:
859
+ loaded_weight += 1.0 / 5
860
+ elif ".pre_mlp_norm" in name:
861
+ loaded_weight += 1.0
862
+ elif ".post_mlp_norm" in name:
863
+ loaded_weight += 1.0 / (5**1.5)
864
+ elif name == "norm.weight":
865
+ loaded_weight += 1.0
866
+
867
+ # Skip layers on other devices.
868
+ if is_pp_missing_parameter(name, self):
869
+ continue
870
+
871
+ param = params_dict[name]
872
+ weight_loader = getattr(param, "weight_loader", default_weight_loader)
873
+ weight_loader(param, loaded_weight)
874
+ loaded_params.add(name)
875
+ return loaded_params
876
+
877
878
class Plamo2ForCausalLM(
879
torch.nn.Module, HasInnerState, SupportsLoRA, SupportsPP, IsHybrid
@@ -906,88 +983,9 @@ class Plamo2ForCausalLM(
983
logits = self.logits_processor(self.lm_head, hidden_states)
984
return logits
985
909
- def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
910
- params_dict = dict(self.named_parameters())
911
- for name, loaded_weight in weights:
912
- # Both tie_word_embeddings=True and lm_head.weight in the safetensor
913
- # at the same time causes dict key access error.
914
- if name == "lm_head.weight" and self.config.tie_word_embeddings:
915
- assert "lm_head.weight" not in params_dict
916
- continue
917
- # Same workaround as AutoWeightsLoader for GPTQModel
918
- if any(
919
- substr in name
920
- for substr in AutoWeightsLoader.ROTARY_EMBEDS_UNUSED_WEIGHTS
921
- ):
922
- continue
923
-
924
- # Update the weight names to be compatible with the vllm version
925
- # of the model.
926
- # Do not change the order of the replacements.
927
- replacements = {
928
- # Rename incompatible weight names.
929
- ".A_log": ".A",
930
- ".B_norm_weight": ".B_norm.weight",
931
- ".C_norm_weight": ".C_norm.weight",
932
- ".dt_norm_weight": ".dt_norm.weight",
933
- ".q_weight": ".q_norm.weight",
934
- ".k_weight": ".k_norm.weight",
935
- }
936
- # Apply replacements based on the defined mappings
937
- for old, new in replacements.items():
938
- if old in name:
939
- name = name.replace(old, new)
940
-
941
- # Reshape the in_proj weights to match the shape expected
942
- # by MergedColumnParallelLinear.
943
- # This works both for unquantized weights and
944
- # for quantized weights.
945
- # In the quantized case, the weights are already transposed.
946
- # Also, in addition to the quantized weights,
947
- # the zero points and scales have to be reshaped as well.
948
- # Packing should not be affected by this.
949
- if (
950
- ".mixer.in_proj.weight" in name
951
- or "mixer.in_proj.qweight" in name
952
- or "mixer.in_proj.scales" in name
953
- or "mixer.in_proj.qzeros" in name
954
- ):
955
- if "mixer.in_proj.weight" in name:
956
- loaded_weight = loaded_weight.transpose(0, 1)
957
- # for weight:
958
- # loaded_weight.shape[0] == self.config.hidden_size
959
- # for qweight:
960
- # loaded_weight.shape[0] == self.config.hidden_size // param.pack_factor # noqa
961
- # for scales and qzeros:
962
- # loaded_weight.shape[0] == self.config.hidden_size // self.vllm_config.quant_config.group_size # noqa
963
- loaded_weight = loaded_weight.reshape(
964
- loaded_weight.shape[0], self.config.mamba_num_heads, -1
965
- )
966
- gate_weight, hidden_states_weight = loaded_weight.chunk(2, dim=-1)
967
- gate_weight = gate_weight.reshape(loaded_weight.shape[0], -1)
968
- hidden_states_weight = hidden_states_weight.reshape(
969
- loaded_weight.shape[0], -1
970
- )
971
- loaded_weight = torch.cat([gate_weight, hidden_states_weight], dim=-1)
972
- if "mixer.in_proj.weight" in name:
973
- loaded_weight = loaded_weight.transpose(0, 1)
974
-
975
- # Offset parameter with vllm's RMSNorm haven't been supported yet.
976
- if ".pre_mixer_norm" in name:
977
- loaded_weight += 1.0
978
- elif ".post_mixer_norm" in name:
979
- loaded_weight += 1.0 / 5
980
- elif ".pre_mlp_norm" in name:
981
- loaded_weight += 1.0
982
- elif ".post_mlp_norm" in name:
983
- loaded_weight += 1.0 / (5**1.5)
984
- elif "model.norm.weight" in name:
985
- loaded_weight += 1.0
986
-
987
- # Skip layers on other devices.
988
- if is_pp_missing_parameter(name, self):
989
- continue
990
-
991
- param = params_dict[name]
992
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
993
- weight_loader(param, loaded_weight)
986
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
987
+ loader = AutoWeightsLoader(
988
+ self,
989
+ skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
990
+ )
991
+ return loader.load_weights(weights)