[pull] main from SandAI-org:main

magi_compiler/magi_backend/piecewise_compiler.py

@@ -14,20 +14,65 @@
 
 import contextlib
 import re
+import shutil
+import uuid
 from abc import abstractmethod
 from collections.abc import Callable
 from pathlib import Path
 from typing import Any
 
 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
 
 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(

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

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(<class 'list'>, {'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