[WIP][Fix] Key compile state by topology and resolve it per call - #49
Open
cennn wants to merge 3 commits into
Open
[WIP][Fix] Key compile state by topology and resolve it per call#49cennn wants to merge 3 commits into
cennn wants to merge 3 commits into
Conversation
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__.
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 join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
背景与问题
推理运行时会在请求之间改变并行拓扑:adaptive DP 根据队列深度在
cp=8/dp=1与cp=4/dp=2之间切换。在这种场景下,
magi_compile会把第一个拓扑编译出的图一直用下去。原因是两处叠加:
_magi_state_for_<entry>上,所有拓扑共用一份;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 做了什么
MAGI_COMPILE_TOPOLOGY_KEY(运行时在每次 mesh 变更时刷新,与 [Fix] Include parallel topology in cache directory to prevent cross-contamination #42 引入的 cache 目录 key 同源)。每个拓扑各自编译一次、各自保留,回访时复用。
验证
4 rank / gloo 的最小 probe(编译一个读取当前 CP 组并据此计算的函数,切换拓扑后再调用):
torch.compile(dynamic False/True/None 三种)magi_compile(本 PR 之前)magi_compile(本 PR 之后)torch.compile一列说明它自己的 guard 本就能识别组的变化,问题只在 fast path 绕过 guard 的这条路径上。收益
torch._dynamo.reset()全量丢弃图来规避,代价是完整重新 trace——在 gaga4 上实测单步 denoise 从秒级变成 2m23s / 2m29s / 2m47s。
其他
无行为变化的路径:未设置
MAGI_COMPILE_TOPOLOGY_KEY时属性名与此前完全一致,单一固定拓扑的用法不受影响。