Skip to content

[WIP][Fix] Key compile state by topology and resolve it per call - #49

Open
cennn wants to merge 3 commits into
mainfrom
fix/compile-state-per-topology
Open

[WIP][Fix] Key compile state by topology and resolve it per call#49
cennn wants to merge 3 commits into
mainfrom
fix/compile-state-per-topology

Conversation

@cennn

@cennn cennn commented Jul 29, 2026

Copy link
Copy Markdown
Collaborator

背景与问题

推理运行时会在请求之间改变并行拓扑:adaptive DP 根据队列深度在 cp=8/dp=1cp=4/dp=2 之间切换。
在这种场景下,magi_compile 会把第一个拓扑编译出的图一直用下去。

原因是两处叠加:

  • compile state 挂在固定属性 _magi_state_for_<entry> 上,所有拓扑共用一份;
  • 该属性名在 wrap 时(模型构造那一刻)解析一次,被 new_call 的闭包捕获,之后不再变化。

而 state 持有的是已捕获的 bytecode / AOT 产物,_run_orchestration 的 fast path 直接 replay 它们,
不经过 dynamo 的 guard;这些产物又把 trace 时的 ProcessGroup 固化在图里(ProcessGroup 不是 tensor,
dynamo 无法把它变成图输入)。于是拓扑切换后,replay 的图仍在旧通信组上通信。

在下游(athena 的 gaga4 disagg 推理)表现为切换后 dist.all_to_all_single
NCCL Error 1: unhandled cuda error / NCCL Error 3: internal error

本 PR 做了什么

  1. compile state 按拓扑分身:属性名附加 MAGI_COMPILE_TOPOLOGY_KEY(运行时在每次 mesh 变更时刷新,
    [Fix] Include parallel topology in cache directory to prevent cross-contamination #42 引入的 cache 目录 key 同源)。每个拓扑各自编译一次、各自保留,回访时复用。
  2. 属性名改为每次调用解析:这是让上一条真正生效的关键——原先在 wrap 时固定,分身逻辑根本接触不到运行时的拓扑变化。
  3. 重复包装的判定改用 wrapper installed flag:原先借 state 属性是否存在来判断,state 分身后该判定不再可靠。

验证

4 rank / gloo 的最小 probe(编译一个读取当前 CP 组并据此计算的函数,切换拓扑后再调用):

编译方式 cp=4 切到 cp=2 切回 cp=4
torch.compile(dynamic False/True/None 三种) 4.0 2.0 跟随切换 4.0
magi_compile(本 PR 之前) 4.0 4.0 复用了 cp=4 的图 4.0
magi_compile(本 PR 之后) 4.0 2.0 跟随切换 4.0

torch.compile 一列说明它自己的 guard 本就能识别组的变化,问题只在 fast path 绕过 guard 的这条路径上。

收益

  • 正确性:消除拓扑切换后在错误通信组上通信的问题。
  • 性能:拓扑回访直接命中已编译产物。此前下游只能在每次切换时用 torch._dynamo.reset() 全量丢弃图来规避,
    代价是完整重新 trace——在 gaga4 上实测单步 denoise 从秒级变成 2m23s / 2m29s / 2m47s。

[附件位置:adaptive DP 切换前后的端到端时间线对比(甘特图 + 原始日志),待补]

其他

无行为变化的路径:未设置 MAGI_COMPILE_TOPOLOGY_KEY 时属性名与此前完全一致,单一固定拓扑的用法不受影响。

cennn added 3 commits July 29, 2026 16:33
A runtime that changes parallel topology between calls -- adaptive DP switching CP/DP
by queue depth -- kept replaying the graph captured under the first topology, because
the compile state lived on a fixed attribute resolved once at wrap time.

The state owns bytecode and AOT artifacts that the fast paths replay without consulting
dynamo guards, and those artifacts have the ProcessGroup they traced with baked in, so
replaying them under another topology communicates on the wrong group.

Verified with a 4-rank probe: plain torch.compile follows a cp=4 -> cp=2 switch on its
own, magi_compile returned the cp=4 result until this change.
Four cases around the compile state: that sharing one across topologies is what lets one
graph serve another, that each topology gets its own, that a single fixed topology keeps
the plain attribute name, and -- the regression -- that the topology live at call time
decides, not the one live when the instance was wrapped.

The last one fails against the previous behaviour with the wrap-time name, which is what
made keying the state by topology inert.
Keying the compile state by topology meant the wrap-time idempotency check could no
longer be "does the state attribute exist", since that name now moves with the topology,
so it became the wrapper-installed flag. But _magi_compile_class sets that flag on the
class and then patches every new instance from __init__, and an instance inherits its
class's attributes -- so _magi_compile_bound_method returned early for each one. Nothing
got patched: forward stayed the class's own, ran eagerly, and no compile state was ever
attached.

That is what the 23 CI failures were, all one cause: a missing _magi_state_for_forward,
magi timings equal to eager, zero CUDA graphs, zero inductor artifacts, spies never
reached, and CPU-offload inputs left on the host because the wrapper that offloads them
was never installed.

The marker used here lives only on instances, so it cannot be inherited, and it names the
method, so one instance can still carry several entry points. The class-level flag keeps
its own job of not re-wrapping __init__.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant