diff --git a/magi_compiler/magi_backend/magi_backend.py b/magi_compiler/magi_backend/magi_backend.py index 973465c..1bdd1a7 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), ] ) @@ -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) 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/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..384d772 --- /dev/null +++ b/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 +``/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_placement_cache_key.py b/tests/feature_tests/cache/test_placement_cache_key.py new file mode 100644 index 0000000..3fe8eac --- /dev/null +++ b/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