diff --git a/magi_compiler/magi_backend/piecewise_compiler.py b/magi_compiler/magi_backend/piecewise_compiler.py index 3e4aa31..2ada107 100644 --- a/magi_compiler/magi_backend/piecewise_compiler.py +++ b/magi_compiler/magi_backend/piecewise_compiler.py @@ -14,6 +14,8 @@ import contextlib import re +import shutil +import uuid from abc import abstractmethod from collections.abc import Callable from pathlib import Path @@ -21,6 +23,7 @@ import torch import torch.fx as fx +from torch.utils._python_dispatch import is_traceable_wrapper_subclass from magi_compiler.magi_depyf.timeline import observe_lifecycle from magi_compiler.utils import compilation_counter, compute_hash, magi_logger @@ -28,6 +31,48 @@ from ._cache_data_cls import CacheEntry, CacheHandle +@contextlib.contextmanager +def _nonce_autograd_cache_key_on_keying_failure(): + """Fall back to a nonce AOTAutograd cache key when a graph cannot be keyed. + + Backport of ``torch._functorch.config.bypass_autograd_cache_key`` (what + ``torch._dynamo.aot_compile`` relies on). Our PyTorch version + (2.9.0a0+nv25.10) does not include it. + + Root cause: AOTAutograd refuses to key graphs whose call_function targets + are outside its allowlist, e.g. the ``prim_redistribute`` / ``prim_to_local`` + closures Dynamo emits for SimpleFSDP DTensor params (upstream Issue #167349). + Without a key nothing is recorded, so ``CompiledArtifact.save()`` fails with + ``aot_autograd: []`` and every subgraph is recompiled on the next run. + + A nonce is safe because MagiCompiler finds artifacts by its own cache path, + never by this key. uuid4 instead of upstream's ``random.random()``: ranks and + runs seeded alike share one cache dir and must not generate the same key. + """ + from torch._functorch._aot_autograd import autograd_cache + + orig_autograd_cache_key = autograd_cache.autograd_cache_key + nonce_keys: list[str] = [] + + def _key_or_nonce(*args, **kwargs): + try: + return orig_autograd_cache_key(*args, **kwargs) + except Exception as e: + nonce_key = "a" + uuid.uuid4().hex + nonce_keys.append(nonce_key) + magi_logger.info("AOTAutograd cannot key this graph, using nonce key %s: %s", nonce_key, e) + return nonce_key, [] + + autograd_cache.autograd_cache_key = _key_or_nonce + try: + yield + finally: + autograd_cache.autograd_cache_key = orig_autograd_cache_key + # Entries under a nonce key can never be looked up again; don't let them pile up in the shared cache dir. + for nonce_key in nonce_keys: + shutil.rmtree(autograd_cache.AOTAutogradCache._get_tmp_dir_for_key(nonce_key), ignore_errors=True) + + def _placeholder_names(graph: fx.GraphModule) -> list[str]: return [str(node.target) for node in graph.graph.nodes if node.op == "placeholder"] @@ -321,7 +366,12 @@ def compile( stride_ctx, captured_strides = _intercept_inductor_output_strides() try: - with scope_asserts, functorch_config.patch(autograd_cache_allow_custom_autograd_functions=True), stride_ctx: + with ( + scope_asserts, + _nonce_autograd_cache_key_on_keying_failure(), + functorch_config.patch(autograd_cache_allow_custom_autograd_functions=True), + stride_ctx, + ): compiled_graph = standalone_compile( graph, example_inputs, dynamic_shapes=dynamic_shapes, options={"config_patches": current_config} ) @@ -346,6 +396,8 @@ def compile( # autograd_cache_allow_custom_autograd_functions=True is required above so that # autograd_function_apply (a HigherOrderOperator) does not bypass AOTAutograd cache # key computation, which would leave aot_autograd_artifacts empty and cause save() to fail. + # _nonce_autograd_cache_key_on_keying_failure() does the same for graphs that still cannot + # be keyed, e.g. those with SimpleFSDP DTensor params. assert key is not None restart_analysis_count = self._restart_analysis_counts.get(key, 0) if hasattr(self, "cache_dir") and self.cache_dir is not None: @@ -365,7 +417,12 @@ def load( assert isinstance(cache_handle.key, str) and cache_handle.key is not None assert isinstance(cache_handle.path, str) and cache_handle.path is not None - expected_arity = _read_generated_code_expected_arity(cache_handle.path) + # AOTAutograd flattens tensor-subclass inputs such as DTensor (PyTorch 2.12 passes two values per DTensor), + # so the generated code's arity only matches the graph's when there are none. + if any(is_traceable_wrapper_subclass(x) for x in example_inputs): + expected_arity = None + else: + expected_arity = _read_generated_code_expected_arity(cache_handle.path) actual_arity = len(example_inputs) summarized_inputs = [_summarize_compile_input(x) for x in example_inputs[:8]] magi_logger.info( diff --git a/tests/feature_tests/cache/cache_reuse_helper/dtensor_cache_save_helper.py b/tests/feature_tests/cache/cache_reuse_helper/dtensor_cache_save_helper.py new file mode 100644 index 0000000..af4ddb5 --- /dev/null +++ b/tests/feature_tests/cache/cache_reuse_helper/dtensor_cache_save_helper.py @@ -0,0 +1,114 @@ +# 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_dtensor_cache_save.py. + +Compiles a small SimpleFSDP-sharded MLP (torchtitan ``data_parallel(mode="fully_shard", +ac_mode="full")``, the same call athena's SimpleFSDP builder makes) with ``magi_compile`` +and runs one inference forward under ``torch.no_grad()``. + +Each rank writes ``/rank.json``: + - num_inductor_compiles / num_compiled_artifacts_saved: from compilation_counter + - aot_bypass_reasons: AOTAutograd "Bypassing autograd cache due to ..." messages + - max_abs_diff: compiled vs eager output + +``--disable-fix`` replaces ``_nonce_autograd_cache_key_on_keying_failure`` with a no-op. +""" + +from __future__ import annotations + +import argparse +import contextlib +import json +import logging +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.piecewise_compiler as piecewise_compiler +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))) + + +class _BypassReasons(logging.Handler): + def __init__(self) -> None: + super().__init__(logging.INFO) + self.reasons: list[str] = [] + + def emit(self, record: logging.LogRecord) -> None: + message = record.getMessage() + if message.startswith("Bypassing autograd cache due to"): + self.reasons.append(message) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--output-dir", required=True) + parser.add_argument("--disable-fix", action="store_true") + args = parser.parse_args() + + if args.disable_fix: + piecewise_compiler._nonce_autograd_cache_key_on_keying_failure = contextlib.nullcontext + + bypass_reasons = _BypassReasons() + cache_logger = logging.getLogger("torch._functorch._aot_autograd.autograd_cache") + cache_logger.setLevel(logging.INFO) + cache_logger.addHandler(bypass_reasons) + + 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="fully_shard", 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, + "num_compiled_artifacts_saved": compilation_counter.num_compiled_artifacts_saved, + "aot_bypass_reasons": bypass_reasons.reasons, + "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_dtensor_cache_save.py b/tests/feature_tests/cache/test_dtensor_cache_save.py new file mode 100644 index 0000000..61b8e6d --- /dev/null +++ b/tests/feature_tests/cache/test_dtensor_cache_save.py @@ -0,0 +1,117 @@ +# 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. + +"""SimpleFSDP (DTensor) subgraphs must produce a saveable, reusable compiled artifact. + +Bug (400B inference bake): through torch 2.9 Dynamo traces SimpleFSDP's +``param.redistribute(...).to_local(...)`` into ``prim_redistribute`` / ``prim_to_local`` +closures. AOTAutograd refuses to key a graph that calls them, records no aot_autograd +artifact, and ``CompiledArtifact.save()`` asserts: + + Failed to save compiled artifact for key 'artifact_shape_None_subgraph_0', skipping cache: + CacheInfo(artifacts=defaultdict(, {'inductor': [...], 'aot_autograd': []})) + +so every bake recompiles every subgraph. + +Fix: ``_nonce_autograd_cache_key_on_keying_failure`` in piecewise_compiler.py. + +Driven through torchrun like the fsdp e2e tests: world=1, and world=2 where the weights are +real Shard(0) DTensors all-gathered inside the graph. +""" + +from __future__ import annotations + +import json +import os +import re +import shutil +import subprocess +from pathlib import Path + +import pytest +import torch + +_HELPER = Path(__file__).parent / "cache_reuse_helper" / "dtensor_cache_save_helper.py" + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +requires_torchrun = pytest.mark.skipif(shutil.which("torchrun") is None, reason="requires torchrun") +world_sizes = pytest.mark.parametrize( + "nproc", [1, pytest.param(2, marks=pytest.mark.skipif(torch.cuda.device_count() < 2, reason="requires >=2 GPUs"))] +) + + +def _run(cache_root: Path, output_dir: Path, nproc: int, port: int, *flags: str) -> tuple[str, list[dict]]: + output_dir.mkdir(parents=True) + env = os.environ.copy() + env["MAGI_LOGGING_LEVEL"] = "info" + env["MAGI_COMPILE_CACHE_ROOT_DIR"] = str(cache_root) + p = subprocess.run( + [ + "torchrun", + f"--nproc_per_node={nproc}", + f"--master_port={port}", + str(_HELPER), + "--output-dir", + str(output_dir), + *flags, + ], + env=env, + capture_output=True, + text=True, + timeout=900, + ) + out = p.stdout + p.stderr + assert p.returncode == 0, f"helper failed:\n{out[-4000:]}" + return out, [json.loads((output_dir / f"rank{rank}.json").read_text()) for rank in range(nproc)] + + +@requires_cuda +@requires_torchrun +@world_sizes +def test_bug_reproduced_without_fix(tmp_path: Path, nproc: int): + """Without the nonce-key fix: the exact bake error, and no artifact is saved.""" + out, payloads = _run(tmp_path / "cache", tmp_path / "out", nproc, 29660 + nproc, "--disable-fix") + reasons = [reason for payload in payloads for reason in payload["aot_bypass_reasons"]] + if not reasons: + pytest.skip("this PyTorch keys SimpleFSDP graphs (DTensor methods traced as call_method)") + + assert any("prim_redistribute" in r or "prim_to_local" in r for r in reasons), reasons + assert "Failed to save compiled artifact" in out, out[-4000:] + assert "'aot_autograd': []" in out, out[-4000:] + assert all(payload["num_compiled_artifacts_saved"] == 0 for payload in payloads), payloads + + +@requires_cuda +@requires_torchrun +@world_sizes +def test_bug_fixed_artifact_saved_and_reused(tmp_path: Path, nproc: int): + """With the fix: every compiled subgraph is saved on the cold run and loaded on the warm run.""" + cache_root = tmp_path / "cache" + + cold_out, cold_payloads = _run(cache_root, tmp_path / "cold", nproc, 29670 + nproc) + assert "Failed to save compiled artifact" not in cold_out, cold_out[-4000:] + for payload in cold_payloads: + assert payload["num_inductor_compiles"] > 0, payload + assert payload["num_compiled_artifacts_saved"] == payload["num_inductor_compiles"], payload + assert payload["max_abs_diff"] < 1e-2, payload + + # A nonce key can never be looked up again, so its entry must not stay in the shared cache dir. + for nonce_key in re.findall(r"using nonce key (a[0-9a-f]{32})", cold_out): + assert not (cache_root / "inductor_cache" / "aotautograd" / nonce_key).exists(), nonce_key + + warm_out, warm_payloads = _run(cache_root, tmp_path / "warm", nproc, 29680 + nproc) + assert "Failed to save compiled artifact" not in warm_out, warm_out[-4000:] + for payload in warm_payloads: + assert payload["num_inductor_compiles"] == 0, f"artifact not reused: {payload}\n{warm_out[-4000:]}" + assert payload["max_abs_diff"] < 1e-2, payload