fix: cache key on placements, AOT JIT fallback, skip FSDP overlap on split graphs

magi_compiler/_api.py

@@ -157,7 +157,24 @@ 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 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__,
+                        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)

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)
 

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",

magi_compiler/utils/hash.py

@@ -14,9 +14,13 @@
 
 import hashlib
 import os
+from collections.abc import Sequence
 from functools import reduce
 from typing import Any, List, Union
 
+import torch.fx as fx
+from torch.distributed.tensor import DTensor
+
 HASH_LENGTH = 10
 
 
@@ -52,6 +56,14 @@ def compute_code_hash_with_content(file_contents: dict[str, str]) -> str:
     return compute_hash("\n".join(hash_content))
 
 
+def compute_graph_hash(graph: fx.GraphModule, example_inputs: Sequence[Any]) -> str:
+    # 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:
     file_contents = {}
     for filepath in files:

tests/api_tests/test_aot_jit_fallback.py

@@ -0,0 +1,65 @@
+# 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
+
+
+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."""
+
+    def _boom(self, *args, **kwargs):
+        raise RuntimeError("Attempted to read undefined local variable")
+
+    monkeypatch.setattr(MagiCompileState, "aot_compile", _boom)
+    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)

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})
 

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
+``<output-dir>/rank<r>.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()

tests/feature_tests/cache/test_graph_hash.py

@@ -0,0 +1,39 @@
+# 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 stay stable across TensorifyScalarRestartAnalysis retrace."""
+
+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_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())
+    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)])