[CI/Build]Reduce the time consumption for LoRA tests (vllm-project#7396)

ShangmingCai · Aug 14, 2024 · 9799280 · 9799280
1 parent 59edd0f
commit 9799280
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 128 deletions.
diff --git a/tests/lora/test_layer_variation.py b/tests/lora/test_layer_variation.py
diff --git a/tests/lora/test_punica_sizes.py b/tests/lora/test_punica_sizes.py
@@ -98,7 +98,7 @@
     128256,
 ]
 #The size of TP
-divisibility = [1, 2, 4, 8, 16, 32, 64]
+divisibility = [1, 2, 8, 16, 64]
 
 all_hidden_size = []
 for div in divisibility:

diff --git a/tests/lora/test_punica_variation.py b/tests/lora/test_punica_variation.py
@@ -20,10 +20,10 @@
 from .utils import (generate_data, generate_data_for_expand_nslices,
                     ref_torch_groupgemm)
 
-HIDDEN_SIZES = [3424, 4096, 4097]
+HIDDEN_SIZES = [4097]
 
 BATCHES = [1, 4, 16, 32]
-NUM_LORA = [1, 4, 8, 16, 32, 64, 128]
+NUM_LORA = [1, 8, 32, 128]
 DTYPES = [torch.float16, torch.bfloat16]
 MAX_RANKS = [1, 4, 8, 16, 32, 64, 128, 256]
 SCALES = [0.5]
@@ -321,22 +321,3 @@ def test_punica_expand_nslices(
 
         slice_offset += hidden_size
     assert_close(our_outputs, ref_outputs)
-
-
-if __name__ == "__main__":
-    from itertools import product
-
-    lst = list(
-        product(
-            BATCHES,
-            NUM_LORA,
-            MAX_RANKS,
-            [1.0],
-            [torch.float16],
-            ["expand"],
-            SEED,
-            CUDA_DEVICES,
-        ))
-    for ele in lst:
-        test_punica_bgmv(*ele)
-        print(f"{ele},pass")