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 10, 20:15 UTC
Add get_expert_mapping to NemotronHModel (for LoRA support) (#31539)
Signed-off-by: Daniel Serebrenik <daserebrenik@nvidia.com>
danisereb committed
Dec 31, 2025 at 07:09 UTC
108a2728f74da693c3007094063a6ee357243a5e
1 file changed
+14
-10
vllm/model_executor/models/nemotron_h.py
+14
-10
@@ -632,14 +632,7 @@ class NemotronHModel(nn.Module):
632
hidden_states, _ = self.norm_f(hidden_states, residual)
633
return hidden_states
634
635
- def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
636
- stacked_params_mapping = [
637
- # (param_name, shard_name, shard_id)
638
- ("qkv_proj", "q_proj", "q"),
639
- ("qkv_proj", "k_proj", "k"),
640
- ("qkv_proj", "v_proj", "v"),
641
- ]
642
-
635
+ def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
636
if self.has_moe:
637
# (param_name, weight_name, expert_id, shard_id)
638
expert_params_mapping = FusedMoE.make_expert_params_mapping(
@@ -653,8 +646,19 @@ class NemotronHModel(nn.Module):
646
num_experts=self.config.n_routed_experts,
647
num_redundant_experts=getattr(self, "num_redundant_experts", 0),
648
)
656
- else:
657
- expert_params_mapping = []
649
+ return expert_params_mapping
650
+
651
+ return []
652
+
653
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
654
+ stacked_params_mapping = [
655
+ # (param_name, shard_name, shard_id)
656
+ ("qkv_proj", "q_proj", "q"),
657
+ ("qkv_proj", "k_proj", "k"),
658
+ ("qkv_proj", "v_proj", "v"),
659
+ ]
660
+
661
+ expert_params_mapping = self.get_expert_mapping()
662
663
params_dict = dict(self.named_parameters())
664
loaded_params: set[str] = set()