Unverified Commit 64384bbc authored by youkaichao's avatar youkaichao Committed by GitHub
Browse files

[torch.compile] upgrade tests (#9858)


Signed-off-by: default avataryoukaichao <youkaichao@gmail.com>
parent 00d91c8a
...@@ -30,18 +30,20 @@ def test_compile_correctness(model, model_args, pp_size, tp_size, attn_backend, ...@@ -30,18 +30,20 @@ def test_compile_correctness(model, model_args, pp_size, tp_size, attn_backend,
pytest.skip("Not correct CUDA devices for the test.") pytest.skip("Not correct CUDA devices for the test.")
import os import os
os.environ["VLLM_ATTENTION_BACKEND"] = attn_backend os.environ["VLLM_ATTENTION_BACKEND"] = attn_backend
if not fullgraph: all_args = [["--enforce-eager"] + model_args + ["-pp", str(pp_size)] +
os.environ["VLLM_TEST_DYNAMO_FULLGRAPH_CAPTURE"] = "0" ["-tp", str(tp_size)]] * 3
all_args = [["--enforce-eager"] + model_args + ["--max_model_len", "1024"]
+ ["-pp", str(pp_size)] + ["-tp", str(tp_size)]] * 3
# don't test VLLM_TORCH_COMPILE_LEVEL == 3 case # don't test VLLM_TORCH_COMPILE_LEVEL == 3 case
# inductor will change the output, so we cannot compare them. # inductor will change the output, so we cannot compare them.
all_envs: List[Optional[Dict[str, str]]] = [{ all_envs: List[Optional[Dict[str, str]]] = []
"VLLM_TORCH_COMPILE_LEVEL": for level in [
str(level)
} for level in [
CompilationLevel.NO_COMPILATION, CompilationLevel.NO_COMPILATION,
CompilationLevel.DYNAMO_AS_IS, CompilationLevel.DYNAMO_AS_IS,
CompilationLevel.DYNAMO_ONCE, CompilationLevel.DYNAMO_ONCE,
]] ]:
all_envs.append({"VLLM_TORCH_COMPILE_LEVEL": str(level)})
if level != CompilationLevel.DYNAMO_ONCE and not fullgraph:
# "DYNAMO_ONCE" will always use fullgraph
all_envs[-1][
"VLLM_TEST_DYNAMO_FULLGRAPH_CAPTURE"] = "0" # type: ignore
compare_all_settings(model, all_args, all_envs, method=method) compare_all_settings(model, all_args, all_envs, method=method)
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment