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 8, 01:15 UTC
[Models] LFM2: Support LoRA (#34921)
Co-authored-by: Piotr Mazurek <piotr635@gmail.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
tianshu-Michael-yu committed
Feb 19, 2026 at 22:07 UTC
ea37530b474fa738a99a53a8975af4e389b968c7
2 files changed
+36
-16
vllm/model_executor/models/lfm2.py
+18
-8
@@ -39,6 +39,7 @@ from .interfaces import HasInnerState, IsHybrid, SupportsLoRA, SupportsPP, Suppo
39
from .utils import (
40
AutoWeightsLoader,
41
PPMissingLayer,
42
+ WeightsMapper,
43
extract_layer_index,
44
is_pp_missing_parameter,
45
make_empty_intermediate_tensors_factory,
@@ -66,12 +67,12 @@ class Lfm2MLP(nn.Module):
67
ff_dim = int(ffn_dim_multiplier * ff_dim)
68
ff_dim = multiple_of * ((ff_dim + multiple_of - 1) // multiple_of)
69
69
- self.w1 = MergedColumnParallelLinear(
70
+ self.w13 = MergedColumnParallelLinear(
71
input_size=dim,
72
output_sizes=[ff_dim] * 2,
73
bias=False,
74
quant_config=quant_config,
74
- prefix=f"{prefix}.w1",
75
+ prefix=f"{prefix}.w13",
76
)
77
self.w2 = RowParallelLinear(
78
input_size=ff_dim,
@@ -83,7 +84,7 @@ class Lfm2MLP(nn.Module):
84
self.act_fn = SiluAndMul()
85
86
def forward(self, x: torch.Tensor) -> torch.Tensor:
86
- gate_up, _ = self.w1(x)
87
+ gate_up, _ = self.w13(x)
88
x = self.act_fn(gate_up)
89
x, _ = self.w2(x)
90
return x
@@ -376,8 +377,8 @@ class Lfm2Model(nn.Module):
377
(".qkv_proj", ".q_proj", "q"),
378
(".qkv_proj", ".k_proj", "k"),
379
(".qkv_proj", ".v_proj", "v"),
379
- (".w1", ".w1", 0),
380
- (".w1", ".w3", 1),
380
+ (".w13", ".w1", 0),
381
+ (".w13", ".w3", 1),
382
]
383
params_dict = dict(self.named_parameters())
384
loaded_params: set[str] = set()
@@ -386,9 +387,11 @@ class Lfm2Model(nn.Module):
387
name = name.replace(".conv.", ".short_conv.", 1)
388
389
for param_name, weight_name, shard_id in stacked_params_mapping:
389
- if weight_name not in name:
390
+ # Use segment-boundary matching (trailing dot) to prevent
391
+ # e.g. ".w1" from matching inside ".w13" in pre-fused keys.
392
+ if weight_name + "." not in name:
393
continue
391
- name = name.replace(weight_name, param_name)
394
+ name = name.replace(weight_name + ".", param_name + ".")
395
396
if is_pp_missing_parameter(name, self):
397
continue
@@ -415,13 +418,20 @@ class Lfm2ForCausalLM(
418
"k_proj",
419
"v_proj",
420
],
418
- "w1": [
421
+ "w13": [
422
"w1",
423
"w3",
424
],
425
"in_proj": ["in_proj"],
426
}
427
428
+ # HF uses .conv. but vLLM uses .short_conv. to avoid LoRA regex collision
429
+ # with the inner .conv.conv child (ShortConv has a child self.conv, so
430
+ # naming the container .conv too makes _match_target_modules match both)
431
+ hf_to_vllm_mapper = WeightsMapper(
432
+ orig_to_new_substr={".conv.": ".short_conv."},
433
+ )
434
+
435
# LoRA specific attributes
436
embedding_modules = {
437
"embed_tokens": "input_embeddings",
vllm/model_executor/models/lfm2_moe.py
+18
-8
@@ -52,6 +52,7 @@ from .interfaces import (
52
from .utils import (
53
AutoWeightsLoader,
54
PPMissingLayer,
55
+ WeightsMapper,
56
extract_layer_index,
57
is_pp_missing_parameter,
58
make_empty_intermediate_tensors_factory,
@@ -69,12 +70,12 @@ class Lfm2MoeMlp(nn.Module):
70
prefix: str = "",
71
):
72
super().__init__()
72
- self.w1 = MergedColumnParallelLinear(
73
+ self.w13 = MergedColumnParallelLinear(
74
input_size=dim,
75
output_sizes=[ff_dim] * 2,
76
bias=False,
77
quant_config=quant_config,
77
- prefix=f"{prefix}.w1",
78
+ prefix=f"{prefix}.w13",
79
)
80
self.w2 = RowParallelLinear(
81
input_size=ff_dim,
@@ -86,7 +87,7 @@ class Lfm2MoeMlp(nn.Module):
87
self.act_fn = SiluAndMul()
88
89
def forward(self, x: torch.Tensor) -> torch.Tensor:
89
- gate_up, _ = self.w1(x)
90
+ gate_up, _ = self.w13(x)
91
x = self.act_fn(gate_up)
92
x, _ = self.w2(x)
93
return x
@@ -501,8 +502,8 @@ class Lfm2MoeModel(nn.Module):
502
(".qkv_proj", ".q_proj", "q"),
503
(".qkv_proj", ".k_proj", "k"),
504
(".qkv_proj", ".v_proj", "v"),
504
- (".w1", ".w1", 0),
505
- (".w1", ".w3", 1),
505
+ (".w13", ".w1", 0),
506
+ (".w13", ".w3", 1),
507
]
508
params_dict = dict(self.named_parameters())
509
loaded_params: set[str] = set()
@@ -516,12 +517,14 @@ class Lfm2MoeModel(nn.Module):
517
518
for param_name, weight_name, shard_id in stacked_params_mapping:
519
# Skip non-stacked layers and experts (experts handled below).
519
- if weight_name not in name:
520
+ # Use segment-boundary matching (trailing dot) to prevent
521
+ # e.g. ".w1" from matching inside ".w13" in pre-fused keys.
522
+ if weight_name + "." not in name:
523
continue
524
525
if ("feed_forward.experts." in name) and name not in params_dict:
526
continue
524
- name = name.replace(weight_name, param_name)
527
+ name = name.replace(weight_name + ".", param_name + ".")
528
# Skip loading extra bias for GPTQ models.
529
if (
530
name.endswith(".bias") or name.endswith("_bias")
@@ -596,13 +599,20 @@ class Lfm2MoeForCausalLM(
599
"k_proj",
600
"v_proj",
601
],
599
- "w1": [
602
+ "w13": [
603
"w1",
604
"w3",
605
],
606
"in_proj": ["in_proj"],
607
}
608
609
+ # HF uses .conv. but vLLM uses .short_conv. to avoid LoRA regex collision
610
+ # with the inner .conv.conv child (ShortConv has a child self.conv, so
611
+ # naming the container .conv too makes _match_target_modules match both)
612
+ hf_to_vllm_mapper = WeightsMapper(
613
+ orig_to_new_substr={".conv.": ".short_conv."},
614
+ )
615
+
616
# LoRA specific attributes
617
embedding_modules = {
618
"embed_tokens": "input_embeddings",