Repository navigation
[pull] main from SandAI-org:main - #28
Merged
Merged
Conversation
…aphs (#82) * fix: backport PyTorch PR #185635 — force eager backward lowering in standalone_compile Root cause: standalone_compile() calls save_cache_artifacts() immediately after compile_fx() returns. For training graphs, AOTAutograd lazily defers backward lowering to the first .backward() call. This means aot_autograd_artifacts is empty when save() runs, causing: AssertionError: CacheInfo(..., aot_autograd_artifacts=[], ...) Our PyTorch version (2.9.0a0+nv25.10) does not include the upstream fix (PyTorch Issue #152022, fixed in PR #185635 by @jansel, merged 2026-06-04). Fix: Add _force_eager_backward_lowering() context manager that temporarily patches AOTConfig.__post_init__ to set force_non_lazy_backward_lowering=True. This makes AOTAutograd eagerly compile the backward before save_cache_artifacts() runs inside standalone_compile. Observed in production: H200 monolithic 400B TransformerBlock bake produces 108 "Failed to save compiled artifact" warnings. Evidence in athena PR #1070 CI. Replaces the invalid placeholder-injection approach (which would make save() pass mechanically but produce an incomplete cache entry that fails on load). * test: rewrite aot_autograd tests — real training model + cache reuse integration Replace fake mocks (FakeCacheInfo/FakeCompiledGraph) with: - Unit tests: directly exercise _force_eager_backward_lowering() on real AOTConfig, verify patch + restore behaviour - Integration test (subprocess, CUDA): real training model through magi_compile → spy on standalone_compile records force_non_lazy_backward_lowering=True → artifact saved → cache reused on second run (0 recompilations) * style: add copyright header + fix trailing whitespace in helper * test: skip integration test on PT 2.12 — ABI arity mismatch is a known PT version difference * test: bug-reproduce-first pattern + inference scenario - Rename unit tests: test_bug_reproduced / test_bug_fixed - Add test_bug_training_cache_save_fails_without_fix: disables _force_eager_backward_lowering via --disable-fix flag, proves flag is False during standalone_compile (the PyTorch Issue #152022 condition) - Change integration test to inference model (real deployment scenario): forward only under torch.no_grad, cache save + reuse verified - Helper now supports --mode train|infer and --disable-fix * test: use inference model for both bug reproduction and fix verification Both tests now use inference (forward only under torch.no_grad) — matching real deployment. Real-world evidence: every subgraph in 400B inference bake fails with aot_autograd:[] without the fix. Remove --mode train/infer flag; helper is inference-only. Remove TrainingModel. Bug reproduction: --disable-fix replaces _force_eager_backward_lowering with nullcontext, spy records flag=False (the PyTorch Issue #152022 condition). * fix: save SimpleFSDP (DTensor) subgraph artifacts via a nonce AOTAutograd key Dynamo traces SimpleFSDP's param.redistribute(...).to_local(...) into prim_redistribute / prim_to_local closures. AOTAutograd cannot key a graph that calls them, records no aot_autograd artifact, and CompiledArtifact.save() fails with "aot_autograd: []", so all 108 SimpleFSDP subgraphs of the 400B inference bake were recompiled on every run. Backport torch 2.12's nonce-key fallback (bypass_autograd_cache_key) for graphs that cannot be keyed. MagiCompiler finds artifacts by its own cache path, never by this key. Drop _force_eager_backward_lowering and its flag-only tests: the bake runs inference, which has no backward to lower, and the failure reproduces with it enabled. * fix: load saved DTensor subgraph artifacts on PyTorch 2.12 InductorStandaloneAdaptor.load() skips an artifact when the arity of the generated call(self, args) differs from the graph's input count. AOTAutograd flattens tensor-subclass inputs before calling the generated code, and on PyTorch 2.12 each DTensor becomes two values (its local tensor and a non-tensor item), so SimpleFSDP subgraph artifacts were saved but never loaded (6 generated-code inputs vs 4 graph inputs). Skip the check when an input is a traceable wrapper subclass; other graphs keep it.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to subscribe to this conversation on GitHub.
Already have an account?
Sign in.
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
See Commits and Changes for more details.
Created by
pull[bot] (v2.0.0-alpha.4)
Can you help keep this open source service alive? 💖 Please sponsor : )