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 5, 12:15 UTC
[Model] Use AutoWeightsLoader for FalconH1 (#39092)
Signed-off-by: Rishapveer Singh <215205492+rishaps@users.noreply.github.com>
Rishapveer Singh committed
Apr 7, 2026 at 10:25 UTC
da4c0e4db93c3d3ad18003cca56c645c91010532
1 file changed
+63
-59
vllm/model_executor/models/falcon_h1.py
+63
-59
@@ -50,6 +50,7 @@ from .interfaces import (
50
SupportsPP,
51
)
52
from .utils import (
53
+ AutoWeightsLoader,
54
PPMissingLayer,
55
is_pp_missing_parameter,
56
make_empty_intermediate_tensors_factory,
@@ -495,6 +496,63 @@ class FalconH1Model(nn.Module):
496
hidden_states = self.final_layernorm(hidden_states)
497
return hidden_states
498
499
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
500
+ stacked_params_mapping = [
501
+ # (param_name, shard_name, shard_id)
502
+ ("qkv_proj", "q_proj", "q"),
503
+ ("qkv_proj", "k_proj", "k"),
504
+ ("qkv_proj", "v_proj", "v"),
505
+ ("gate_up_proj", "gate_proj", 0),
506
+ ("gate_up_proj", "up_proj", 1),
507
+ ]
508
+
509
+ params_dict = dict(self.named_parameters())
510
+ loaded_params: set[str] = set()
511
+ for name, loaded_weight in weights:
512
+ if "rotary_emb.inv_freq" in name:
513
+ continue
514
+
515
+ if "A_log" in name:
516
+ name = name.replace("A_log", "A")
517
+
518
+ if "mamba" in name:
519
+ name = name.replace("mamba", "mamba.mamba")
520
+
521
+ if "scale" in name:
522
+ # Remapping the name of kv-scale.
523
+ name = maybe_remap_kv_scale_name(name, params_dict)
524
+ if name is None:
525
+ continue
526
+
527
+ for param_name, weight_name, shard_id in stacked_params_mapping:
528
+ if weight_name not in name:
529
+ continue
530
+
531
+ name = name.replace(weight_name, param_name)
532
+ # Skip loading extra bias for GPTQ models.
533
+ if name.endswith(".bias") and name not in params_dict:
534
+ continue
535
+ # Skip layers on other devices.
536
+ if is_pp_missing_parameter(name, self):
537
+ continue
538
+ param = params_dict[name]
539
+ weight_loader = param.weight_loader
540
+ weight_loader(param, loaded_weight, shard_id)
541
+ break
542
+ else:
543
+ # Skip loading extra bias for GPTQ models.
544
+ if name.endswith(".bias") and name not in params_dict:
545
+ continue
546
+ if is_pp_missing_parameter(name, self):
547
+ continue
548
+
549
+ param = params_dict[name]
550
+ weight_loader = getattr(param, "weight_loader", default_weight_loader)
551
+ weight_loader(param, loaded_weight)
552
+ loaded_params.add(name)
553
+
554
+ return loaded_params
555
+
556
557
class FalconH1ForCausalLM(
558
nn.Module,
@@ -632,62 +690,8 @@ class FalconH1ForCausalLM(
690
return logits
691
692
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
635
- stacked_params_mapping = [
636
- # (param_name, shard_name, shard_id)
637
- ("qkv_proj", "q_proj", "q"),
638
- ("qkv_proj", "k_proj", "k"),
639
- ("qkv_proj", "v_proj", "v"),
640
- ("gate_up_proj", "gate_proj", 0),
641
- ("gate_up_proj", "up_proj", 1),
642
- ]
643
-
644
- params_dict = dict(self.named_parameters())
645
- loaded_params: set[str] = set()
646
- for name, loaded_weight in weights:
647
- if "rotary_emb.inv_freq" in name:
648
- continue
649
-
650
- if "A_log" in name:
651
- name = name.replace("A_log", "A")
652
-
653
- if "mamba" in name:
654
- name = name.replace("mamba", "mamba.mamba")
655
-
656
- if "scale" in name:
657
- # Remapping the name of kv-scale.
658
- name = maybe_remap_kv_scale_name(name, params_dict)
659
- if name is None:
660
- continue
661
-
662
- for param_name, weight_name, shard_id in stacked_params_mapping:
663
- if weight_name not in name:
664
- continue
665
-
666
- name = name.replace(weight_name, param_name)
667
- # Skip loading extra bias for GPTQ models.
668
- if name.endswith(".bias") and name not in params_dict:
669
- continue
670
- # Skip layers on other devices.
671
- if is_pp_missing_parameter(name, self):
672
- continue
673
- param = params_dict[name]
674
- weight_loader = param.weight_loader
675
- weight_loader(param, loaded_weight, shard_id)
676
- break
677
- else:
678
- # Skip loading extra bias for GPTQ models.
679
- if name.endswith(".bias") and name not in params_dict:
680
- continue
681
- if is_pp_missing_parameter(name, self):
682
- continue
683
- if self.tie_word_embeddings and "lm_head" in name:
684
- continue
685
-
686
- param = params_dict[name]
687
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
688
- weight_loader(param, loaded_weight)
689
- loaded_params.add(name)
690
-
691
- if self.tie_word_embeddings:
692
- loaded_params.add("lm_head.weight")
693
- return loaded_params
693
+ loader = AutoWeightsLoader(
694
+ self,
695
+ skip_prefixes=(["lm_head."] if self.tie_word_embeddings else None),
696
+ )
697
+ return loader.load_weights(weights)