@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / aa7f37ccfa
Fork 0

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, 00:15 UTC

Add support for LoRA adapters in Nemotron-H models (#30802)

Signed-off-by: Daniel Serebrenik <daserebrenik@nvidia.com>

danisereb committed Jan 19, 2026 at 16:30 UTC aa7f37ccfa168d40c89d1a104a870757666fcb56
10 files changed +497 -27
tests/lora/test_layers.py
+297
@@ -17,6 +17,7 @@ from vllm.lora.layers import (
17 ColumnParallelLinearWithShardedLoRA,
18 LogitsProcessorWithLoRA,
19 LoRAMapping,
20 + MergedColumnParallelLinearVariableSliceWithLoRA,
21 MergedColumnParallelLinearWithLoRA,
22 MergedColumnParallelLinearWithShardedLoRA,
23 MergedQKVParallelLinearWithLoRA,
@@ -850,6 +851,116 @@ def test_column_parallel_packed(
851 torch.testing.assert_close(lora_result, expected_result, rtol=rtol, atol=atol)
852
853
854 +@torch.inference_mode()
855 +@pytest.mark.parametrize("num_loras", [1, 2, 4])
856 +@pytest.mark.parametrize("num_slices", [3, 5])
857 +@pytest.mark.parametrize("device", DEVICES)
858 +@pytest.mark.parametrize("stage", STAGES)
859 +def test_merged_column_parallel_variable_slice(
860 + default_vllm_config, dist_init, num_loras, num_slices, device, stage
861 +) -> None:
862 + if current_platform.is_cuda_alike():
863 + torch.cuda.set_device(device)
864 +
865 + max_loras = 8
866 + torch.set_default_device(device)
867 + lora_config = LoRAConfig(
868 + max_loras=max_loras, max_lora_rank=8, lora_dtype=torch.float16
869 + )
870 + punica_wrapper = get_punica_wrapper(8192, 256, device, lora_config=lora_config)
871 +
872 + # Set number of output slices
873 + output_sizes = [1024 + i * 256 for i in range(num_slices)]
874 + total_output = sum(output_sizes)
875 +
876 + def create_layer():
877 + # Create linear layer
878 + linear = MergedColumnParallelLinear(
879 + 4096, output_sizes, bias=False, params_dtype=torch.float16
880 + )
881 + linear.weight.data = torch.rand_like(linear.weight.data)
882 +
883 + # Create linear layer with LoRA adapter
884 + lora_linear = MergedColumnParallelLinearVariableSliceWithLoRA(linear)
885 + lora_linear.create_lora_weights(max_loras, lora_config)
886 + return linear, lora_linear
887 +
888 + for i in range(NUM_RANDOM_SEEDS):
889 + set_random_seed(i)
890 + id_to_index = get_random_id_to_index(num_loras, max_loras)
891 + linear, lora_linear = create_layer()
892 + lora_linear.set_mapping(punica_wrapper)
893 +
894 + # Populate LoRA weights
895 + lora_dict, sublora_dict = {}, {}
896 + for slot_idx, lora_id in enumerate(id_to_index):
897 + if lora_id is not None:
898 + # Create random LoRA weights
899 + lora_a = torch.rand(8, 4096, dtype=torch.float16, device=device)
900 + lora_b = torch.rand(total_output, 8, dtype=torch.float16, device=device)
901 + lora_linear.set_lora(slot_idx, lora_a, lora_b)
902 + lora_dict[lora_id] = (lora_a, lora_b)
903 +
904 + # Split lora_b for expected computation
905 + sublora_dict[lora_id] = torch.split(lora_b, output_sizes, dim=0)
906 +
907 + inputs, index_mapping, prompt_mapping = create_random_inputs(
908 + active_lora_ids=list(lora_dict.keys()),
909 + num_inputs=32 * num_loras,
910 + input_size=(1, 4096),
911 + input_range=(0, 1),
912 + input_type=torch.float16,
913 + device=device,
914 + )
915 + lora_mapping = LoRAMapping(index_mapping, prompt_mapping, is_prefill=stage)
916 + punica_wrapper.update_metadata(lora_mapping, id_to_index, max_loras, 512)
917 +
918 + # Compute LoRA result
919 + lora_result = lora_linear(torch.cat(inputs))[0]
920 +
921 + # Compute expected result
922 + expected_results = []
923 + for input_, lora_id in zip(inputs, prompt_mapping):
924 + result = linear(input_)[0]
925 + lora_a, _ = lora_dict[lora_id]
926 + offset = 0
927 + # Compute expected result for each sublora
928 + for lora_b_slice in sublora_dict[lora_id]:
929 + sz = lora_b_slice.shape[0]
930 + result[:, offset : offset + sz] += input_ @ lora_a.T @ lora_b_slice.T
931 + offset += sz
932 + expected_results.append(result)
933 +
934 + # Check that the LoRA result is close to the expected result
935 + rtol, atol = TOLERANCES[lora_result.dtype]
936 + torch.testing.assert_close(
937 + lora_result, torch.cat(expected_results), rtol=rtol, atol=atol
938 + )
939 +
940 + # Reset LoRA weights and check results with zero LoRA weights
941 + for slot_idx in range(max_loras):
942 + lora_linear.reset_lora(slot_idx)
943 +
944 + inputs, index_mapping, prompt_mapping = create_random_inputs(
945 + active_lora_ids=[0],
946 + num_inputs=32 * num_loras,
947 + input_size=(1, 4096),
948 + input_range=(0, 1),
949 + input_type=torch.float16,
950 + device=device,
951 + )
952 + lora_mapping = LoRAMapping(index_mapping, prompt_mapping, is_prefill=stage)
953 + punica_wrapper.update_metadata(lora_mapping, id_to_index, max_loras, 512)
954 +
955 + # After resetting LoRA weights,
956 + # lora_linear should behave like the base linear layer
957 + lora_result = lora_linear(torch.cat(inputs))[0]
958 + expected_result = linear(torch.cat(inputs))[0]
959 +
960 + rtol, atol = TOLERANCES[lora_result.dtype]
961 + torch.testing.assert_close(lora_result, expected_result, rtol=rtol, atol=atol)
962 +
963 +
964 @pytest.mark.parametrize("tp_size", [1, 2, 4, 8])
965 @pytest.mark.parametrize(
966 "seed", list(range(VOCAB_PARALLEL_EMBEDDING_TEST_NUM_RANDOM_SEEDS))
@@ -1119,3 +1230,189 @@ def test_get_masked_input_and_mask():
1230 assert torch.equal(
1231 modified_x_rank_3, torch.tensor([0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 4])
1232 )
1233 +
1234 +
1235 +def test_variable_slice_lora_class_selection(default_vllm_config, dist_init):
1236 + """Test that MergedColumnParallelLinearVariableSliceWithLoRA is selected
1237 + only for nemotron-h style models (checkpoint has single weight but layer
1238 + has 3+ output slices).
1239 +
1240 + This verifies that from_layer selects
1241 + MergedColumnParallelLinearVariableSliceWithLoRA
1242 + before ColumnParallelLinearWithLoRA for layers with 3+ output sizes, since
1243 + ColumnParallelLinearWithLoRA's slice_lora_b assumes exactly 2 slices.
1244 + """
1245 + from vllm.lora.utils import from_layer
1246 +
1247 + lora_config = LoRAConfig(max_loras=8, max_lora_rank=8, lora_dtype=torch.float16)
1248 +
1249 + # Case 1: MergedColumnParallelLinear with 3+ output sizes and
1250 + # packed_modules_list with 1 item (nemotron-h style)
1251 + # -> MergedColumnParallelLinearVariableSliceWithLoRA should be selected
1252 + layer_3_slices = MergedColumnParallelLinear(
1253 + 4096, [1024, 1280, 1536], bias=False, params_dtype=torch.float16
1254 + )
1255 + packed_modules_single = ["mlp"]
1256 +
1257 + assert MergedColumnParallelLinearVariableSliceWithLoRA.can_replace_layer(
1258 + source_layer=layer_3_slices,
1259 + lora_config=lora_config,
1260 + packed_modules_list=packed_modules_single,
1261 + ), "MergedColumnParallelLinearVariableSliceWithLoRA should handle 3+ slices"
1262 +
1263 + # ColumnParallelLinearWithLoRA should NOT match 3+ slices
1264 + # (its slice_lora_b assumes exactly 2 slices)
1265 + assert not ColumnParallelLinearWithLoRA.can_replace_layer(
1266 + source_layer=layer_3_slices,
1267 + lora_config=lora_config,
1268 + packed_modules_list=packed_modules_single,
1269 + ), (
1270 + "ColumnParallelLinearWithLoRA should NOT handle 3+ slices "
1271 + "(slice_lora_b assumes 2 slices)"
1272 + )
1273 +
1274 + # Verify from_layer selects the correct class (Variable, not base)
1275 + selected_layer = from_layer(
1276 + layer_3_slices,
1277 + max_loras=8,
1278 + lora_config=lora_config,
1279 + packed_modules_list=packed_modules_single,
1280 + )
1281 + assert isinstance(
1282 + selected_layer, MergedColumnParallelLinearVariableSliceWithLoRA
1283 + ), (
1284 + f"from_layer should select MergedColumnParallelLinearVariableSliceWithLoRA "
1285 + f"for 3+ slices, got {type(selected_layer).__name__}"
1286 + )
1287 +
1288 + # Case 2: MergedColumnParallelLinear with 2 output sizes and
1289 + # packed_modules_list with 1 item (standard gate_up style)
1290 + # -> ColumnParallelLinearWithLoRA should be selected
1291 + # -> MergedColumnParallelLinearVariableSliceWithLoRA should NOT match
1292 + layer_2_slices = MergedColumnParallelLinear(
1293 + 4096, [2048, 2048], bias=False, params_dtype=torch.float16
1294 + )
1295 +
1296 + assert ColumnParallelLinearWithLoRA.can_replace_layer(
1297 + source_layer=layer_2_slices,
1298 + lora_config=lora_config,
1299 + packed_modules_list=packed_modules_single,
1300 + ), "ColumnParallelLinearWithLoRA should handle 2 slices"
1301 +
1302 + assert not MergedColumnParallelLinearVariableSliceWithLoRA.can_replace_layer(
1303 + source_layer=layer_2_slices,
1304 + lora_config=lora_config,
1305 + packed_modules_list=packed_modules_single,
1306 + ), "MergedColumnParallelLinearVariableSliceWithLoRA should NOT handle 2 slices"
1307 +
1308 + # Verify from_layer selects ColumnParallelLinearWithLoRA for 2 slices
1309 + selected_layer_2 = from_layer(
1310 + layer_2_slices,
1311 + max_loras=8,
1312 + lora_config=lora_config,
1313 + packed_modules_list=packed_modules_single,
1314 + )
1315 + assert isinstance(selected_layer_2, ColumnParallelLinearWithLoRA), (
1316 + f"from_layer should select ColumnParallelLinearWithLoRA "
1317 + f"for 2 slices, got {type(selected_layer_2).__name__}"
1318 + )
1319 + # But NOT the Variable subclass
1320 + assert not isinstance(
1321 + selected_layer_2, MergedColumnParallelLinearVariableSliceWithLoRA
1322 + ), (
1323 + "from_layer should NOT select "
1324 + "MergedColumnParallelLinearVariableSliceWithLoRA for 2 slices"
1325 + )
1326 +
1327 + # Case 3: MergedColumnParallelLinear with 3+ items in packed_modules_list
1328 + # -> MergedColumnParallelLinearVariableSliceWithLoRA should be selected
1329 + packed_modules_three = ["gate_proj", "up_proj", "down_proj"]
1330 +
1331 + assert MergedColumnParallelLinearVariableSliceWithLoRA.can_replace_layer(
1332 + source_layer=layer_3_slices,
1333 + lora_config=lora_config,
1334 + packed_modules_list=packed_modules_three,
1335 + ), "MergedColumnParallelLinearVariableSliceWithLoRA should handle 3+ packed modules"
1336 +
1337 + # Case 4: MergedColumnParallelLinear with 2 items in packed_modules_list
1338 + # -> MergedColumnParallelLinearWithLoRA should handle this (not Variable)
1339 + packed_modules_two = ["gate_proj", "up_proj"]
1340 +
1341 + assert not MergedColumnParallelLinearVariableSliceWithLoRA.can_replace_layer(
1342 + source_layer=layer_2_slices,
1343 + lora_config=lora_config,
1344 + packed_modules_list=packed_modules_two,
1345 + ), (
1346 + "MergedColumnParallelLinearVariableSliceWithLoRA"
1347 + " should NOT handle 2 packed modules"
1348 + )
1349 +
1350 + assert MergedColumnParallelLinearWithLoRA.can_replace_layer(
1351 + source_layer=layer_2_slices,
1352 + lora_config=lora_config,
1353 + packed_modules_list=packed_modules_two,
1354 + ), "MergedColumnParallelLinearWithLoRA should handle 2 packed modules"
1355 +
1356 + # Verify from_layer selects MergedColumnParallelLinearWithLoRA for 2 packed modules
1357 + selected_layer_merged = from_layer(
1358 + layer_2_slices,
1359 + max_loras=8,
1360 + lora_config=lora_config,
1361 + packed_modules_list=packed_modules_two,
1362 + )
1363 + assert isinstance(selected_layer_merged, MergedColumnParallelLinearWithLoRA), (
1364 + f"from_layer should select MergedColumnParallelLinearWithLoRA "
1365 + f"for 2 packed modules, got {type(selected_layer_merged).__name__}"
1366 + )
1367 +
1368 + # Case 5: Plain ColumnParallelLinear (not merged) - common in many models
1369 + # -> ColumnParallelLinearWithLoRA should be selected
1370 + plain_column_parallel = ColumnParallelLinear(
1371 + 4096, 4096, bias=False, params_dtype=torch.float16
1372 + )
1373 +
1374 + assert ColumnParallelLinearWithLoRA.can_replace_layer(
1375 + source_layer=plain_column_parallel,
1376 + lora_config=lora_config,
1377 + packed_modules_list=packed_modules_single,
1378 + ), "ColumnParallelLinearWithLoRA should handle plain ColumnParallelLinear"
1379 +
1380 + assert not MergedColumnParallelLinearVariableSliceWithLoRA.can_replace_layer(
1381 + source_layer=plain_column_parallel,
1382 + lora_config=lora_config,
1383 + packed_modules_list=packed_modules_single,
1384 + ), (
1385 + "MergedColumnParallelLinearVariableSliceWithLoRA "
1386 + "should NOT handle plain ColumnParallelLinear"
1387 + )
1388 +
1389 + # Verify from_layer selects ColumnParallelLinearWithLoRA for plain layer
1390 + selected_plain = from_layer(
1391 + plain_column_parallel,
1392 + max_loras=8,
1393 + lora_config=lora_config,
1394 + packed_modules_list=packed_modules_single,
1395 + )
1396 + assert isinstance(selected_plain, ColumnParallelLinearWithLoRA), (
1397 + f"from_layer should select ColumnParallelLinearWithLoRA "
1398 + f"for plain ColumnParallelLinear, got {type(selected_plain).__name__}"
1399 + )
1400 +
1401 + # Case 6: MergedColumnParallelLinear with exactly 2 output sizes
1402 + # and empty packed_modules_list
1403 + # -> ColumnParallelLinearWithLoRA should NOT match (packed_modules_list != 1)
1404 + # -> MergedColumnParallelLinearVariableSliceWithLoRA should NOT match (< 3 slices)
1405 + assert not ColumnParallelLinearWithLoRA.can_replace_layer(
1406 + source_layer=layer_2_slices,
1407 + lora_config=lora_config,
1408 + packed_modules_list=[],
1409 + ), "ColumnParallelLinearWithLoRA should NOT handle empty packed_modules_list"
1410 +
1411 + assert not MergedColumnParallelLinearVariableSliceWithLoRA.can_replace_layer(
1412 + source_layer=layer_2_slices,
1413 + lora_config=lora_config,
1414 + packed_modules_list=[],
1415 + ), (
1416 + "MergedColumnParallelLinearVariableSliceWithLoRA "
1417 + "should NOT handle 2 slices even with empty packed_modules_list"
1418 + )
vllm/lora/layers/__init__.py
+2
@@ -4,6 +4,7 @@ from vllm.lora.layers.base import BaseLayerWithLoRA
4 from vllm.lora.layers.column_parallel_linear import (
5 ColumnParallelLinearWithLoRA,
6 ColumnParallelLinearWithShardedLoRA,
7 + MergedColumnParallelLinearVariableSliceWithLoRA,
8 MergedColumnParallelLinearWithLoRA,
9 MergedColumnParallelLinearWithShardedLoRA,
10 MergedQKVParallelLinearWithLoRA,
@@ -29,6 +30,7 @@ __all__ = [
30 "ColumnParallelLinearWithShardedLoRA",
31 "MergedColumnParallelLinearWithLoRA",
32 "MergedColumnParallelLinearWithShardedLoRA",
33 + "MergedColumnParallelLinearVariableSliceWithLoRA",
34 "MergedQKVParallelLinearWithLoRA",
35 "MergedQKVParallelLinearWithShardedLoRA",
36 "QKVParallelLinearWithLoRA",
vllm/lora/layers/column_parallel_linear.py
+85 -4
@@ -155,10 +155,19 @@ class ColumnParallelLinearWithLoRA(BaseLinearLayerWithLoRA):
155 packed_modules_list: list,
156 model_config: PretrainedConfig | None = None,
157 ) -> bool:
158 - return type(source_layer) is ColumnParallelLinear or (
159 - type(source_layer) is MergedColumnParallelLinear
160 - and len(packed_modules_list) == 1
161 - )
158 + if type(source_layer) is ColumnParallelLinear:
159 + return True
160 + if type(source_layer) is MergedColumnParallelLinear:
161 + if len(packed_modules_list) != 1:
162 + return False
163 + # Exclude layers with 3+ output sizes - those are handled by
164 + # MergedColumnParallelLinearVariableSliceWithLoRA since this
165 + # class's slice_lora_b assumes exactly 2 slices.
166 + return not (
167 + hasattr(source_layer, "output_sizes")
168 + and len(source_layer.output_sizes) >= 3
169 + )
170 + return False
171
172
173 class MergedColumnParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):
@@ -575,3 +584,75 @@ class MergedQKVParallelLinearWithShardedLoRA(MergedQKVParallelLinearWithLoRA):
584 model_config=model_config,
585 decorate=False,
586 )
587 +
588 +
589 +class MergedColumnParallelLinearVariableSliceWithLoRA(
590 + MergedColumnParallelLinearWithLoRA
591 +):
592 + """MergedColumnParallelLinear with variable number of slices (3+).
593 +
594 + This handles cases where the checkpoint has a single weight for the whole
595 + module (not split into slices), but the layer itself has multiple slices.
596 + """
597 +
598 + @classmethod
599 + @_not_fully_sharded_can_replace
600 + def can_replace_layer(
601 + cls,
602 + source_layer: nn.Module,
603 + lora_config: LoRAConfig,
604 + packed_modules_list: list,
605 + model_config: PretrainedConfig | None = None,
606 + ) -> bool:
607 + # Support MergedColumnParallelLinear with 3 or more slices
608 + # (2 slices are handled by MergedColumnParallelLinearWithLoRA)
609 + if type(source_layer) is not MergedColumnParallelLinear:
610 + return False
611 +
612 + # If packed_modules_list has 3+ items, use this class
613 + if len(packed_modules_list) >= 3:
614 + return True
615 +
616 + # If packed_modules_list has exactly 2 items, let
617 + # MergedColumnParallelLinearWithLoRA handle it
618 + if len(packed_modules_list) == 2:
619 + return False
620 +
621 + # If packed_modules_list is empty or has 1 item,
622 + # check the layer's output_sizes.
623 + # This handles cases where the checkpoint has a single weight
624 + # but the layer has multiple slices (3+)
625 + return (
626 + hasattr(source_layer, "output_sizes")
627 + and len(source_layer.output_sizes) >= 3
628 + )
629 +
630 + def set_lora(
631 + self,
632 + index: int,
633 + lora_a: torch.Tensor | list[torch.Tensor],
634 + lora_b: torch.Tensor | list[torch.Tensor],
635 + ):
636 + """Override to handle single tensor weights
637 + that need to be split into slices."""
638 + self.reset_lora(index)
639 +
640 + # Handle case where checkpoint has single tensor weights
641 + # lora_a shape: (rank, input_size) - same for all slices, duplicate it
642 + if isinstance(lora_a, torch.Tensor):
643 + lora_a = [lora_a] * self.n_slices
644 +
645 + # lora_b shape: (total_output_size, rank) -
646 + # split along dim 0 based on output_sizes
647 + if isinstance(lora_b, torch.Tensor):
648 + output_sizes = self.base_layer.output_sizes
649 + lora_b_list = []
650 + start_idx = 0
651 + for output_size in output_sizes:
652 + end_idx = start_idx + output_size
653 + lora_b_list.append(lora_b[start_idx:end_idx, :])
654 + start_idx = end_idx
655 + lora_b = lora_b_list
656 +
657 + # Now call parent's set_lora which expects lists
658 + super().set_lora(index, lora_a, lora_b)
vllm/lora/layers/fused_moe.py
+25 -17
@@ -52,7 +52,9 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
52 self.tp_size = get_tensor_model_parallel_world_size()
53 self.tp_rank = get_tensor_model_parallel_rank()
54 self.device = _get_lora_device(base_layer)
55 - self._w13_slices = 2
55 + # For non-gated MoE (is_act_and_mul=False), only 1 slice is needed
56 + # since there's only up_proj (w1), not gate_proj + up_proj (w1 + w3)
57 + self._w13_slices = 2 if base_layer.moe_config.is_act_and_mul else 1
58 self._inject_lora_into_fused_moe()
59
60 def _normalize_keys(self, config: dict[str, int | None]) -> dict[str, int | None]:
@@ -400,7 +402,8 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
402 self.lora_b_stacked = []
403 for lora_id in range(max_loras):
404 for experts_id in range(self.base_layer.local_num_experts):
403 - # gate_proj,down_proj,up_proj
405 + # For gated MoE: gate_proj (w1), down_proj (w2), up_proj (w3)
406 + # For non-gated MoE: up_proj (w1), down_proj (w2)
407 self.lora_a_stacked.append(
408 self.w13_lora_a_stacked[0][lora_id][experts_id]
409 )
@@ -415,12 +418,14 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
418 self.w2_lora_b_stacked[0][lora_id][experts_id]
419 )
420
418 - self.lora_a_stacked.append(
419 - self.w13_lora_a_stacked[1][lora_id][experts_id]
420 - )
421 - self.lora_b_stacked.append(
422 - self.w13_lora_b_stacked[1][lora_id][experts_id]
423 - )
421 + # Only add w3 (up_proj) for gated MoE (_w13_slices == 2)
422 + if self._w13_slices == 2:
423 + self.lora_a_stacked.append(
424 + self.w13_lora_a_stacked[1][lora_id][experts_id]
425 + )
426 + self.lora_b_stacked.append(
427 + self.w13_lora_b_stacked[1][lora_id][experts_id]
428 + )
429
430 def _slice_w13_a(self, w13_lora_a: torch.Tensor) -> torch.Tensor:
431 """
@@ -515,8 +520,6 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
520
521 slliced_w1_lora_a = self._slice_w13_a(w1_lora_a)
522 slliced_w1_lora_b = self._slice_w13_b(w1_lora_b)
518 - slliced_w3_lora_a = self._slice_w13_a(w3_lora_a)
519 - slliced_w3_lora_b = self._slice_w13_b(w3_lora_b)
523
524 sliced_w2_lora_a = self._slice_w2_a(w2_lora_a)
525 sliced_w2_lora_b = self._slice_w2_b(w2_lora_b)
@@ -525,17 +528,22 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
528 index, :, : slliced_w1_lora_a.shape[1], : slliced_w1_lora_a.shape[2]
529 ].copy_(slliced_w1_lora_a, non_blocking=True)
530
528 - self.w13_lora_a_stacked[1][
529 - index, :, : slliced_w3_lora_a.shape[1], : slliced_w3_lora_a.shape[2]
530 - ].copy_(slliced_w3_lora_a, non_blocking=True)
531 -
531 self.w13_lora_b_stacked[0][
532 index, :, : slliced_w1_lora_b.shape[1], : slliced_w1_lora_b.shape[2]
533 ].copy_(slliced_w1_lora_b, non_blocking=True)
534
536 - self.w13_lora_b_stacked[1][
537 - index, :, : slliced_w3_lora_b.shape[1], : slliced_w3_lora_b.shape[2]
538 - ].copy_(slliced_w3_lora_b, non_blocking=True)
535 + # Only copy w3 (up_proj) for gated MoE (_w13_slices == 2)
536 + if self._w13_slices == 2:
537 + slliced_w3_lora_a = self._slice_w13_a(w3_lora_a)
538 + slliced_w3_lora_b = self._slice_w13_b(w3_lora_b)
539 +
540 + self.w13_lora_a_stacked[1][
541 + index, :, : slliced_w3_lora_a.shape[1], : slliced_w3_lora_a.shape[2]
542 + ].copy_(slliced_w3_lora_a, non_blocking=True)
543 +
544 + self.w13_lora_b_stacked[1][
545 + index, :, : slliced_w3_lora_b.shape[1], : slliced_w3_lora_b.shape[2]
546 + ].copy_(slliced_w3_lora_b, non_blocking=True)
547
548 self.w2_lora_a_stacked[0][
549 index, :, : sliced_w2_lora_a.shape[1], : sliced_w2_lora_a.shape[2]
vllm/lora/lora_weights.py
+26 -3
@@ -154,7 +154,10 @@ class PackedLoRALayerWeights(LoRALayerWeights):
154
155 @classmethod
156 def pack_moe(
157 - cls, loras: GenericSequence[Optional["LoRALayerWeights"]], module_name: str
157 + cls,
158 + loras: GenericSequence[Optional["LoRALayerWeights"]],
159 + module_name: str,
160 + is_non_gated_moe: bool = False,
161 ) -> "PackedLoRALayerWeights":
162 """Pack a list of LoRAs into a single LoRA.
163
@@ -177,6 +180,11 @@ class PackedLoRALayerWeights(LoRALayerWeights):
180 w1_lora = loras[eid * 3]
181 w2_lora = loras[eid * 3 + 1]
182 w3_lora = loras[eid * 3 + 2]
183 + # For non-gated MoE, w3 is not used, so we use w1's LoRA weights
184 + # This is determined by checking the expert mapping (get_expert_mapping)
185 + # which indicates when ckpt_up_proj_name is empty.
186 + if w3_lora is None and is_non_gated_moe:
187 + w3_lora = w1_lora
188 assert w1_lora is not None
189 assert w2_lora is not None
190 assert w3_lora is not None
@@ -191,10 +199,24 @@ class PackedLoRALayerWeights(LoRALayerWeights):
199
200 w1_lora_a = torch.stack(w1_lora_a_lst, dim=0) # (num_experts,rank,input_size)
201 w2_lora_a = torch.stack(w2_lora_a_lst, dim=0)
194 - w3_lora_a = torch.stack(w3_lora_a_lst, dim=0)
202 w1_lora_b = torch.stack(w1_lora_b_lst, dim=0) # (num_experts,output_size,rank)
203 w2_lora_b = torch.stack(w2_lora_b_lst, dim=0)
197 - w3_lora_b = torch.stack(w3_lora_b_lst, dim=0)
204 +
205 + # All w1, w2, w3 have the same scaling factor.
206 + scaling = lora_alpha / rank
207 + last_scaling = scaling
208 +
209 + if is_non_gated_moe:
210 + # For non-gated MoE, reuse w1 tensors for w3 to avoid memory waste
211 + # w3_lora_a_lst and w3_lora_b_lst are not relevant in this case
212 + w3_lora_a = w1_lora_a
213 + w3_lora_b = w1_lora_b
214 +
215 + # For non-gated MoE, avoid double-scaling by setting w3's scaling to 1.
216 + last_scaling = 1.0
217 + else:
218 + w3_lora_a = torch.stack(w3_lora_a_lst, dim=0)
219 + w3_lora_b = torch.stack(w3_lora_b_lst, dim=0)
220
221 obj = cls(
222 module_name,
@@ -202,6 +224,7 @@ class PackedLoRALayerWeights(LoRALayerWeights):
224 [lora_alpha, lora_alpha, lora_alpha],
225 [w1_lora_a, w2_lora_a, w3_lora_a],
226 [w1_lora_b, w2_lora_b, w3_lora_b],
227 + scaling=[scaling, scaling, last_scaling],
228 )
229 return obj
230
vllm/lora/model_manager.py
+47 -3
@@ -104,7 +104,9 @@ class LoRAModelManager:
104 self.modules: dict[str, BaseLayerWithLoRA] = {}
105 # Dict instead of a set for compatibility with LRUCache.
106 self._last_mapping: LoRAMapping | None = None
107 - self._is_3d_moe_model = is_moe_model(self.model) and self.model.is_3d_moe_weight
107 + is_moe = is_moe_model(self.model)
108 + self._is_3d_moe_model = is_moe and self.model.is_3d_moe_weight
109 + self._is_non_gated_moe = is_moe and self.model.is_non_gated_moe
110 self._init_punica_wrapper(max_num_batched_tokens, vllm_config)
111 self._create_lora_modules()
112
@@ -339,6 +341,20 @@ class LoRAModelManager:
341 )
342 continue
343
344 + # TODO: Remove this restriction
345 + # peft error when generating LoRA adapter with "gate" module:
346 + # "Target module NemotronHTopkRouter() is not supported."
347 + # Working LoRA adapter was created using peft with:
348 + # LoraConfig(target_modules="all-linear", ...)
349 + if self._is_non_gated_moe and module_name.endswith("mixer.gate"):
350 + logger.debug_once(
351 + "LoRA is not supported for non-gated MoE gate module."
352 + " %s will be ignored.",
353 + module_name,
354 + scope="local",
355 + )
356 + continue
357 +
358 parts = module_name.split(".")[-1]
359 packed_moduled_lst = self.packed_modules_mapping.get(parts, [])
360 if isinstance(module, FusedMoE):
@@ -405,6 +421,22 @@ class LoRAModelManager:
421 )
422 self.modules[module_name] = module
423
424 + @staticmethod
425 + def _pad_lora_pairs_to_triplets(
426 + loras: list[LoRALayerWeights | None],
427 + ) -> list[LoRALayerWeights | None]:
428 + """Pad LoRA weight pairs to triplets for non-gated MoE.
429 +
430 + For non-gated MoE, each expert has 2 entries (w1, w2) that need to be
431 + padded to triplets (w1, w2, None) to match pack_moe expectations.
432 + """
433 + assert len(loras) % 2 == 0, "Expected pairs of LoRA weights for non-gated MoE."
434 + padded: list[LoRALayerWeights | None] = []
435 + for i in range(0, len(loras), 2):
436 + padded.extend(loras[i : i + 2])
437 + padded.append(None)
438 + return padded
439 +
440 def create_dummy_lora(
441 self,
442 lora_id: int,
@@ -491,7 +523,13 @@ class LoRAModelManager:
523 )
524 subloras.append(lora)
525 if module.__class__.__name__ == "FusedMoEWithLoRA":
494 - lora = PackedLoRALayerWeights.pack_moe(subloras, module_name)
526 + # For non-gated MoE, pad subloras to 3 elements per expert
527 + # to match pack_moe expectations (w1, w2, None for w3)
528 + if self._is_non_gated_moe and len(subloras) > 0:
529 + subloras = self._pad_lora_pairs_to_triplets(subloras)
530 + lora = PackedLoRALayerWeights.pack_moe(
531 + subloras, module_name, is_non_gated_moe=self._is_non_gated_moe
532 + )
533 else:
534 lora = PackedLoRALayerWeights.pack(subloras)
535 model.loras[module_name] = lora
@@ -559,8 +597,14 @@ class LoRAModelManager:
597 if lora_model.check_lora_name(module_name):
598 module_name = replaced_module_name
599 if module_name.endswith(".experts"):
600 + if self._is_non_gated_moe and len(replacement_loras) > 0:
601 + replacement_loras = self._pad_lora_pairs_to_triplets(
602 + replacement_loras
603 + )
604 lora_model.loras[module_name] = PackedLoRALayerWeights.pack_moe(
563 - replacement_loras, module_name
605 + replacement_loras,
606 + module_name,
607 + is_non_gated_moe=self._is_non_gated_moe,
608 )
609 else:
610 lora_model.loras[module_name] = PackedLoRALayerWeights.pack(
vllm/lora/utils.py
+6
@@ -25,6 +25,7 @@ from vllm.lora.layers import (
25 FusedMoE3DWithLoRA,
26 FusedMoEWithLoRA,
27 LogitsProcessorWithLoRA,
28 + MergedColumnParallelLinearVariableSliceWithLoRA,
29 MergedColumnParallelLinearWithLoRA,
30 MergedColumnParallelLinearWithShardedLoRA,
31 MergedQKVParallelLinearWithLoRA,
@@ -68,6 +69,7 @@ _all_lora_classes: set[type[BaseLayerWithLoRA]] = {
69 ColumnParallelLinearWithShardedLoRA,
70 QKVParallelLinearWithShardedLoRA,
71 MergedColumnParallelLinearWithShardedLoRA,
72 + MergedColumnParallelLinearVariableSliceWithLoRA,
73 MergedQKVParallelLinearWithShardedLoRA,
74 RowParallelLinearWithShardedLoRA,
75 FusedMoEWithLoRA,
@@ -266,9 +268,13 @@ def process_packed_modules_mapping(model: nn.Module) -> dict[str, list[str]]:
268 packed_modules_mapping = get_packed_modules_mapping(model)
269 if not model.is_3d_moe_weight:
270 # 3D MoE LoRA does not need `packed_modules_mapping`
271 + # Filter out malformed entries: non-gated MoE has empty
272 + # ckpt_up_proj_name which results in weight_name containing ".."
273 + # (e.g., "experts.0.." instead of "experts.0.layer_name.")
274 packed_modules_mapping["experts"] = [
275 weight_name.rstrip(".")
276 for _, weight_name, _, _ in moe_packed_mapping
277 + if ".." not in weight_name
278 ]
279
280 return packed_modules_mapping
vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py
+5
@@ -227,6 +227,11 @@ class FlashInferExperts(mk.FusedMoEPermuteExpertsUnpermute):
227 use_deepseek_fp8_block_scale=self.use_deepseek_fp8_block_scale,
228 )
229
230 + def moe_sum(self, input: torch.Tensor, output: torch.Tensor) -> None:
231 + # No support for LoRA in flashinfer_cutlass_fused_moe.
232 + # See TODOs in flashinfer functions runMoe and runMoeMinLantency.
233 + raise NotImplementedError("LoRA is not supported for flashinfer_cutlass_moe")
234 +
235
236 def flashinfer_cutlass_moe_fp4(
237 hidden_states: torch.Tensor,
vllm/model_executor/models/interfaces.py
+1
@@ -376,6 +376,7 @@ class SupportsLoRA(Protocol):
376 MRO of your model class.
377 """
378 is_3d_moe_weight: ClassVar[bool] = False
379 + is_non_gated_moe: ClassVar[bool] = False
380 # The `embedding_module` and `embedding_padding_modules`
381 # are empty by default.
382 embedding_modules: ClassVar[dict[str, str]] = {}
vllm/model_executor/models/nemotron_h.py
+3
@@ -747,6 +747,9 @@ class NemotronHForCausalLM(
747 MixtureOfExperts,
748 SupportsMambaPrefixCaching,
749 ):
750 + # Relevant only if self.has_moe is True
751 + is_non_gated_moe: bool = True
752 +
753 hf_to_vllm_mapper = WeightsMapper(
754 orig_to_new_prefix={"backbone": "model"},
755 orig_to_new_substr={"A_log": "A", "embeddings": "embed_tokens"},