Skip to content

Commit 7be141b

Browse files
authored
[CI] Enable encoder model compilation test (#24442)
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
1 parent 8d7f39b commit 7be141b

File tree

1 file changed

+17
-14
lines changed

1 file changed

+17
-14
lines changed

tests/compile/test_basic_correctness.py

Lines changed: 17 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -62,26 +62,28 @@ class TestSetting:
6262
TestSetting(
6363
model="BAAI/bge-multilingual-gemma2",
6464
model_args=[
65-
"--runner", "pooling", "--dtype", "bfloat16",
66-
"--max-model-len", "2048"
65+
"--runner",
66+
"pooling",
67+
"--dtype",
68+
"bfloat16",
69+
"--max-model-len",
70+
"2048",
6771
],
6872
pp_size=1,
6973
tp_size=1,
7074
attn_backend="FLASH_ATTN",
7175
method="encode",
7276
fullgraph=True,
7377
),
74-
# TODO: bert models are not supported in V1 yet
75-
# # encoder-based embedding model (BERT)
76-
# TestSetting(
77-
# model="BAAI/bge-base-en-v1.5",
78-
# model_args=["--runner", "pooling"],
79-
# pp_size=1,
80-
# tp_size=1,
81-
# attn_backend="XFORMERS",
82-
# method="encode",
83-
# fullgraph=True,
84-
# ),
78+
TestSetting(
79+
model="BAAI/bge-base-en-v1.5",
80+
model_args=["--runner", "pooling"],
81+
pp_size=1,
82+
tp_size=1,
83+
attn_backend="FLASH_ATTN",
84+
method="encode",
85+
fullgraph=True,
86+
),
8587
# vision language model
8688
TestSetting(
8789
model="microsoft/Phi-3.5-vision-instruct",
@@ -92,7 +94,8 @@ class TestSetting:
9294
method="generate_with_image",
9395
fullgraph=False,
9496
),
95-
])
97+
],
98+
)
9699
def test_compile_correctness(
97100
monkeypatch: pytest.MonkeyPatch,
98101
test_setting: TestSetting,

0 commit comments

Comments
 (0)