[torch.compile] support encoder based models (#10613)

Signed-off-by: youkaichao <youkaichao@gmail.com>
2024-11-24 21:24:33 -08:00
parent 7ea3cd7c3e
commit 571841b7fc
2 changed files with 17 additions and 10 deletions
--- a/tests/compile/test_basic_correctness.py
+++ b/tests/compile/test_basic_correctness.py
@@ -62,6 +62,16 @@ test_settings = [
        method="encode",
        fullgraph=True,
    ),
+    # encoder-based embedding model (BERT)
+    TestSetting(
+        model="BAAI/bge-base-en-v1.5",
+        model_args=["--task", "embedding"],
+        pp_size=1,
+        tp_size=1,
+        attn_backend="XFORMERS",
+        method="encode",
+        fullgraph=True,
+    ),
    # vision language model
    TestSetting(
        model="microsoft/Phi-3.5-vision-instruct",