[BugFix] Work around graph partition x torch.compile cache issue (#26956)
Signed-off-by: Richard Zou <zou3519@gmail.com>
This commit is contained in:
@@ -337,9 +337,8 @@ def run_model(llama_config, compile_config: CompilationConfig) -> torch.Tensor:
|
||||
def test_toy_llama(
|
||||
backend: str, use_inductor_graph_partition: bool, monkeypatch, tmp_path
|
||||
):
|
||||
# We disable the vLLM compile cache into a new tmp dir for 2 reasons:
|
||||
# We disable the vLLM compile cache into a new tmp dir for 1 reason:
|
||||
# 1. To make sure we can properly track the number of Inductor compilations.
|
||||
# 2. Inductor partitioning does not play nicely with Autograd cache (below)
|
||||
monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1")
|
||||
|
||||
if use_inductor_graph_partition and not is_torch_equal_or_newer("2.9.0.dev"):
|
||||
@@ -369,15 +368,6 @@ def test_toy_llama(
|
||||
cudagraph_capture_sizes=[1, 2],
|
||||
)
|
||||
|
||||
# FIXME(luka/boyuan): the graph from the previous test case
|
||||
# (no inductor partition) gets cached by AotAutograd so then the
|
||||
# compilation with inductor partitioning incorrectly loads an unpartitioned
|
||||
# graph and never partitions. I think this is a bug with custom inductor
|
||||
# partitioning but does not affect vLLM more generally as vLLM uses its own
|
||||
# cache (which takes inductor partitioning into account).
|
||||
if use_inductor_graph_partition:
|
||||
compile_config_no_split.inductor_compile_config["force_disable_caches"] = True
|
||||
|
||||
compile_config_split = deepcopy(compile_config_no_split)
|
||||
compile_config_split.splitting_ops = ["silly::attention"]
|
||||
|
||||
|
||||
Reference in New Issue
Block a user