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 14, 05:30 UTC
[CI/Build] Reduce LoRA layer test cases (#23721)
Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>
Jee Jee Li committed
Aug 27, 2025 at 18:59 UTC
e03940762b43812fccd3c214bda60201cff9d16a
1 file changed
+33
-39
tests/lora/test_layers.py
+33
-39
@@ -243,7 +243,7 @@ def check_punica_wrapper(punica_wrapper) -> bool:
243
244
245
@torch.inference_mode()
246
-@pytest.mark.parametrize("num_loras", [1, 2, 4, 8])
246
+@pytest.mark.parametrize("num_loras", [1, 2, 4])
247
@pytest.mark.parametrize("device", DEVICES)
248
@pytest.mark.parametrize("vocab_size", [512, 32000, 64000, 128000])
249
@pytest.mark.parametrize("stage", STAGES)
@@ -347,7 +347,7 @@ def test_embeddings(dist_init, num_loras, device, vocab_size, stage) -> None:
347
@torch.inference_mode()
348
# @pytest.mark.skip(
349
# reason="Fails when loras are in any slot other than the first.")
350
-@pytest.mark.parametrize("num_loras", [1, 2, 4, 8])
350
+@pytest.mark.parametrize("num_loras", [1, 2, 4])
351
@pytest.mark.parametrize("device", DEVICES)
352
@pytest.mark.parametrize("vocab_size", [512, 32000, 64000, 128000])
353
@pytest.mark.parametrize("stage", STAGES)
@@ -486,7 +486,7 @@ def test_embeddings_with_new_embeddings(dist_init, num_loras, device,
486
487
488
@torch.inference_mode()
489
-@pytest.mark.parametrize("num_loras", [1, 2, 4, 8])
489
+@pytest.mark.parametrize("num_loras", [1, 2, 4])
490
@pytest.mark.parametrize("device", DEVICES)
491
@pytest.mark.parametrize("vocab_size", [512, 32000, 64000, 256512])
492
@pytest.mark.parametrize("stage", STAGES)
@@ -620,12 +620,15 @@ def test_lm_head_logits_processor(dist_init, num_loras, device, vocab_size,
620
621
622
@torch.inference_mode()
623
-@pytest.mark.parametrize("num_loras", [1, 2, 4, 8])
623
+@pytest.mark.parametrize("num_loras", [1, 2, 4])
624
@pytest.mark.parametrize("device", DEVICES)
625
@pytest.mark.parametrize("stage", STAGES)
626
-@pytest.mark.parametrize("bias_enabled", [True, False])
627
-def test_linear_replicated(dist_init, num_loras, device, stage,
628
- bias_enabled) -> None:
626
+def test_linear_replicated(
627
+ dist_init,
628
+ num_loras,
629
+ device,
630
+ stage,
631
+) -> None:
632
633
if current_platform.is_cuda_alike():
634
torch.cuda.set_device(device)
@@ -634,10 +637,11 @@ def test_linear_replicated(dist_init, num_loras, device, stage,
637
torch.set_default_device(device)
638
punica_wrapper = get_punica_wrapper(8192, 256, device, max_loras=max_loras)
639
assert check_punica_wrapper(punica_wrapper)
637
- lora_config = LoRAConfig(max_loras=max_loras,
638
- max_lora_rank=8,
639
- lora_dtype=torch.float16,
640
- bias_enabled=bias_enabled)
640
+ lora_config = LoRAConfig(
641
+ max_loras=max_loras,
642
+ max_lora_rank=8,
643
+ lora_dtype=torch.float16,
644
+ )
645
646
def create_random_linear_replicated_layer():
647
@@ -651,10 +655,6 @@ def test_linear_replicated(dist_init, num_loras, device, stage,
655
lora_linear.create_lora_weights(max_loras, lora_config)
656
assert (lora_linear.n_slices == len(lora_linear.lora_a_stacked) == len(
657
lora_linear.lora_b_stacked) == 1)
654
- if bias_enabled:
655
- assert len(lora_linear.lora_bias_stacked) == lora_linear.n_slices
656
- else:
657
- assert lora_linear.lora_bias_stacked is None
658
return linear, lora_linear
659
660
for i in range(NUM_RANDOM_SEEDS):
@@ -734,14 +734,13 @@ def test_linear_replicated(dist_init, num_loras, device, stage,
734
735
736
@torch.inference_mode()
737
-@pytest.mark.parametrize("num_loras", [1, 2, 4, 8])
737
+@pytest.mark.parametrize("num_loras", [1, 2, 4])
738
@pytest.mark.parametrize("orientation", ["row", "column"])
739
@pytest.mark.parametrize("fully_shard", [True, False])
740
@pytest.mark.parametrize("device", DEVICES)
741
@pytest.mark.parametrize("stage", STAGES)
742
-@pytest.mark.parametrize("bias_enabled", [True, False])
742
def test_linear_parallel(dist_init, num_loras, orientation, fully_shard,
744
- device, stage, bias_enabled) -> None:
743
+ device, stage) -> None:
744
745
if current_platform.is_cuda_alike():
746
torch.cuda.set_device(device)
@@ -750,11 +749,12 @@ def test_linear_parallel(dist_init, num_loras, orientation, fully_shard,
749
torch.set_default_device(device)
750
punica_wrapper = get_punica_wrapper(8192, 256, device, max_loras=max_loras)
751
assert check_punica_wrapper(punica_wrapper)
753
- lora_config = LoRAConfig(max_loras=max_loras,
754
- max_lora_rank=8,
755
- fully_sharded_loras=fully_shard,
756
- lora_dtype=torch.float16,
757
- bias_enabled=bias_enabled)
752
+ lora_config = LoRAConfig(
753
+ max_loras=max_loras,
754
+ max_lora_rank=8,
755
+ fully_sharded_loras=fully_shard,
756
+ lora_dtype=torch.float16,
757
+ )
758
759
def create_random_linear_parallel_layer():
760
if orientation == "row":
@@ -777,10 +777,7 @@ def test_linear_parallel(dist_init, num_loras, orientation, fully_shard,
777
lora_linear.create_lora_weights(max_loras, lora_config)
778
assert (lora_linear.n_slices == len(lora_linear.lora_a_stacked) == len(
779
lora_linear.lora_b_stacked) == 1)
780
- if bias_enabled:
781
- assert len(lora_linear.lora_bias_stacked) == lora_linear.n_slices
782
- else:
783
- assert lora_linear.lora_bias_stacked is None
780
+
781
return linear, lora_linear
782
783
for i in range(NUM_RANDOM_SEEDS):
@@ -860,14 +857,13 @@ def test_linear_parallel(dist_init, num_loras, orientation, fully_shard,
857
858
859
@torch.inference_mode()
863
-@pytest.mark.parametrize("num_loras", [1, 2, 4, 8])
860
+@pytest.mark.parametrize("num_loras", [1, 2, 4])
861
@pytest.mark.parametrize("repeats", [1, 2, 3])
862
@pytest.mark.parametrize("fully_shard", [True, False])
863
@pytest.mark.parametrize("device", DEVICES)
864
@pytest.mark.parametrize("stage", STAGES)
868
-@pytest.mark.parametrize("bias_enabled", [True, False])
865
def test_column_parallel_packed(dist_init, num_loras, repeats, fully_shard,
870
- device, stage, bias_enabled) -> None:
866
+ device, stage) -> None:
867
868
if current_platform.is_cuda_alike():
869
torch.cuda.set_device(device)
@@ -876,11 +872,12 @@ def test_column_parallel_packed(dist_init, num_loras, repeats, fully_shard,
872
torch.set_default_device(device)
873
punica_wrapper = get_punica_wrapper(8192, 256, device, max_loras=max_loras)
874
assert check_punica_wrapper(punica_wrapper)
879
- lora_config = LoRAConfig(max_loras=max_loras,
880
- max_lora_rank=8,
881
- fully_sharded_loras=fully_shard,
882
- lora_dtype=torch.float16,
883
- bias_enabled=bias_enabled)
875
+ lora_config = LoRAConfig(
876
+ max_loras=max_loras,
877
+ max_lora_rank=8,
878
+ fully_sharded_loras=fully_shard,
879
+ lora_dtype=torch.float16,
880
+ )
881
882
def create_column_parallel_packed_layer():
883
if repeats == 2:
@@ -924,10 +921,7 @@ def test_column_parallel_packed(dist_init, num_loras, repeats, fully_shard,
921
model_config=FakeConfig())
922
assert (lora_linear.n_slices == len(lora_linear.lora_a_stacked) == len(
923
lora_linear.lora_b_stacked) == n_slices)
927
- if bias_enabled:
928
- assert len(lora_linear.lora_bias_stacked) == lora_linear.n_slices
929
- else:
930
- assert lora_linear.lora_bias_stacked is None
924
+
925
return linear, lora_linear
926
927
for i in range(NUM_RANDOM_SEEDS):