Fix quantized Falcon-H1 model loading issues (vllm-project#32728)

Signed-off-by: Shengliang Xu <[email protected]>
Co-authored-by: Cyrus Leung <[email protected]>

Author: shengliangxu

vllm/model_executor/models/falcon_h1.py

@@ -35,7 +35,10 @@
     ParallelLMHead,
     VocabParallelEmbedding,
 )
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+from vllm.model_executor.model_loader.weight_utils import (
+    default_weight_loader,
+    maybe_remap_kv_scale_name,
+)
 from vllm.sequence import IntermediateTensors
 from vllm.transformers_utils.config import set_default_rope_theta
 
@@ -278,6 +281,7 @@ def __init__(
             self.scaling,
             num_kv_heads=self.num_kv_heads,
             cache_config=cache_config,
+            quant_config=quant_config,
             prefix=f"{prefix}.attn",
         )
         self.key_multiplier = config.key_multiplier
@@ -360,7 +364,9 @@ def __init__(
         self.attention_in_multiplier = config.attention_in_multiplier
         self.attn_out_multiplier = config.attention_out_multiplier
 
-        self.feed_forward = FalconH1MLP(config, prefix=f"{prefix}.feed_forward")
+        self.feed_forward = FalconH1MLP(
+            config, quant_config=quant_config, prefix=f"{prefix}.feed_forward"
+        )
 
         self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
         self.pre_ff_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
@@ -647,6 +653,12 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
             if "mamba" in name:
                 name = name.replace("mamba", "mamba.mamba")
 
+            if "scale" in name:
+                # Remapping the name of kv-scale.
+                name = maybe_remap_kv_scale_name(name, params_dict)
+                if name is None:
+                    continue
+
             for param_name, weight_name, shard_id in stacked_params_mapping:
                 if weight_name not in name:
                     continue