fix(cache): key compiled artifacts on the fx graph and input placements

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),
             ]
         )
 
@@ -736,7 +737,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,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:

tests/feature_tests/cache/cache_reuse_helper/placement_cache_key_helper.py

@@ -0,0 +1,89 @@
+# Copyright (c) 2025 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_placement_cache_key.py

@@ -0,0 +1,65 @@
+# Copyright (c) 2025 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