From 3c4e3e09af82398dfa8ebe5b2c278c67ad7ecf85 Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Wed, 14 Oct 2026 15:46:06 +0000 Subject: [PATCH 1/4] fix: cache key on graph+placements, AOT JIT fallback, skip FSDP overlap on split graphs Unmerged remnants of #88 / #86 / #84 / #79, rebased onto current main: - Cache key includes graph.code and DTensor placements. Changing SimpleFSDP fully_shard -> replicate used to reuse the old artifact (1-rank AllGathers served as DtoD copies). - AOT compile failures fall back to JIT so MAGI_COMPILE_AOT=1 can stay process-wide (VAE AOT, SimpleFSDP DiT JIT). - enable_fsdp no longer runs fullgraph overlap when disable_graph_split is False (VAE in a monolithic FSDP process). #79's restart-analysis half already landed as #75. - Timing bench dim 128->1024 so launch jitter does not flake max/min < 1.2 on B300/H100. --- magi_compiler/_api.py | 18 +++- magi_compiler/magi_backend/magi_backend.py | 12 +-- magi_compiler/utils/__init__.py | 3 +- magi_compiler/utils/hash.py | 18 ++++ tests/api_tests/test_aot_jit_fallback.py | 46 ++++++++++ tests/api_tests/test_magi_compile.py | 14 +-- .../placement_cache_key_helper.py | 89 +++++++++++++++++++ tests/feature_tests/cache/test_graph_hash.py | 35 ++++++++ .../cache/test_placement_cache_key.py | 75 ++++++++++++++++ .../fsdp/test_fsdp_skip_split.py | 50 +++++++++++ 10 files changed, 346 insertions(+), 14 deletions(-) create mode 100644 tests/api_tests/test_aot_jit_fallback.py create mode 100644 tests/feature_tests/cache/cache_reuse_helper/placement_cache_key_helper.py create mode 100644 tests/feature_tests/cache/test_graph_hash.py create mode 100644 tests/feature_tests/cache/test_placement_cache_key.py create mode 100644 tests/feature_tests/fsdp/test_fsdp_skip_split.py diff --git a/magi_compiler/_api.py b/magi_compiler/_api.py index 6f91ddf..f54cc0c 100644 --- a/magi_compiler/_api.py +++ b/magi_compiler/_api.py @@ -157,7 +157,23 @@ def _run_orchestration(state: MagiCompileState, args, kwargs): state._ensure_compiled() if state.compile_config.aot: - state.aot_compile(*args, **kwargs) + try: + state.aot_compile(*args, **kwargs) + except Exception as e: + # AOT still cannot compile SimpleFSDP/DTensor graphs (undefined local in + # aot_compile_fullgraph). Keep MAGI_COMPILE_AOT=1 process-wide for VAE. + magi_logger.warning( + "AOT compile failed (%s: %s), falling back to JIT path for %s", + type(e).__name__, + e, + state.original_code_for_hook, + ) + state.compile_config = state.compile_config.model_copy(update={"aot": False}) + state.compiled_entry = None + torch._dynamo.reset() + state._ensure_compiled() + with state._jit_capture_compiled_bytecode(): + return state.compiled_entry(*args, **kwargs) else: with state._jit_capture_compiled_bytecode(): return state.compiled_entry(*args, **kwargs) diff --git a/magi_compiler/magi_backend/magi_backend.py b/magi_compiler/magi_backend/magi_backend.py index 973465c..48843b5 100644 --- a/magi_compiler/magi_backend/magi_backend.py +++ b/magi_compiler/magi_backend/magi_backend.py @@ -34,7 +34,7 @@ from magi_compiler.magi_depyf.timeline import observe_lifecycle, observe_lifecycle_context from magi_compiler.offload.offload_warpper import OffloadWrapper from magi_compiler.passes import CustomJointGraphPartitionFn, FullGraphPassManager, PostGradPassManager, pass_context -from magi_compiler.utils import compilation_counter, compute_code_hash, compute_hash, magi_logger +from magi_compiler.utils import compilation_counter, compute_code_hash, compute_graph_hash, compute_hash, magi_logger from magi_compiler.utils.visualize import save_fx_graph_visualization from ._cache_data_cls import CacheEntry, CacheHandle @@ -587,13 +587,14 @@ def _configure_custom_passes(self): post_grad_pass_manager.snapshot_original_inductor_configs(self.inductor_compile_config) - def _init_cache(self) -> str: + def _init_cache(self, graph: fx.GraphModule, example_inputs) -> str: hash_key = compute_hash( [ self.compile_config.hash, inductor_compile_config_hash(self.inductor_compile_config), self.compiler_manager.hash, compute_code_hash(self.traced_files), + compute_graph_hash(graph, example_inputs), ] ) @@ -677,8 +678,9 @@ def _split_graph(self, graph: fx.GraphModule, example_inputs) -> tuple[fx.GraphM magi_logger.info(f"Setting up FX-level graph split with ops: {fx_split_ops=}") magi_logger.info(f"Resolved splitting ops for FX-level graph split: {resolved_ops=}") - # Step 1.4: whole-graph FSDP overlap. - if self.compile_config.fsdp_config.enable_fsdp: + # Step 1.4: whole-graph FSDP overlap. enable_fsdp is a process-wide flag; VAE and other + # split-graph models share it with FSDP DiT, so skip unless this graph is fullgraph. + if self.compile_config.fsdp_config.enable_fsdp and self.compile_config.disable_graph_split: self._apply_fsdp_fullgraph_overlap(graph, example_inputs) # Step 2: split graph by ops, we split graph based on resolved_ops, which becomes the partitioned single graph. @@ -736,7 +738,7 @@ def __call__(self, graph: fx.GraphModule, example_inputs) -> MagiSerializableFun magi_logger.info("Dynamo traced files (for compilation cache):\n%s", "\n".join(self.traced_files)) compilation_counter.num_graphs_seen += 1 - self._init_cache() + self._init_cache(graph=graph, example_inputs=example_inputs) self.full_graph_pass_manager(graph) diff --git a/magi_compiler/utils/__init__.py b/magi_compiler/utils/__init__.py index c93a742..b1df2b0 100644 --- a/magi_compiler/utils/__init__.py +++ b/magi_compiler/utils/__init__.py @@ -17,7 +17,7 @@ from .compile_counter import compilation_counter from .device import device_capability, device_capability_major from .envs import set_env_var -from .hash import compute_code_hash, compute_code_hash_with_content, compute_hash +from .hash import compute_code_hash, compute_code_hash_with_content, compute_graph_hash, compute_hash from .logger import logger, magi_logger from .nvtx import add_nvtx_event, instrument_nvtx from .ordered_set import OrderedSet @@ -28,6 +28,7 @@ "set_env_var", "compute_code_hash", "compute_code_hash_with_content", + "compute_graph_hash", "compute_hash", "logger", "magi_logger", diff --git a/magi_compiler/utils/hash.py b/magi_compiler/utils/hash.py index 035c047..6e6b2eb 100644 --- a/magi_compiler/utils/hash.py +++ b/magi_compiler/utils/hash.py @@ -14,9 +14,14 @@ import hashlib import os +from collections.abc import Sequence from functools import reduce from typing import Any, List, Union +import torch +import torch.fx as fx +from torch.distributed.tensor import DTensor + HASH_LENGTH = 10 @@ -52,6 +57,19 @@ def compute_code_hash_with_content(file_contents: dict[str, str]) -> str: return compute_hash("\n".join(hash_content)) +def _input_signature(x: Any) -> str: + if isinstance(x, DTensor): + return f"DTensor({x.dtype}, {x.placements}, {tuple(x.device_mesh.shape)})" + if isinstance(x, torch.Tensor): + return f"Tensor({x.dtype})" + return type(x).__name__ + + +def compute_graph_hash(graph: fx.GraphModule, example_inputs: Sequence[Any]) -> str: + # DTensor placements live in the inputs, not in the traced source files (e.g. SimpleFSDP fully_shard vs replicate). + return compute_hash([graph.code, [_input_signature(x) for x in example_inputs]]) + + def compute_code_hash(files: set[str]) -> str: file_contents = {} for filepath in files: diff --git a/tests/api_tests/test_aot_jit_fallback.py b/tests/api_tests/test_aot_jit_fallback.py new file mode 100644 index 0000000..af56373 --- /dev/null +++ b/tests/api_tests/test_aot_jit_fallback.py @@ -0,0 +1,46 @@ +# Copyright (c) 2026 SandAI. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""AOT compile failure must fall back to JIT instead of aborting the first request.""" + +from __future__ import annotations + +import pytest +import torch +import torch.nn as nn + +from magi_compiler import magi_compile +from magi_compiler.magi_backend.magi_compiler_base import MagiCompileState + + +class _Tiny(nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x + 1 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_aot_failure_falls_back_to_jit(monkeypatch: pytest.MonkeyPatch) -> None: + """Without the fallback, aot_compile raising aborts warmup for every model in the process.""" + + def _boom(self, *args, **kwargs): + raise RuntimeError("Attempted to read undefined local variable") + + monkeypatch.setattr(MagiCompileState, "aot_compile", _boom) + model = magi_compile( + _Tiny().cuda().eval(), dynamic_arg_dims={"x": 0}, config_patch=lambda c: c.model_copy(update={"aot": True}) + ) + x = torch.randn(4, 4, device="cuda") + with torch.no_grad(): + out = model(x) + torch.testing.assert_close(out, x + 1) diff --git a/tests/api_tests/test_magi_compile.py b/tests/api_tests/test_magi_compile.py index 93ef398..7b0a8a3 100644 --- a/tests/api_tests/test_magi_compile.py +++ b/tests/api_tests/test_magi_compile.py @@ -552,7 +552,7 @@ def test_simple_model_timing_class_function_instance_method(self): class SimpleModel(nn.Module): def __init__(self): super().__init__() - self.dim = 128 + self.dim = 1024 self.layers = nn.ModuleList([nn.Linear(self.dim, self.dim) for _ in range(4)]) self.norms = nn.ModuleList([nn.LayerNorm(self.dim) for _ in range(4)]) @@ -567,8 +567,8 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return x device = torch.device("cuda:0") - seq_len = 256 - test_input = torch.randn(seq_len, 128, device=device) + seq_len = 512 + test_input = torch.randn(seq_len, 1024, device=device) native = SimpleModel().to(device).eval() @@ -632,16 +632,16 @@ class CompiledNonModulePerf(NonModulePerf): def forward(self, x: torch.Tensor) -> torch.Tensor: return super().forward(x) - non_module_native = NonModulePerf(128) + non_module_native = NonModulePerf(1024) - non_module_cls = CompiledNonModulePerf(128) + non_module_cls = CompiledNonModulePerf(1024) non_module_cls.copy_from(non_module_native) - non_module_inst_obj = NonModulePerf(128) + non_module_inst_obj = NonModulePerf(1024) non_module_inst_obj.copy_from(non_module_native) non_module_inst = magi_compile(non_module_inst_obj, dynamic_arg_dims={"x": 0}) - non_module_mtd_obj = NonModulePerf(128) + non_module_mtd_obj = NonModulePerf(1024) non_module_mtd_obj.copy_from(non_module_native) non_module_mtd_obj.step = magi_compile(non_module_mtd_obj.step, dynamic_arg_dims={"x": 0}) diff --git a/tests/feature_tests/cache/cache_reuse_helper/placement_cache_key_helper.py b/tests/feature_tests/cache/cache_reuse_helper/placement_cache_key_helper.py new file mode 100644 index 0000000..20f5ec1 --- /dev/null +++ b/tests/feature_tests/cache/cache_reuse_helper/placement_cache_key_helper.py @@ -0,0 +1,89 @@ +# Copyright (c) 2026 SandAI. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""torchrun helper for test_placement_cache_key.py. + +Compiles a SimpleFSDP MLP wrapped with ``data_parallel(mode=--mode)`` and writes +``/rank.json`` with ``num_inductor_compiles`` and ``max_abs_diff``. + +``--disable-fix`` drops the graph signature from the magi cache key. +""" + +from __future__ import annotations + +import argparse +import json +import os +from pathlib import Path + +import torch +import torch.distributed as dist +import torch.nn as nn +from torch.distributed.device_mesh import init_device_mesh + +import magi_compiler.magi_backend.magi_backend as magi_backend +from magi_compiler import magi_compile +from magi_compiler.utils import compilation_counter + +HIDDEN = 64 + + +class MLP(nn.Module): + def __init__(self) -> None: + super().__init__() + self.fc1 = nn.Linear(HIDDEN, 4 * HIDDEN, bias=False) + self.fc2 = nn.Linear(4 * HIDDEN, HIDDEN, bias=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.fc2(torch.nn.functional.gelu(self.fc1(x))) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--output-dir", required=True) + parser.add_argument("--mode", choices=["fully_shard", "replicate"], required=True) + parser.add_argument("--disable-fix", action="store_true") + args = parser.parse_args() + + if args.disable_fix: + magi_backend.compute_graph_hash = lambda graph, example_inputs: "" + + dist.init_process_group("cpu:gloo,cuda:nccl") + rank = dist.get_rank() + torch.cuda.set_device(int(os.environ.get("LOCAL_RANK", rank))) + mesh = init_device_mesh("cuda", (dist.get_world_size(),)) + + from torchtitan.experiments.simple_fsdp.simple_fsdp import data_parallel + + torch.manual_seed(0) + reference = MLP().cuda().to(torch.bfloat16) + torch.manual_seed(0) + model = data_parallel(MLP().cuda().to(torch.bfloat16), mesh, mode=args.mode, ac_mode="full") + compiled = magi_compile(model, dynamic_arg_dims={"x": 0}) + + x = torch.randn(8, HIDDEN, device="cuda", dtype=torch.bfloat16) + with torch.no_grad(): + out = compiled(x) + expected = reference(x) + + payload = { + "num_inductor_compiles": compilation_counter.num_inductor_compiles, + "max_abs_diff": (out.float() - expected.float()).abs().max().item(), + } + Path(args.output_dir, f"rank{rank}.json").write_text(json.dumps(payload)) + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/tests/feature_tests/cache/test_graph_hash.py b/tests/feature_tests/cache/test_graph_hash.py new file mode 100644 index 0000000..bd59806 --- /dev/null +++ b/tests/feature_tests/cache/test_graph_hash.py @@ -0,0 +1,35 @@ +# Copyright (c) 2026 SandAI. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""compute_graph_hash must change when input dtypes change; placements are covered by test_placement_cache_key.""" + +from __future__ import annotations + +import torch +import torch.nn as nn + +from magi_compiler.utils import compute_graph_hash + + +class _Add(nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x + 1 + + +def test_compute_graph_hash_changes_with_input_dtype() -> None: + graph = torch.fx.symbolic_trace(_Add()) + fp32 = compute_graph_hash(graph, [torch.zeros(2, dtype=torch.float32)]) + bf16 = compute_graph_hash(graph, [torch.zeros(2, dtype=torch.bfloat16)]) + assert fp32 != bf16 + assert fp32 == compute_graph_hash(graph, [torch.ones(8, dtype=torch.float32)]) diff --git a/tests/feature_tests/cache/test_placement_cache_key.py b/tests/feature_tests/cache/test_placement_cache_key.py new file mode 100644 index 0000000..a727656 --- /dev/null +++ b/tests/feature_tests/cache/test_placement_cache_key.py @@ -0,0 +1,75 @@ +# Copyright (c) 2026 SandAI. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Changing parameter placements alone must not reuse a compiled artifact from the old placements.""" + +from __future__ import annotations + +import json +import os +import shutil +import subprocess +from pathlib import Path + +import pytest +import torch + +_HELPER = Path(__file__).parent / "cache_reuse_helper" / "placement_cache_key_helper.py" + +pytestmark = [ + pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA"), + pytest.mark.skipif(shutil.which("torchrun") is None, reason="requires torchrun"), +] + + +def _compiles(cache_root: Path, output_dir: Path, mode: str, port: int, *flags: str) -> int: + output_dir.mkdir(parents=True) + env = os.environ.copy() + env["MAGI_COMPILE_CACHE_ROOT_DIR"] = str(cache_root) + p = subprocess.run( + [ + "torchrun", + "--nproc_per_node=1", + f"--master_port={port}", + str(_HELPER), + "--output-dir", + str(output_dir), + "--mode", + mode, + *flags, + ], + env=env, + capture_output=True, + text=True, + timeout=900, + ) + assert p.returncode == 0, f"helper failed:\n{(p.stdout + p.stderr)[-4000:]}" + payload = json.loads((output_dir / "rank0.json").read_text()) + assert payload["max_abs_diff"] < 1e-2, payload + return payload["num_inductor_compiles"] + + +def test_bug_reproduced_without_fix(tmp_path: Path): + """Without the graph signature in the key, the replicate run loads the fully_shard artifact.""" + cache_root = tmp_path / "cache" + assert _compiles(cache_root, tmp_path / "shard", "fully_shard", 29690, "--disable-fix") > 0 + assert _compiles(cache_root, tmp_path / "replicate", "replicate", 29691, "--disable-fix") == 0 + + +def test_bug_fixed_placement_change_recompiles(tmp_path: Path): + """A placement change recompiles; rerunning the same placements still reuses the cache.""" + cache_root = tmp_path / "cache" + assert _compiles(cache_root, tmp_path / "shard", "fully_shard", 29692) > 0 + assert _compiles(cache_root, tmp_path / "replicate", "replicate", 29693) > 0 + assert _compiles(cache_root, tmp_path / "replicate_warm", "replicate", 29694) == 0 diff --git a/tests/feature_tests/fsdp/test_fsdp_skip_split.py b/tests/feature_tests/fsdp/test_fsdp_skip_split.py new file mode 100644 index 0000000..af8140c --- /dev/null +++ b/tests/feature_tests/fsdp/test_fsdp_skip_split.py @@ -0,0 +1,50 @@ +# Copyright (c) 2026 SandAI. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""enable_fsdp is process-wide; split-graph models (VAE) must not enter fullgraph overlap.""" + +from __future__ import annotations + +import torch +import torch.nn as nn + +from magi_compiler.config import CompileConfig, FSDPConfig +from magi_compiler.magi_backend.magi_backend import MagiBackend +from magi_compiler.utils.ordered_set import OrderedSet + + +class _Add(nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x + 1 + + +def test_split_graph_skips_fsdp_overlap_when_graph_split_enabled(monkeypatch) -> None: + called: list[bool] = [] + cfg = CompileConfig(fsdp_config=FSDPConfig(enable_fsdp=True)) + assert cfg.disable_graph_split is False + backend = MagiBackend(cfg, 0, "vae", OrderedSet(), {}) + monkeypatch.setattr(backend, "_apply_fsdp_fullgraph_overlap", lambda *a, **k: called.append(True)) + graph = torch.fx.symbolic_trace(_Add()) + backend._split_graph(graph, (torch.zeros(2),)) + assert called == [] + + +def test_fullgraph_still_applies_fsdp_overlap(monkeypatch) -> None: + called: list[bool] = [] + cfg = CompileConfig(disable_graph_split=True, fsdp_config=FSDPConfig(enable_fsdp=True)) + backend = MagiBackend(cfg, 0, "base", OrderedSet(), {}) + monkeypatch.setattr(backend, "_apply_fsdp_fullgraph_overlap", lambda *a, **k: called.append(True)) + graph = torch.fx.symbolic_trace(_Add()) + backend._split_graph(graph, (torch.zeros(2),)) + assert called == [True] From daf45f1bc4a674d055625c78071b66afbdbdc312 Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Wed, 14 Oct 2026 16:00:03 +0000 Subject: [PATCH 2/4] fix(cache): hash input placements only, not graph.code Restart-analysis retraces with different FX node names, so including graph.code in the magi cache key made the second pass miss the first-pass directory (test_assert_cache_hit_succeeds_with_restart_analysis / test_restart_analysis_cache on pt212). DTensor fully_shard vs replicate is still distinguished by the input placement signature. --- magi_compiler/utils/hash.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/magi_compiler/utils/hash.py b/magi_compiler/utils/hash.py index 6e6b2eb..5f31ca9 100644 --- a/magi_compiler/utils/hash.py +++ b/magi_compiler/utils/hash.py @@ -66,8 +66,9 @@ def _input_signature(x: Any) -> str: def compute_graph_hash(graph: fx.GraphModule, example_inputs: Sequence[Any]) -> str: - # DTensor placements live in the inputs, not in the traced source files (e.g. SimpleFSDP fully_shard vs replicate). - return compute_hash([graph.code, [_input_signature(x) for x in example_inputs]]) + # Placements live on the inputs, not in traced source files. graph.code is omitted: + # restart-analysis retraces with different node names and would miss the first-pass cache. + return compute_hash([_input_signature(x) for x in example_inputs]) def compute_code_hash(files: set[str]) -> str: From 9ac1a9de8e97b95dcd4f94d70f67e5893b2fd566 Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Wed, 14 Oct 2026 16:16:36 +0000 Subject: [PATCH 3/4] fix(cache): hash only DTensor placements so restart-analysis still hits --- magi_compiler/utils/hash.py | 17 +++++------------ tests/feature_tests/cache/test_graph_hash.py | 16 ++++++++++------ 2 files changed, 15 insertions(+), 18 deletions(-) diff --git a/magi_compiler/utils/hash.py b/magi_compiler/utils/hash.py index 5f31ca9..0f15f94 100644 --- a/magi_compiler/utils/hash.py +++ b/magi_compiler/utils/hash.py @@ -18,7 +18,6 @@ from functools import reduce from typing import Any, List, Union -import torch import torch.fx as fx from torch.distributed.tensor import DTensor @@ -57,18 +56,12 @@ def compute_code_hash_with_content(file_contents: dict[str, str]) -> str: return compute_hash("\n".join(hash_content)) -def _input_signature(x: Any) -> str: - if isinstance(x, DTensor): - return f"DTensor({x.dtype}, {x.placements}, {tuple(x.device_mesh.shape)})" - if isinstance(x, torch.Tensor): - return f"Tensor({x.dtype})" - return type(x).__name__ - - def compute_graph_hash(graph: fx.GraphModule, example_inputs: Sequence[Any]) -> str: - # Placements live on the inputs, not in traced source files. graph.code is omitted: - # restart-analysis retraces with different node names and would miss the first-pass cache. - return compute_hash([_input_signature(x) for x in example_inputs]) + # graph.code and Tensor/scalar types both change across TensorifyScalarRestartAnalysis. + # DTensor placements do not appear in traced source files. + return compute_hash( + [f"DTensor({x.dtype}, {x.placements}, {tuple(x.device_mesh.shape)})" for x in example_inputs if isinstance(x, DTensor)] + ) def compute_code_hash(files: set[str]) -> str: diff --git a/tests/feature_tests/cache/test_graph_hash.py b/tests/feature_tests/cache/test_graph_hash.py index bd59806..ff2d485 100644 --- a/tests/feature_tests/cache/test_graph_hash.py +++ b/tests/feature_tests/cache/test_graph_hash.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""compute_graph_hash must change when input dtypes change; placements are covered by test_placement_cache_key.""" +"""compute_graph_hash must stay stable across TensorifyScalarRestartAnalysis retrace.""" from __future__ import annotations @@ -27,9 +27,13 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return x + 1 -def test_compute_graph_hash_changes_with_input_dtype() -> None: +def test_compute_graph_hash_stable_across_tensorify_scalar() -> None: + """Restart-analysis retrace can replace a Tensor with a Python scalar. + + Including those types in the cache key makes the second pass miss first-pass artifacts. + """ graph = torch.fx.symbolic_trace(_Add()) - fp32 = compute_graph_hash(graph, [torch.zeros(2, dtype=torch.float32)]) - bf16 = compute_graph_hash(graph, [torch.zeros(2, dtype=torch.bfloat16)]) - assert fp32 != bf16 - assert fp32 == compute_graph_hash(graph, [torch.ones(8, dtype=torch.float32)]) + tensor_inputs = compute_graph_hash(graph, [torch.zeros(2, dtype=torch.float32)]) + scalar_inputs = compute_graph_hash(graph, [1]) + assert tensor_inputs == scalar_inputs + assert tensor_inputs == compute_graph_hash(graph, [torch.zeros(2, dtype=torch.bfloat16)]) From 3fc8a63f5b650f5557d0c8f97258060c454584dd Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Thu, 15 Oct 2026 04:00:04 +0000 Subject: [PATCH 4/4] fix(aot): do not swallow AssertionError/OOM when AOT falls back to JIT --- magi_compiler/_api.py | 7 ++++--- tests/api_tests/test_aot_jit_fallback.py | 25 +++++++++++++++++++++--- 2 files changed, 26 insertions(+), 6 deletions(-) diff --git a/magi_compiler/_api.py b/magi_compiler/_api.py index f54cc0c..5ec4323 100644 --- a/magi_compiler/_api.py +++ b/magi_compiler/_api.py @@ -159,9 +159,10 @@ def _run_orchestration(state: MagiCompileState, args, kwargs): if state.compile_config.aot: try: state.aot_compile(*args, **kwargs) - except Exception as e: - # AOT still cannot compile SimpleFSDP/DTensor graphs (undefined local in - # aot_compile_fullgraph). Keep MAGI_COMPILE_AOT=1 process-wide for VAE. + except RuntimeError as e: + # SimpleFSDP/DTensor aot_compile_fullgraph is unsupported; VAE still uses MAGI_COMPILE_AOT=1. + if isinstance(e, torch.cuda.OutOfMemoryError): + raise magi_logger.warning( "AOT compile failed (%s: %s), falling back to JIT path for %s", type(e).__name__, diff --git a/tests/api_tests/test_aot_jit_fallback.py b/tests/api_tests/test_aot_jit_fallback.py index af56373..f9e893f 100644 --- a/tests/api_tests/test_aot_jit_fallback.py +++ b/tests/api_tests/test_aot_jit_fallback.py @@ -29,6 +29,12 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return x + 1 +def _compile_tiny(): + return magi_compile( + _Tiny().cuda().eval(), dynamic_arg_dims={"x": 0}, config_patch=lambda c: c.model_copy(update={"aot": True}) + ) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") def test_aot_failure_falls_back_to_jit(monkeypatch: pytest.MonkeyPatch) -> None: """Without the fallback, aot_compile raising aborts warmup for every model in the process.""" @@ -37,10 +43,23 @@ def _boom(self, *args, **kwargs): raise RuntimeError("Attempted to read undefined local variable") monkeypatch.setattr(MagiCompileState, "aot_compile", _boom) - model = magi_compile( - _Tiny().cuda().eval(), dynamic_arg_dims={"x": 0}, config_patch=lambda c: c.model_copy(update={"aot": True}) - ) + model = _compile_tiny() x = torch.randn(4, 4, device="cuda") with torch.no_grad(): out = model(x) torch.testing.assert_close(out, x + 1) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_aot_assertion_error_is_not_swallowed(monkeypatch: pytest.MonkeyPatch) -> None: + """Programmer invariants in aot_compile must not be converted into a silent JIT fallback.""" + + def _boom(self, *args, **kwargs): + raise AssertionError("programmer invariant") + + monkeypatch.setattr(MagiCompileState, "aot_compile", _boom) + model = _compile_tiny() + x = torch.randn(4, 4, device="cuda") + with pytest.raises(AssertionError, match="programmer invariant"): + with torch.no_grad(): + model(x)