From 4b0978bf37ded66074c385dd08553307f10d3d50 Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Wed, 29 Jul 2026 16:33:29 +0800 Subject: [PATCH 1/3] [Fix] Key compile state by topology and resolve it per call 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. --- magi_compiler/_api.py | 29 ++++++++++++++++++++++++----- 1 file changed, 24 insertions(+), 5 deletions(-) diff --git a/magi_compiler/_api.py b/magi_compiler/_api.py index 43ae1da..b29c5f2 100644 --- a/magi_compiler/_api.py +++ b/magi_compiler/_api.py @@ -77,7 +77,22 @@ def get_attr_name_for_wrapper_installed_flag() -> str: def get_attr_name_for_state(entry_name: str) -> str: - return f"_magi_state_for_{entry_name}" + """Name the attribute holding an entry point's compile state, one per topology. + + The state owns the captured bytecode and AOT artifacts that _run_orchestration's fast + paths replay directly, without going through dynamo's guards. Those artifacts have the + ProcessGroup they traced with baked in, because a ProcessGroup is not something dynamo + can turn into a graph input. So a runtime that changes CP/DP between calls -- what + adaptive DP does between requests -- needs one state per topology; a shared one would + replay the first topology's graph under every later one. Keyed this way each topology + compiles once and is reused whenever the runtime returns to it. + + MAGI_COMPILE_TOPOLOGY_KEY is set by the runtime on every mesh change; empty means a + single fixed topology, i.e. the plain name. + """ + topology = os.environ.get("MAGI_COMPILE_TOPOLOGY_KEY", "") + suffix = f"__{topology}" if topology else "" + return f"_magi_state_for_{entry_name}{suffix}" def _run_orchestration(state: MagiCompileState, args, kwargs): @@ -209,14 +224,17 @@ def _magi_compile_bound_method( if not callable(getattr(instance, method_name, None)): raise AttributeError(f"{instance.__class__.__name__} instance has no callable method '{method_name}'") - state_attr = get_attr_name_for_state(method_name) - if getattr(instance, state_attr, None) is not None: + if getattr(instance, get_attr_name_for_wrapper_installed_flag(), False): return instance old_method = getattr(instance, method_name) @torch.compiler.disable() def new_call(*args, **kwargs): + # Per call, not per wrap: the name carries the topology, and adaptive DP changes + # topology between calls. Binding it at construction pinned every later call to + # the topology that happened to be live back then. + state_attr = get_attr_name_for_state(method_name) state = getattr(instance, state_attr, None) if state is None: _lazy_init_magi_state(instance, instance, dynamic_arg_dims, conf, model_tag, method_name, state_attr) @@ -238,13 +256,13 @@ def new_call(*args, **kwargs): def _magi_compile_function(func: Callable, dynamic_arg_dims: dict[str, int | list[int]], conf: CompileConfig, model_tag: str): """Wrap a function entry with compiled routing.""" - state_attr = get_attr_name_for_state("function") - if getattr(func, state_attr, None) is not None: + if getattr(func, get_attr_name_for_wrapper_installed_flag(), False): return func @torch.compiler.disable() @functools.wraps(func) # for the original function name and docstring def wrapper(*args, **kwargs): + state_attr = get_attr_name_for_state("function") # per call: see _magi_compile_bound_method state = getattr(wrapper, state_attr, None) if state is None: _lazy_init_magi_state(wrapper, func, dynamic_arg_dims, conf, model_tag, None, state_attr) @@ -255,6 +273,7 @@ def wrapper(*args, **kwargs): return _run_orchestration(state, args, kwargs) + setattr(wrapper, get_attr_name_for_wrapper_installed_flag(), True) return wrapper From d09f80f7f0a4b40a3e244e6c2143d9cd65566579 Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Wed, 29 Jul 2026 17:00:34 +0800 Subject: [PATCH 2/3] [Test] Cover the in-memory half of topology isolation 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. --- .../test_cache_topology_isolation.py | 79 ++++++++++++++++++- 1 file changed, 77 insertions(+), 2 deletions(-) diff --git a/tests/feature_tests/test_cache_topology_isolation.py b/tests/feature_tests/test_cache_topology_isolation.py index bcc0b21..9f8da5a 100644 --- a/tests/feature_tests/test_cache_topology_isolation.py +++ b/tests/feature_tests/test_cache_topology_isolation.py @@ -3,16 +3,24 @@ # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. -"""Verify that different parallel topologies produce isolated cache directories, -preventing EP/CP cross-contamination of compiled artifacts.""" +"""Verify that different parallel topologies stay isolated, on disk and in memory. + +Compiled artifacts have the ProcessGroup they traced with baked in, so a runtime that +changes CP/DP between calls -- adaptive DP reacting to queue depth -- must not reuse one +topology's artifacts under another. Two things have to be keyed by topology for that: +the cache directory, and the compile state held on the instance. +""" from __future__ import annotations +import os from pathlib import Path +from types import SimpleNamespace from unittest.mock import patch import torch +from magi_compiler._api import get_attr_name_for_state from magi_compiler.config import _get_parallel_topology, magi_cache_dump_path, model_rank_dir_name @@ -94,3 +102,70 @@ def test_same_topology_same_path(self): name1 = model_rank_dir_name(0, None) name2 = model_rank_dir_name(0, None) assert name1 == name2 + + +class TestCompileStateIsolation: + """The compile state on an instance is the in-memory counterpart of the cache dir. + + It owns the bytecode and AOT artifacts that the fast paths in _run_orchestration replay + directly, without going through dynamo's guards -- so if two topologies share one state, + the second silently runs the first one's graph, on the first one's process group. + """ + + def test_a_shared_state_is_what_lets_one_graph_serve_two_topologies(self): + """The failure mode: with the key pinned, both topologies land on one attribute.""" + with patch.dict(os.environ, {"MAGI_COMPILE_TOPOLOGY_KEY": "pinned"}): + under_cp8 = get_attr_name_for_state("forward") + under_cp4 = get_attr_name_for_state("forward") + assert under_cp8 == under_cp4 + + def test_each_topology_gets_its_own_state(self): + with patch.dict(os.environ, {"MAGI_COMPILE_TOPOLOGY_KEY": "cp8_dp1"}): + under_cp8 = get_attr_name_for_state("forward") + with patch.dict(os.environ, {"MAGI_COMPILE_TOPOLOGY_KEY": "cp4_dp2"}): + under_cp4 = get_attr_name_for_state("forward") + assert under_cp8 != under_cp4 + assert "cp8_dp1" in under_cp8 and "cp4_dp2" in under_cp4 + + def test_a_single_fixed_topology_keeps_the_plain_name(self): + """No key means no adaptive runtime, and the attribute must not change shape.""" + with patch.dict(os.environ, {}, clear=False): + os.environ.pop("MAGI_COMPILE_TOPOLOGY_KEY", None) + assert get_attr_name_for_state("forward") == "_magi_state_for_forward" + + def test_the_topology_at_call_time_decides_not_the_one_at_wrap_time(self): + """The regression this guards. + + The attribute used to be resolved once while wrapping the instance, so a model built + under cp=8 kept reaching for the cp=8 state no matter what the runtime switched to + afterwards -- which made keying by topology inert. + """ + from magi_compiler._api import _magi_compile_bound_method + + reached = [] + + def fake_init(holder, target, dims, conf, tag, method, state_attr): + reached.append(state_attr) + setattr( + holder, + state_attr, + SimpleNamespace( + compile_config=SimpleNamespace(offload_config=SimpleNamespace(model_cpu_offload=False)), + jit_compiled_code=None, + ), + ) + + class Probe(torch.nn.Module): + def forward(self, x): + return x + + probe = Probe() + with patch.dict(os.environ, {"MAGI_COMPILE_TOPOLOGY_KEY": "cp8_dp1"}): + _magi_compile_bound_method(probe, {"x": 0}, SimpleNamespace(), "probe", method_name="forward") + + with patch.dict(os.environ, {"MAGI_COMPILE_TOPOLOGY_KEY": "cp4_dp2"}): + with patch("magi_compiler._api._lazy_init_magi_state", side_effect=fake_init): + with patch("magi_compiler._api._run_orchestration", return_value=None): + probe.forward(torch.zeros(2)) + + assert reached == ["_magi_state_for_forward__cp4_dp2"], reached From e68d5ee59d4b0d02f6580e75f754b537b329752f Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Wed, 29 Jul 2026 18:40:45 +0800 Subject: [PATCH 3/3] [Fix] Keep the instance-patched marker off the class 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__. --- magi_compiler/_api.py | 16 ++++++++++++-- .../test_cache_topology_isolation.py | 22 ++++++++++++++++++- 2 files changed, 35 insertions(+), 3 deletions(-) diff --git a/magi_compiler/_api.py b/magi_compiler/_api.py index b29c5f2..c6dc57f 100644 --- a/magi_compiler/_api.py +++ b/magi_compiler/_api.py @@ -76,6 +76,17 @@ def get_attr_name_for_wrapper_installed_flag() -> str: return "_magi_wrapper_installed" +def get_attr_name_for_bound_wrapper_flag(entry_name: str) -> str: + """Name the marker saying *entry_name* has already been patched on this instance. + + Deliberately not the class-level flag: the class decorator sets that one and then + patches each new instance from __init__, and an instance inherits it, so reading it + here would skip the patching it is supposed to guard. Per method, so an instance can + carry more than one entry point. + """ + return f"_magi_wrapper_installed_for_{entry_name}" + + def get_attr_name_for_state(entry_name: str) -> str: """Name the attribute holding an entry point's compile state, one per topology. @@ -224,7 +235,8 @@ def _magi_compile_bound_method( if not callable(getattr(instance, method_name, None)): raise AttributeError(f"{instance.__class__.__name__} instance has no callable method '{method_name}'") - if getattr(instance, get_attr_name_for_wrapper_installed_flag(), False): + installed_attr = get_attr_name_for_bound_wrapper_flag(method_name) + if getattr(instance, installed_attr, False): return instance old_method = getattr(instance, method_name) @@ -250,7 +262,7 @@ def new_call(*args, **kwargs): return _run_orchestration(state, args, kwargs) setattr(instance, method_name, new_call) - setattr(instance, get_attr_name_for_wrapper_installed_flag(), True) + setattr(instance, installed_attr, True) return instance diff --git a/tests/feature_tests/test_cache_topology_isolation.py b/tests/feature_tests/test_cache_topology_isolation.py index 9f8da5a..c7df040 100644 --- a/tests/feature_tests/test_cache_topology_isolation.py +++ b/tests/feature_tests/test_cache_topology_isolation.py @@ -20,7 +20,7 @@ import torch -from magi_compiler._api import get_attr_name_for_state +from magi_compiler._api import get_attr_name_for_state, get_attr_name_for_wrapper_installed_flag from magi_compiler.config import _get_parallel_topology, magi_cache_dump_path, model_rank_dir_name @@ -169,3 +169,23 @@ def forward(self, x): probe.forward(torch.zeros(2)) assert reached == ["_magi_state_for_forward__cp4_dp2"], reached + + def test_the_class_decorator_still_patches_every_instance(self): + """A mark on the class must not answer for an instance. + + magi_compile on a class marks the class and patches each instance from __init__. + Reading that inherited mark to decide whether an instance was patched made every + instance skip it: forward stayed the class's own and ran eagerly, and no state was + ever attached -- which is invisible except as timings that match eager. + """ + from magi_compiler import magi_compile + + @magi_compile(dynamic_arg_dims={"x": 0}) + class Block(torch.nn.Module): + def forward(self, x): + return x + + block = Block() + + assert getattr(Block, get_attr_name_for_wrapper_installed_flag(), False), "the class went unmarked" + assert "forward" in vars(block), "the instance kept the class's forward, so nothing was compiled"