fix: cache key on placements, AOT JIT fallback, skip FSDP overlap on split graphs

[cennn]

Summary

Consolidates the still-useful parts of four open PRs onto current main, then those PRs can close:

• fix(cache): key compiled artifacts on the fx graph and input placements #88 cache key: compute_graph_hash(example_inputs) is part of the magi cache directory name. Only DTensor (dtype, placements, mesh.shape) is hashed — not graph.code and not plain Tensor/scalar types. Changing SimpleFSDP fully_shard → replicate used to reuse the old artifact (source files unchanged); on 400B 8xH200 that kept issuing 1-rank AllGathers (NCCL as DtoD copies, ~474 GB / 223 ms per request). graph.code / Tensor types would make TensorifyScalarRestartAnalysis's second MagiBackend __call__ miss the bake directory (assert_cache_hit / fix: allow restart-analysis replay in assert_cache_hit mode #75 skip-retry never starts).
• feat(aot): graceful fallback to JIT when AOT compile fails #86 AOT → JIT fallback: aot_compile RuntimeError (SimpleFSDP/DTensor, "Attempted to read undefined local variable") falls back to JIT on that MagiCompileState so MAGI_COMPILE_AOT=1 can stay process-wide. AssertionError and CUDA OOM are not treated as "AOT unsupported". VAE still AOT-compiles; base/SR do not abort warmup.
• fix: skip FSDP fullgraph overlap for non-fullgraph models (e.g. VAE) #79 FSDP overlap skip only: enable_fsdp is process-wide; _apply_fsdp_fullgraph_overlap still asserts disable_graph_split=True. VAE (graph split) in a monolithic FSDP process no longer hits that assert. The restart-analysis half of fix: skip FSDP fullgraph overlap for non-fullgraph models (e.g. VAE) #79 already landed as fix: allow restart-analysis replay in assert_cache_hit mode #75.
• fix(test): increase model dim for stable timing benchmark #84 timing bench: SimpleModel dim 128→1024 / seq 256→512 so GPU compute dominates launch jitter and max/min < 1.2 stays stable on B300/H100.

Dropped: #79's assert_cache_hit restart-analysis change (already on main).

Test plan

• test_split_graph_skips_fsdp_overlap_when_graph_split_enabled fails on unfixed _split_graph, passes after the guard
• test_aot_failure_falls_back_to_jit raises on unfixed _api.py, passes after fallback
• test_aot_assertion_error_is_not_swallowed fails on except Exception, passes after except RuntimeError
• test_compute_graph_hash_stable_across_tensorify_scalar (CPU)
• CI: test_placement_cache_key.py (torchrun, 1 GPU, SimpleFSDP fully_shard vs replicate)
• CI: existing test_simple_model_timing_class_function_instance_method (pt29+pt212 Integration Test green on 9ac1a9d; re-run on 3fc8a63)

[cennn]

Code Review: MagiCompiler PR #89

• PR: #89
• Head: 3fc8a63 (fix/consolidate-cenn-prs)
• Base: main
• Reviewer note: 审查过程中发现 [P0-BLOCKER]（except Exception 吞掉 AssertionError），已在 3fc8a63 修复并推送。下文按修复后的 HEAD 写；未修的项不构成 merge blocker。

Findings

• [P0-BLOCKER] 已修复 (3fc8a63) — _run_orchestration 原 except Exception 把 AssertionError 当成 AOT-unsupported，静默 JIT。现改为 except RuntimeError，并 re-raise torch.cuda.OutOfMemoryError。回归测试：test_aot_assertion_error_is_not_swallowed（unfixed 不 raise；fixed 后 raise）。
• [P2-MAINTAIN] magi_compiler/magi_backend/magi_backend.py Step 1.4：enable_fsdp=True 且 disable_graph_split=False 时 skip fullgraph overlap 没有 warning。VAE 同进程是合法 skip；FSDP DiT 若忘了 config_patch 强制 fullgraph，overlap 会静默失效。应打 magi_logger.warning（不要全局 assert / 不要强制 disable_graph_split）。
• [P3-STYLE] compute_graph_hash(graph, example_inputs) 的 graph 未使用。签名留给 _init_cache 调用点；请在参数上标明故意忽略，避免下一次又把 graph.code 塞回去。
• [P4-PROCESS] PR body 仍写 cache key = graph.code + Tensor dtype + placements，且测试名还是 test_compute_graph_hash_changes_with_input_dtype。HEAD 已是 只 hash DTensor placement；test_graph_hash.py 测的是 tensorify scalar 稳定性。Body 与 helper docstring（--disable-fix drops the graph signature）需改到与代码一致。

无 [P1-PERF]。

---

1. PR Summary & Technical Background

本 PR 把尚未合入的 #88 / #86 / #84 / #79 中现在仍正确的部分叠到当前 main（#75 的 restart-analysis / assert_cache_hit 已在 main，本 PR 不再带那一半）。

四件事：

1. Cache key 补 DTensor placement。 Magi cache 目录名原先是 compile_config + inductor + compiler + traced_files。SimpleFSDP fully_shard（Shard(0) → forward 前 AllGather）改成 replicate（Replicate() → AllGather 退化成 copy）时源文件不变，会错 load 旧 artifact。400B 8×H200 上表现为 1-rank AllGather 当 NCCL DtoD copy。compute_graph_hash 把 example_inputs 里 DTensor 的 (dtype, placements, mesh.shape) 编进目录名。graph.code 和普通 Tensor/scalar 类型不能进 key：TensorifyScalarRestartAnalysis 会让 Dynamo 用新 GraphModule 再 __call__，verify 第一次必须打开 bake 成功那次的同一目录，否则 #75 的 restart_analysis_count skip/retry 在 cache_entry not in self.cache 处直接炸。
2. AOT → JIT fallback。 MAGI_COMPILE_AOT=1 是进程级；SimpleFSDP/DTensor 的 aot_compile_fullgraph 会 RuntimeError: Attempted to read undefined local variable。失败时本 state aot=False + dynamo.reset() + JIT capture，VAE 仍可 AOT。
3. FSDP overlap 只在 fullgraph。 _apply_fsdp_fullgraph_overlap 仍 assert disable_graph_split=True。enable_fsdp 来自 env，VAE 的 @magi_compile 无 config_patch 会继承它。Skip 避免 monolithic 进程里 VAE 被 assert 打死。
4. Timing bench dim 128→1024 / seq 256→512。 小图 launch-bound，max/min < 1.2 在 B300/H100 上 flake。

变更是必要的：placement 错 cache、AOT 中断 warmup、VAE×FSDP assert 都是生产问题，不是空转重构。最终 cache key 比 #88 初稿更窄，是 CI 证明过的约束，不是画蛇添足。

---

2. Code Quality & Correctness

逐文件。P0 已在 3fc8a63 落地。

magi_compiler/_api.py

aot_compile 单次调用外包 try，保护域够窄。3fc8a63 前 except Exception 违反 Fail Fast / 禁止 over-catch：AssertionError、其它非 AOT-unsupported 的 Exception 被转成 JIT。现 except RuntimeError + OOM re-raise；AssertionError / TypeError 继续往外抛。Fallback 路径：model_copy(aot=False)、清 compiled_entry、torch._dynamo.reset()、再 _ensure_compiled + JIT capture，状态复位完整。

magi_compiler/utils/hash.py + magi_backend.py _init_cache

compute_graph_hash 只扫 example_inputs 里的 DTensor。placement 测试（torchrun 1 GPU，--disable-fix 把该项打成 ""）在 pt29/pt212 绿过：example_inputs 里确实有不同 placement 的 DTensor，否则 test_bug_fixed_placement_change_recompiles 过不了。graph 参数未使用（[P3-STYLE]）。不要把 graph.code 加回来。

magi_backend.py _split_graph Step 1.4

enable_fsdp and disable_graph_split 才进 overlap，契约正确（Athena DiT config_patch 在 enable 时强制 fullgraph；VAE 保持 split）。缺 warning（[P2-MAINTAIN]）。全局 assert enable_fsdp ⇒ disable_graph_split 或强制 split 会把 VAE 合法路径打死，不要改成 raise/强制。

Tests

• test_placement_cache_key.py：reproduce（--disable-fix → 错 hit compiles==0）+ fixed（换 mode 必须重编）。信号是 num_inductor_compiles，不是 max_abs_diff（1 卡错 load 数值仍可能过）。
• test_graph_hash.py：钉死 Tensor vs int / dtype 不变 hash。不覆盖 DTensor；那部分在 torchrun 测试。
• test_aot_jit_fallback.py：RuntimeError → JIT；3fc8a63 增加 AssertionError 必须冒泡。
• test_fsdp_skip_split.py：monkeypatch _apply_fsdp_fullgraph_overlap，测的是门闩不是 overlap 正确性，对这个 PR 够用。
• timing dim 只改常量，无新断言逻辑。

---

3. Goal Completeness

对照 PR 描述（按 HEAD 行为，不是过时 body）：

目标	是否完成
fully_shard vs replicate 不再错 reuse cache	是。CI test_placement_cache_key 双 pipeline 绿
restart-analysis bake-verify 仍 hit	是。graph.code / Tensor 类型已移出 key；#75 测试绿
AOT 失败不中断进程内其它模型	是。state 级 aot=False；3fc8a63 不再吞 programmer error
VAE split-graph 不进 fullgraph overlap	是。缺的是可观测性（warning），不是行为
timing 在 B300 稳定	是。pt29/pt212 perf 绿

未完成 / 易误导：PR body 与 helper docstring 仍描述 #88 初稿（graph.code + dtype）。请改 body，避免 reviewer 按错误契约审。

---

4. AI-Generated Code Detection

Consolidation 痕迹明显，但不是无脑粘贴：cache key 被 CI（restart-analysis miss → 去掉 graph.code → 仍 miss → 只留 DTensor）收窄过两轮，placement / AOT / FSDP 都有「先红后绿」测试。

仍像 LLM 初稿的点：

• 原 except Exception（已修）。
• helper / PR body 的 “graph signature” 在 key 语义改了之后没同步。
• compute_graph_hash 留着不用的 graph，像是从 hash graph.code 的版本删改残留。

整体能看出作者理解 Magi cache 目录 vs #75 skip/retry、进程级 enable_fsdp vs per-graph split。不构成「vibe coding」否决；把 body 和 unused graph 收干净即可。

[cennn]

Code Review: PR #89 fix/consolidate-cenn-prs @ 3fc8a63（第二轮）

结论：本轮没有 P0-BLOCKER。 上一轮的 P0（except Exception 吞掉 AssertionError）已在 3fc8a63 修复，CI run 38679281013 全绿。这一轮针对 AOT→JIT fallback 剩下的副作用做了 3 个实测，都没有发现正确性问题；剩下的都是 P2–P4。

1. PR Summary & Technical Background

合并 #88/#86/#84/#79 中仍然有用的三处修改：

• Cache key 加入 DTensor placement（utils/hash.py::compute_graph_hash，magi_backend.py::_init_cache）：SimpleFSDP 从 fully_shard（Shard(0)，forward 里 AllGather）切换到 replicate（Replicate()）时，traced source 完全一样，旧 key 会复用错误的 artifact。只哈希 DTensor 的 dtype/placements/mesh.shape，不哈希 graph.code 和普通 Tensor/scalar 类型，这样 #75 的 TensorifyScalarRestartAnalysis 两次 trace 仍然落在同一个 cache 目录里。
• FSDP fullgraph overlap 只对 disable_graph_split=True 的图生效：enable_fsdp 是进程级开关，VAE 这类 split-graph 模型不应该进入 overlap。
• AOT 失败回退 JIT（_api.py::_run_orchestration）：SimpleFSDP/DTensor 不支持 aot_compile_fullgraph，而 MAGI_COMPILE_AOT=1 是进程级的。
• 测试：placement 端到端（含 --disable-fix 复现旧 bug）、restart 稳定性单测、FSDP gate 单测、AOT fallback / AssertionError 不被吞的单测；timing 测试改大尺寸以降低抖动。

2. Code Quality & Correctness

本轮实测（均在 cenn-5 @ 3fc8a63）

疑点	实测	结论
torch._dynamo.reset() 是进程级的，会不会让先编译好的模型 A 失效	A 先 JIT 编译 → B 触发 AOT 失败 fallback → A 以新 shape 和旧 shape 再调用	(graphs_seen, backend_compilations) 保持 (2, 2)，A 未重编译，输出正确
aot=True + assert_cache_hit=True 时，AOT 的 cache-miss RuntimeError 被 fallback 接住，会不会绕过断言	空 cache 下调用	JIT 重试再次报 MAGI_COMPILE_ASSERT_CACHE_HIT: cache miss，以 BackendCompilerFailed 抛出，断言没有被绕过
AOT 半途失败会不会留下损坏的 artifact	读代码	save_aot_compile_artifacts() 只在 aot_compile 成功返回后调用，失败路径不会写盘
CompileConfig.hash 包含 aot，fallback 的目录是否稳定	aot=False→481c3441c5，aot=True→18de276159	两者目录不同，但 fallback 每次都固定用 aot=False 的目录，跨进程稳定

Issues

• [P2-MAINTAIN] magi_compiler/_api.py:162 — except RuntimeError 的异常类型仍然偏宽。 保护区只有一行（符合“窄保护区”），但所有 RuntimeError 都会被当成 “AOT 不支持”：真实的 inductor/triton 编译错误（BackendCompilerFailed）和 assert_cache_hit 的 miss 都会先打一条 AOT compile failed ... falling back 的 warning，再走一遍 JIT。正确性不受影响（上表第 2 行），但会有两个后果：VAE 真实的 AOT 回归会被降级成一条 warning，每次启动都多付一次失败的 trace；assert_cache_hit 的报错前面多一条误导性日志。建议：在 fallback 前对 assert_cache_hit 直接 raise；或者只在 example inputs 含 DTensor 时才回退，其余情况 re-raise。
• [P2-MAINTAIN] magi_backend.py:681-684 — FSDP overlap 被跳过时没有日志。 enable_fsdp=True 但 disable_graph_split=False 时会静默跳过。对 VAE 这是预期行为，但 DiT 如果漏配 config_patch，FSDP overlap 就会无声失效。建议：加一行 magi_logger.info("enable_fsdp is set but graph split is enabled for %s; skipping fullgraph FSDP overlap", self.model_tag)。
• [P3-STYLE] utils/hash.py:59 — compute_graph_hash 的 graph 参数没有用到。 函数名和签名暗示图结构参与了哈希，实际没有，容易被后人误以为 graph.code 已经在 key 里。建议：改名为 compute_input_placement_hash(example_inputs)，或者在 docstring 里写明“故意不哈希 graph”。
• [P3-STYLE] _api.py:172 — fallback 修改了 state.compile_config。 这是 state 上的持久变更，之后同一个 state 永远走 JIT。行为符合预期，但值得在 warning 文案里写出 “for the rest of this process”。
• [P4-PROCESS] placement_cache_key_helper.py:277 — docstring 过期。 “drops the graph signature” 已经不准确，实际 key 只含 DTensor placement。建议：改为 “drops the DTensor placement hash from the magi cache key”。
• [P4-PROCESS] 升级会让所有已有 cache 失效。 _init_cache 的 key 从 4 项变成 5 项（即使没有 DTensor 也会多一个空列表的哈希），所有预烘焙的 cache 都会 miss。在 assert_cache_hit 的生产环境里升级必须重新烘焙。建议：在 PR 描述或 release note 里写明。

3. Goal Completeness

• placement 复用 bug：test_bug_reproduced_without_fix 能复现，test_bug_fixed_placement_change_recompiles 验证修复，并验证同 placement 仍命中，完整。
• restart-analysis 兼容：test_graph_hash 单测加上 CI 中已有的 test_assert_cache_hit_succeeds_with_restart_analysis 都通过，完整。
• FSDP gate：两条分支都有单测，完整；缺少跳过时的可观测性（见 P2）。
• AOT fallback：正常回退和 AssertionError 不被吞都有测试，完整；尚未覆盖 assert_cache_hit 与 fallback 的交互（本轮已手工验证行为正确，可以考虑补一个测试）。

4. AI-Generated Code Detection

• 注释密度适中，没有“下一行做什么”式的注释。hash.py 和 _api.py 的注释说明的是约束（restart retrace、SimpleFSDP 不支持 AOT），合格。
• except RuntimeError + OutOfMemoryError 特判，是典型的“防御性兜底”写法。结合 P0 规则 “Do NOT overprotect”，需要在 PR 描述里讲清楚为什么必须兜底（SimpleFSDP AOT 确定性失败、VAE 依赖 AOT），这一点目前已在注释中说明。
• 测试命名和结构与仓库现有风格一致，没有模板化的冗余。

[jiahy0825]

LGTM