[Misc] Split up pooling tasks (#10820)

Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>
2024-12-11 17:28:00 +08:00
parent 40766ca1b8
commit 8f10d5e393
27 changed files with 527 additions and 168 deletions
--- a/tests/compile/test_basic_correctness.py
+++ b/tests/compile/test_basic_correctness.py
@@ -55,7 +55,7 @@ test_settings = [
    # embedding model
    TestSetting(
        model="BAAI/bge-multilingual-gemma2",
-        model_args=["--task", "embedding"],
+        model_args=["--task", "embed"],
        pp_size=1,
        tp_size=1,
        attn_backend="FLASHINFER",
@@ -65,7 +65,7 @@ test_settings = [
    # encoder-based embedding model (BERT)
    TestSetting(
        model="BAAI/bge-base-en-v1.5",
-        model_args=["--task", "embedding"],
+        model_args=["--task", "embed"],
        pp_size=1,
        tp_size=1,
        attn_backend="XFORMERS",