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 9, 13:33 UTC
[LoRA][Spec Decode] Support LoRA for Nemotron-H MTP models (#32265)
Signed-off-by: Daniel Serebrenik <daserebrenik@nvidia.com> Co-authored-by: Jee Jee Li <pandaleefree@gmail.com>
danisereb committed
Jan 27, 2026 at 17:53 UTC
f3a5ee705fa9d655c34c476dfed45e55293ab4f9
4 files changed
+31
vllm/lora/lora_model.py
+22
@@ -62,6 +62,14 @@ class LoRAModel:
62
def check_lora_name(self, lora_name: str) -> bool:
63
return lora_name in self.loras
64
65
+ @staticmethod
66
+ def _should_skip_module(module_name: str, skip_prefixes: list[str]) -> bool:
67
+ """Check if a module should be skipped based on skip prefixes"""
68
+ for prefix in skip_prefixes:
69
+ if f".{prefix}" in module_name or module_name.startswith(prefix):
70
+ return True
71
+ return False
72
+
73
@classmethod
74
def from_lora_tensors(
75
cls,
@@ -72,6 +80,7 @@ class LoRAModel:
80
dtype: torch.dtype | None = None,
81
model_vocab_size: int | None = None,
82
weights_mapper: WeightsMapper | None = None,
83
+ skip_prefixes: list[str] | None = None,
84
) -> "LoRAModel":
85
"""Create a LoRAModel from a dictionary of tensors."""
86
pin_memory = str(device) == "cpu" and is_pin_memory_available()
@@ -79,6 +88,9 @@ class LoRAModel:
88
for tensor_name, tensor in tensors.items():
89
if is_base_embeddding_weights(tensor_name):
90
continue
91
+ # Skip modules based on model-defined prefixes (e.g., MTP layers)
92
+ if skip_prefixes and cls._should_skip_module(tensor_name, skip_prefixes):
93
+ continue
94
module_name, is_lora_a = parse_fine_tuned_lora_name(
95
tensor_name, weights_mapper
96
)
@@ -121,6 +133,7 @@ class LoRAModel:
133
model_vocab_size: int | None = None,
134
weights_mapper: WeightsMapper | None = None,
135
tensorizer_config_dict: dict | None = None,
136
+ skip_prefixes: list[str] | None = None,
137
) -> "LoRAModel":
138
"""Create a LoRAModel from a local checkpoint.
139
@@ -133,6 +146,9 @@ class LoRAModel:
146
a global counter.
147
device: Device where the lora model is loaded.
148
dtype: dtype of the lora model weights.
149
+ skip_prefixes: List of module name prefixes to skip during loading.
150
+ Models can define this to skip modules not used in inference
151
+ (e.g., MTP layers). Format: ["mtp."]
152
153
Returns:
154
Loaded LoRA Model.
@@ -152,6 +168,11 @@ class LoRAModel:
168
# gate_up_proj and experts is the down_proj
169
if "base_layer" in lora_module:
170
continue
171
+ # Skip modules based on model-defined prefixes
172
+ if skip_prefixes and cls._should_skip_module(
173
+ lora_module, skip_prefixes
174
+ ):
175
+ continue
176
module_name, _ = parse_fine_tuned_lora_name(lora_module, weights_mapper)
177
# Case for expert lora weights
178
if ".experts" in module_name:
@@ -218,4 +239,5 @@ class LoRAModel:
239
dtype=dtype,
240
model_vocab_size=model_vocab_size,
241
weights_mapper=weights_mapper,
242
+ skip_prefixes=skip_prefixes,
243
)
vllm/lora/worker_manager.py
+4
@@ -114,6 +114,9 @@ class WorkerLoRAManager:
114
model = self._adapter_manager.model
115
hf_to_vllm_mapper = getattr(model, "hf_to_vllm_mapper", None)
116
117
+ # Get model-defined prefixes to skip during LoRA loading.
118
+ lora_skip_prefixes = getattr(model, "lora_skip_prefixes", None)
119
+
120
lora = self._lora_model_cls.from_local_checkpoint(
121
lora_path,
122
expected_lora_modules,
@@ -124,6 +127,7 @@ class WorkerLoRAManager:
127
model_vocab_size=self.vocab_size,
128
tensorizer_config_dict=lora_request.tensorizer_config_dict,
129
weights_mapper=hf_to_vllm_mapper,
130
+ skip_prefixes=lora_skip_prefixes,
131
)
132
133
except FileNotFoundError as e:
vllm/model_executor/models/interfaces.py
+2
@@ -521,6 +521,8 @@ class SupportsLoRA(Protocol):
521
# are empty by default.
522
embedding_modules: ClassVar[dict[str, str]] = {}
523
packed_modules_mapping: dict[str, list[str]] = {}
524
+ # Module prefixes to skip during LoRA loading (e.g., ["mtp."] for MTP layers)
525
+ lora_skip_prefixes: ClassVar[list[str]] = []
526
527
528
# We can't use runtime_checkable with ClassVar for issubclass checks
vllm/model_executor/models/nemotron_h.py
+3
@@ -771,6 +771,9 @@ class NemotronHForCausalLM(
771
"lm_head": "output_embeddings",
772
}
773
774
+ # Skip MTP (Multi-Token Prediction) layers during LoRA loading
775
+ lora_skip_prefixes = ["mtp."]
776
+
777
@classmethod
778
def get_mamba_state_dtype_from_config(
779
cls,