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 6, 14:30 UTC
[Misc][LoRA] Add --lora-target-modules to restrict LoRA to specific modules (#34984)
Signed-off-by: Bhoomit Vasani <bhoomit.2010@gmail.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Bhoomit committed
Mar 17, 2026 at 07:36 UTC
3717a4dd475e6a936df0c84b043743310368e766
9 files changed
+404
-10
docs/features/lora.md
+14
@@ -389,3 +389,17 @@ vllm serve model --enable-lora --max-lora-rank 64
389
# Bad: unnecessarily high, wastes memory
390
vllm serve model --enable-lora --max-lora-rank 256
391
```
392
+
393
+### Restricting LoRA to Specific Modules
394
+
395
+The `--lora-target-modules` parameter allows you to restrict which model modules have LoRA applied at deployment time. This is useful for performance tuning when you only need LoRA on specific layers:
396
+
397
+```bash
398
+# Apply LoRA only to output projection layers
399
+vllm serve model --enable-lora --lora-target-modules o_proj
400
+
401
+# Apply LoRA to multiple specific modules
402
+vllm serve model --enable-lora --lora-target-modules o_proj qkv_proj down_proj
403
+```
404
+
405
+When `--lora-target-modules` is not specified, LoRA will be applied to all supported modules in the model. This parameter accepts module suffixes (the last component of the module name), such as `o_proj`, `qkv_proj`, `gate_proj`, etc.
tests/entrypoints/openai/test_cli_args.py
+29
@@ -291,3 +291,32 @@ def test_served_model_name_parsing(tmp_path, vllm_parser, args, raises):
291
else:
292
with pytest.raises(raises):
293
vllm_parser.parse_args(args=args)
294
+
295
+
296
+### Tests for LoRA target modules parsing
297
+def test_lora_target_modules_single(serve_parser):
298
+ """Test parsing single lora-target-modules argument"""
299
+ args = serve_parser.parse_args(
300
+ args=["--enable-lora", "--lora-target-modules", "o_proj"]
301
+ )
302
+ assert args.lora_target_modules == ["o_proj"]
303
+
304
+
305
+def test_lora_target_modules_multiple(serve_parser):
306
+ """Test parsing multiple lora-target-modules arguments"""
307
+ args = serve_parser.parse_args(
308
+ args=[
309
+ "--enable-lora",
310
+ "--lora-target-modules",
311
+ "o_proj",
312
+ "qkv_proj",
313
+ "down_proj",
314
+ ]
315
+ )
316
+ assert args.lora_target_modules == ["o_proj", "qkv_proj", "down_proj"]
317
+
318
+
319
+def test_lora_target_modules_default_none(serve_parser):
320
+ """Test that lora-target-modules defaults to None"""
321
+ args = serve_parser.parse_args(args=[])
322
+ assert args.lora_target_modules is None
tests/lora/test_lora_manager.py
+189
@@ -711,3 +711,192 @@ def test_packed_loras(default_vllm_config, dist_init, dummy_model_gate_up, devic
711
torch.testing.assert_close(
712
packed_lora1.lora_b[1], model_lora_clone1.get_lora("up_proj").lora_b
713
)
714
+
715
+
716
+def _test_target_modules(
717
+ model,
718
+ target_modules: list[str] | None,
719
+ device: str,
720
+ expected_lora: list[tuple[str, type]],
721
+ expected_no_lora: list[tuple[str, type]],
722
+):
723
+ """Create a LoRAModelManager and assert which modules have LoRA applied."""
724
+ LoRAModelManager(
725
+ model,
726
+ 2,
727
+ 2,
728
+ 2,
729
+ LoRAConfig(
730
+ max_lora_rank=8,
731
+ max_cpu_loras=2,
732
+ max_loras=2,
733
+ lora_dtype=DEFAULT_DTYPE,
734
+ target_modules=target_modules,
735
+ ),
736
+ device=device,
737
+ )
738
+ for module_path, lora_cls in expected_lora:
739
+ assert isinstance(model.get_submodule(module_path), lora_cls)
740
+ for module_path, lora_cls in expected_no_lora:
741
+ assert not isinstance(model.get_submodule(module_path), lora_cls)
742
+
743
+
744
+@pytest.mark.parametrize("device", DEVICES)
745
+def test_target_modules_config(default_vllm_config, dist_init, dummy_model, device):
746
+ """Test that target_modules config restricts which modules get LoRA applied."""
747
+ _test_target_modules(
748
+ dummy_model,
749
+ ["dense1"],
750
+ device,
751
+ expected_lora=[
752
+ ("dense1", ColumnParallelLinearWithLoRA),
753
+ ("layer1.dense1", ColumnParallelLinearWithLoRA),
754
+ ],
755
+ expected_no_lora=[
756
+ ("dense2", RowParallelLinearWithLoRA),
757
+ ("layer1.dense2", RowParallelLinearWithLoRA),
758
+ ],
759
+ )
760
+
761
+
762
+@pytest.mark.parametrize("device", DEVICES)
763
+def test_target_modules_multiple(default_vllm_config, dist_init, dummy_model, device):
764
+ """Test that multiple target_modules work correctly."""
765
+ _test_target_modules(
766
+ dummy_model,
767
+ ["dense1", "dense2"],
768
+ device,
769
+ expected_lora=[
770
+ ("dense1", ColumnParallelLinearWithLoRA),
771
+ ("layer1.dense1", ColumnParallelLinearWithLoRA),
772
+ ("dense2", RowParallelLinearWithLoRA),
773
+ ("layer1.dense2", RowParallelLinearWithLoRA),
774
+ ],
775
+ expected_no_lora=[],
776
+ )
777
+
778
+
779
+@pytest.mark.parametrize("device", DEVICES)
780
+def test_target_modules_none_uses_all(
781
+ default_vllm_config, dist_init, dummy_model, device
782
+):
783
+ """Test that target_modules=None uses all supported modules."""
784
+ _test_target_modules(
785
+ dummy_model,
786
+ None,
787
+ device,
788
+ expected_lora=[
789
+ ("dense1", ColumnParallelLinearWithLoRA),
790
+ ("layer1.dense1", ColumnParallelLinearWithLoRA),
791
+ ("dense2", RowParallelLinearWithLoRA),
792
+ ("layer1.dense2", RowParallelLinearWithLoRA),
793
+ ],
794
+ expected_no_lora=[],
795
+ )
796
+
797
+
798
+@pytest.mark.parametrize("device", DEVICES)
799
+def test_load_adapter_warns_on_unsupported_modules(
800
+ default_vllm_config, dist_init, dummy_model_gate_up, device, tmp_path
801
+):
802
+ """Test that _load_adapter warns when a LoRA adapter contains modules
803
+ not in the model's supported LoRA target modules."""
804
+ from unittest.mock import patch
805
+
806
+ import vllm.lora.worker_manager as wm_module
807
+
808
+ lora_config = LoRAConfig(
809
+ max_lora_rank=8, max_cpu_loras=4, max_loras=4, lora_dtype=DEFAULT_DTYPE
810
+ )
811
+
812
+ dummy_lora_files = f"{tmp_path}/lora_adapter"
813
+ os.makedirs(dummy_lora_files, exist_ok=True)
814
+ create_peft_lora(
815
+ dummy_model_gate_up,
816
+ save_dir=dummy_lora_files,
817
+ target_modules=["layer1.dense1", "dense2"],
818
+ lora_dtype=DEFAULT_DTYPE,
819
+ )
820
+
821
+ model_config = ModelConfig(max_model_len=16)
822
+ vllm_config = VllmConfig(model_config=model_config, lora_config=lora_config)
823
+ vllm_config.scheduler_config.max_num_seqs = 4
824
+ vllm_config.scheduler_config.max_num_batched_tokens = 2
825
+
826
+ worker_manager = WorkerLoRAManager(vllm_config, device, EMBEDDING_MODULES)
827
+ worker_manager.vocab_size = dummy_model_gate_up.unpadded_vocab_size
828
+ worker_manager.create_lora_manager(dummy_model_gate_up)
829
+
830
+ # Patch from_local_checkpoint to inject an unsupported module
831
+ original_from_checkpoint = LoRAModel.from_local_checkpoint
832
+
833
+ def patched_from_checkpoint(*args, **kwargs):
834
+ lora = original_from_checkpoint(*args, **kwargs)
835
+ lora.loras["unsupported_module"] = LoRALayerWeights(
836
+ module_name="unsupported_module",
837
+ rank=8,
838
+ lora_alpha=16,
839
+ lora_a=torch.randn(8, 10),
840
+ lora_b=torch.randn(10, 8),
841
+ )
842
+ return lora
843
+
844
+ lora_request = LoRARequest("test", 1, dummy_lora_files)
845
+ with (
846
+ patch.object(LoRAModel, "from_local_checkpoint", patched_from_checkpoint),
847
+ patch.object(wm_module.logger, "warning_once") as mock_warning,
848
+ ):
849
+ worker_manager._load_adapter(lora_request)
850
+ warning_args = mock_warning.call_args_list
851
+ found = any("unsupported_module" in str(call) for call in warning_args)
852
+ assert found, (
853
+ f"Expected warning about 'unsupported_module', got: {warning_args}"
854
+ )
855
+
856
+
857
+@pytest.mark.parametrize("device", DEVICES)
858
+def test_load_adapter_warns_on_target_modules_restriction(
859
+ default_vllm_config, dist_init, dummy_model_gate_up, device, tmp_path
860
+):
861
+ """Test that _load_adapter warns when a LoRA adapter contains modules
862
+ excluded by the deployment-time target_modules restriction."""
863
+ from unittest.mock import patch
864
+
865
+ import vllm.lora.worker_manager as wm_module
866
+
867
+ # Restrict to only dense2 — adapter has dense1 which will be excluded
868
+ lora_config = LoRAConfig(
869
+ max_lora_rank=8,
870
+ max_cpu_loras=4,
871
+ max_loras=4,
872
+ lora_dtype=DEFAULT_DTYPE,
873
+ target_modules=["dense2"],
874
+ )
875
+
876
+ dummy_lora_files = f"{tmp_path}/lora_adapter"
877
+ os.makedirs(dummy_lora_files, exist_ok=True)
878
+ create_peft_lora(
879
+ dummy_model_gate_up,
880
+ save_dir=dummy_lora_files,
881
+ target_modules=["layer1.dense1", "dense2"],
882
+ lora_dtype=DEFAULT_DTYPE,
883
+ )
884
+
885
+ model_config = ModelConfig(max_model_len=16)
886
+ vllm_config = VllmConfig(model_config=model_config, lora_config=lora_config)
887
+ vllm_config.scheduler_config.max_num_seqs = 4
888
+ vllm_config.scheduler_config.max_num_batched_tokens = 2
889
+
890
+ worker_manager = WorkerLoRAManager(vllm_config, device, EMBEDDING_MODULES)
891
+ worker_manager.vocab_size = dummy_model_gate_up.unpadded_vocab_size
892
+ worker_manager.create_lora_manager(dummy_model_gate_up)
893
+
894
+ lora_request = LoRARequest("test", 1, dummy_lora_files)
895
+ with patch.object(wm_module.logger, "warning_once") as mock_warning:
896
+ worker_manager._load_adapter(lora_request)
897
+ warning_args = mock_warning.call_args_list
898
+ # dense1 is supported by the model but excluded by target_modules
899
+ found = any("target_modules" in str(call) for call in warning_args)
900
+ assert found, (
901
+ f"Expected warning about target_modules restriction, got: {warning_args}"
902
+ )
tests/lora/test_lora_utils.py
new
+60
@@ -0,0 +1,60 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+
5
+from vllm.lora.utils import is_in_target_modules, is_supported_lora_module
6
+
7
+
8
+class TestIsSupportedLoraModule:
9
+ """Tests for is_supported_lora_module (model-definition check)."""
10
+
11
+ def test_suffix_match(self):
12
+ assert is_supported_lora_module(
13
+ "model.layers.0.self_attn.o_proj", ["o_proj", "q_proj"]
14
+ )
15
+
16
+ def test_no_match(self):
17
+ assert not is_supported_lora_module(
18
+ "model.layers.0.self_attn.o_proj", ["q_proj", "k_proj"]
19
+ )
20
+
21
+ def test_exact_match(self):
22
+ assert is_supported_lora_module("o_proj", ["o_proj"])
23
+
24
+ def test_regex_suffix_matching(self):
25
+ """Regex anchors to end — partial suffix should not match."""
26
+ assert not is_supported_lora_module("model.layers.0.self_attn.o_proj", ["proj"])
27
+
28
+ def test_empty_supported_modules(self):
29
+ assert not is_supported_lora_module("model.layers.0.self_attn.o_proj", [])
30
+
31
+ def test_multiple_supported_modules(self):
32
+ supported = ["q_proj", "k_proj", "v_proj", "o_proj"]
33
+ assert is_supported_lora_module("model.layers.0.self_attn.v_proj", supported)
34
+ assert not is_supported_lora_module("model.layers.0.mlp.gate_proj", supported)
35
+
36
+
37
+class TestIsInTargetModules:
38
+ """Tests for is_in_target_modules (deployment-time filter)."""
39
+
40
+ def test_none_allows_all(self):
41
+ assert is_in_target_modules("model.layers.0.self_attn.o_proj", None)
42
+
43
+ def test_suffix_in_target(self):
44
+ assert is_in_target_modules(
45
+ "model.layers.0.self_attn.o_proj", ["o_proj", "q_proj"]
46
+ )
47
+
48
+ def test_suffix_not_in_target(self):
49
+ assert not is_in_target_modules(
50
+ "model.layers.0.self_attn.o_proj", ["q_proj", "k_proj"]
51
+ )
52
+
53
+ def test_empty_target_modules(self):
54
+ assert not is_in_target_modules("model.layers.0.self_attn.o_proj", [])
55
+
56
+ def test_exact_name_match(self):
57
+ assert is_in_target_modules("dense1", ["dense1", "dense2"])
58
+
59
+ def test_exact_name_no_match(self):
60
+ assert not is_in_target_modules("dense3", ["dense1", "dense2"])
vllm/config/lora.py
+8
@@ -43,6 +43,10 @@ class LoRAConfig:
43
`max_loras`."""
44
lora_dtype: torch.dtype | LoRADType = "auto"
45
"""Data type for LoRA. If auto, will default to base model dtype."""
46
+ target_modules: list[str] | None = None
47
+ """Restrict LoRA to specific module suffixes (e.g., ["o_proj", "qkv_proj"]).
48
+ If None, all supported LoRA modules are used. This allows deployment-time
49
+ control over which modules have LoRA applied, useful for performance tuning."""
50
default_mm_loras: dict[str, str] | None = None
51
"""Dictionary mapping specific modalities to LoRA model paths; this field
52
is only applicable to multimodal models and should be leveraged when a
@@ -84,6 +88,10 @@ class LoRAConfig:
88
factors.append(self.fully_sharded_loras)
89
factors.append(self.lora_dtype)
90
factors.append(self.enable_tower_connector_lora)
91
+ # target_modules affects which modules get LoRA applied
92
+ factors.append(
93
+ tuple(sorted(self.target_modules)) if self.target_modules else None
94
+ )
95
96
hash_str = safe_hash(str(factors).encode(), usedforsecurity=False).hexdigest()
97
return hash_str
vllm/engine/arg_utils.py
+5
@@ -506,6 +506,7 @@ class EngineArgs:
506
fully_sharded_loras: bool = LoRAConfig.fully_sharded_loras
507
max_cpu_loras: int | None = LoRAConfig.max_cpu_loras
508
lora_dtype: str | torch.dtype | None = LoRAConfig.lora_dtype
509
+ lora_target_modules: list[str] | None = LoRAConfig.target_modules
510
enable_tower_connector_lora: bool = LoRAConfig.enable_tower_connector_lora
511
specialize_active_lora: bool = LoRAConfig.specialize_active_lora
512
@@ -1107,6 +1108,9 @@ class EngineArgs:
1108
lora_group.add_argument(
1109
"--fully-sharded-loras", **lora_kwargs["fully_sharded_loras"]
1110
)
1111
+ lora_group.add_argument(
1112
+ "--lora-target-modules", **lora_kwargs["target_modules"]
1113
+ )
1114
lora_group.add_argument("--default-mm-loras", **lora_kwargs["default_mm_loras"])
1115
lora_group.add_argument(
1116
"--specialize-active-lora", **lora_kwargs["specialize_active_lora"]
@@ -1800,6 +1804,7 @@ class EngineArgs:
1804
default_mm_loras=self.default_mm_loras,
1805
fully_sharded_loras=self.fully_sharded_loras,
1806
lora_dtype=self.lora_dtype,
1807
+ target_modules=self.lora_target_modules,
1808
enable_tower_connector_lora=self.enable_tower_connector_lora,
1809
specialize_active_lora=self.specialize_active_lora,
1810
max_cpu_loras=self.max_cpu_loras
vllm/lora/model_manager.py
+19
-9
@@ -5,7 +5,6 @@ import math
5
from collections.abc import Callable
6
from typing import TypeVar
7
8
-import regex as re
8
import torch
9
from torch import nn
10
@@ -25,7 +24,9 @@ from vllm.lora.utils import (
24
from_layer,
25
from_layer_logits_processor,
26
get_supported_lora_modules,
27
+ is_in_target_modules,
28
is_moe_model,
29
+ is_supported_lora_module,
30
process_packed_modules_mapping,
31
replace_submodule,
32
)
@@ -541,14 +542,23 @@ class LoRAModelManager:
542
model.loras[module_name] = lora
543
return model
544
544
- def _match_target_modules(self, module_name: str):
545
- return any(
546
- re.match(
547
- r".*\.{target_module}$".format(target_module=target_module), module_name
548
- )
549
- or target_module == module_name
550
- for target_module in self.supported_lora_modules
551
- )
545
+ def _match_target_modules(self, module_name: str) -> bool:
546
+ """Check if a module should have LoRA applied.
547
+
548
+ This method first checks if the module is in vLLM's supported LoRA
549
+ modules, then applies deployment-time restrictions based on
550
+ LoRAConfig.target_modules.
551
+
552
+ Args:
553
+ module_name: Full dot-separated module name (e.g.,
554
+ "model.layers.0.self_attn.o_proj")
555
+
556
+ Returns:
557
+ True if LoRA should be applied to this module, False otherwise.
558
+ """
559
+ if not is_supported_lora_module(module_name, self.supported_lora_modules):
560
+ return False
561
+ return is_in_target_modules(module_name, self.lora_config.target_modules)
562
563
def _get_punica_wrapper(self, module_name: str) -> PunicaWrapperBase | None:
564
"""
vllm/lora/utils.py
+52
@@ -5,6 +5,7 @@ import os
5
from typing import TYPE_CHECKING
6
7
import huggingface_hub
8
+import regex as re
9
from huggingface_hub.utils import HfHubHTTPError, HFValidationError
10
from torch import nn
11
from transformers import PretrainedConfig
@@ -226,6 +227,57 @@ def get_supported_lora_modules(model: nn.Module) -> list[str]:
227
return list(supported_lora_modules)
228
229
230
+def is_supported_lora_module(
231
+ module_name: str,
232
+ supported_lora_modules: list[str],
233
+) -> bool:
234
+ """Check if a module is in the model's supported LoRA modules.
235
+
236
+ Uses regex suffix matching against the model-defined supported modules
237
+ list (e.g., matching "model.layers.0.self_attn.o_proj" against
238
+ "o_proj").
239
+
240
+ Args:
241
+ module_name: Full dot-separated module name.
242
+ supported_lora_modules: List of module suffixes supported by the
243
+ model.
244
+
245
+ Returns:
246
+ True if the module is supported, False otherwise.
247
+ """
248
+ return any(
249
+ re.match(
250
+ r".*\.{target_module}$".format(target_module=target_module),
251
+ module_name,
252
+ )
253
+ or target_module == module_name
254
+ for target_module in supported_lora_modules
255
+ )
256
+
257
+
258
+def is_in_target_modules(
259
+ module_name: str,
260
+ target_modules: list[str] | None,
261
+) -> bool:
262
+ """Check if a module passes the deployment-time target_modules filter.
263
+
264
+ When target_modules is None (no restriction), all modules pass.
265
+ Otherwise, the module's suffix must be in the target_modules list.
266
+
267
+ Args:
268
+ module_name: Full dot-separated module name.
269
+ target_modules: Optional deployment-time restriction list from
270
+ LoRAConfig.target_modules.
271
+
272
+ Returns:
273
+ True if the module passes the filter, False otherwise.
274
+ """
275
+ if target_modules is None:
276
+ return True
277
+ module_suffix = module_name.split(".")[-1]
278
+ return module_suffix in set(target_modules)
279
+
280
+
281
def get_adapter_absolute_path(lora_path: str) -> str:
282
"""
283
Resolves the given lora_path to an absolute local path.
vllm/lora/worker_manager.py
+28
-1
@@ -17,7 +17,11 @@ from vllm.lora.model_manager import (
17
)
18
from vllm.lora.peft_helper import PEFTHelper
19
from vllm.lora.request import LoRARequest
20
-from vllm.lora.utils import get_adapter_absolute_path
20
+from vllm.lora.utils import (
21
+ get_adapter_absolute_path,
22
+ is_in_target_modules,
23
+ is_supported_lora_module,
24
+)
25
26
logger = init_logger(__name__)
27
@@ -142,6 +146,29 @@ class WorkerLoRAManager:
146
skip_prefixes=lora_skip_prefixes,
147
)
148
149
+ # Warn about adapter modules that will be ignored.
150
+ target_modules = self.lora_config.target_modules
151
+ for module_name in lora.loras:
152
+ if not is_supported_lora_module(module_name, supported_lora_modules):
153
+ logger.warning_once(
154
+ "LoRA module '%s' in adapter '%s' is not in the "
155
+ "model's supported LoRA target modules [%s]. "
156
+ "These parameters will be ignored, which may "
157
+ "cause abnormal model behavior.",
158
+ module_name,
159
+ lora_request.lora_path,
160
+ ", ".join(sorted(supported_lora_modules)),
161
+ )
162
+ elif not is_in_target_modules(module_name, target_modules):
163
+ logger.warning_once(
164
+ "LoRA module '%s' in adapter '%s' is not in the "
165
+ "deployment-time target_modules restriction [%s]."
166
+ " These parameters will be ignored.",
167
+ module_name,
168
+ lora_request.lora_path,
169
+ ", ".join(sorted(target_modules)),
170
+ )
171
+
172
except FileNotFoundError as e:
173
# FileNotFoundError should be raised if both
174
# - No adapter found to download from huggingface (or in