已合并
[feat] deterministic_level: add graph-mode guard and EX scope API #42191
Lyric创建于 7月20日
[feat] deterministic_level: add graph-mode guard and EX scope API #42191
已合并
共 6 个文件变更+117-1
| @@ -0,0 +1,37 @@ | |||
| 1 | +import unittest | ||
| 2 | + | ||
| 3 | +import torch | ||
| 4 | + | ||
| 5 | +import torch_npu | ||
| 6 | + | ||
| 7 | + | ||
| 8 | +class TestDeterministicLevelGraphMode(unittest.TestCase): | ||
| 9 | + def tearDown(self): | ||
| 10 | + torch._dynamo.reset() | ||
| 11 | + torch_npu.npu.set_deterministic_level(0) | ||
| 12 | + super().tearDown() | ||
| 13 | + | ||
| 14 | + def test_deterministic_level_guard(self): | ||
| 15 | + from torch_npu.dynamo._deterministic_guard import install_npu_deterministic_level_guard | ||
| 16 | + | ||
| 17 | + compile_count = 0 | ||
| 18 | + | ||
| 19 | + def backend(gm, example_inputs): | ||
| 20 | + nonlocal compile_count | ||
| 21 | + self.assertTrue(install_npu_deterministic_level_guard()) | ||
| 22 | + compile_count += 1 | ||
| 23 | + return gm.forward | ||
| 24 | + | ||
| 25 | + def fn(x): | ||
| 26 | + return x + 1 | ||
| 27 | + | ||
| 28 | + compiled_fn = torch.compile(fn, backend=backend, fullgraph=True, dynamic=False) | ||
| 29 | + for level in (1, 2, 1, 2): | ||
| 30 | + torch_npu.npu.set_deterministic_level(level) | ||
| 31 | + torch.testing.assert_close(compiled_fn(torch.ones(2)), torch.full((2,), 2.0)) | ||
| 32 | + | ||
| 33 | + self.assertEqual(compile_count, 2) | ||
| 34 | + | ||
| 35 | + | ||
| 36 | +if __name__ == "__main__": | ||
| 37 | + unittest.main() | ||
| @@ -2309,6 +2309,9 @@ | |||
| 2309 | "torch_npu.npu.npugraph_ex.scope.limit_core_num": { | 2309 | "torch_npu.npu.npugraph_ex.scope.limit_core_num": { |
| 2310 | "signature": "(op_aicore_num: int, op_vectorcore_num: int, stream=None)" | 2310 | "signature": "(op_aicore_num: int, op_vectorcore_num: int, stream=None)" |
| 2311 | }, | 2311 | }, |
| 2312 | + "torch_npu.npu.npugraph_ex.scope.deterministic": { | ||
| 2313 | + "signature": "(level: int)" | ||
| 2314 | + }, | ||
| 2312 | "torch_npu.npu.npugraph_ex.compile_fx": { | 2315 | "torch_npu.npu.npugraph_ex.compile_fx": { |
| 2313 | "signature": "(gm, example_inputs=None, options=None)" | 2316 | "signature": "(gm, example_inputs=None, options=None)" |
| 2314 | }, | 2317 | }, |
| @@ -0,0 +1,32 @@ | |||
| 1 | +import os | ||
| 2 | + | ||
| 3 | + | ||
| 4 | +def _is_ascendc_backend() -> bool: | ||
| 5 | + return os.getenv("TORCHINDUCTOR_NPU_BACKEND") == "ascendc" | ||
| 6 | + | ||
| 7 | + | ||
| 8 | +def patch_npu_deterministic_level_cache_keys(): | ||
| 9 | + """Add the exact NPU deterministic level to AscendC Inductor cache keys.""" | ||
| 10 | + import torch_npu | ||
| 11 | + from torch._functorch._aot_autograd.autograd_cache import AOTAutogradCacheDetails | ||
| 12 | + from torch._inductor.codecache import FxGraphHashDetails | ||
| 13 | + | ||
| 14 | + if getattr(FxGraphHashDetails, "_npu_deterministic_level_patched", False): | ||
| 15 | + return | ||
| 16 | + | ||
| 17 | + fx_graph_hash_details_init = FxGraphHashDetails.__init__ | ||
| 18 | + aot_autograd_cache_details_init = AOTAutogradCacheDetails.__init__ | ||
| 19 | + | ||
| 20 | + def fx_graph_hash_details_init_with_npu_deterministic_level(self, *args, **kwargs): | ||
| 21 | + fx_graph_hash_details_init(self, *args, **kwargs) | ||
| 22 | + if _is_ascendc_backend(): | ||
| 23 | + self.npu_deterministic_level = torch_npu.npu._get_deterministic_level() | ||
| 24 | + | ||
| 25 | + def aot_autograd_cache_details_init_with_npu_deterministic_level(self, *args, **kwargs): | ||
| 26 | + aot_autograd_cache_details_init(self, *args, **kwargs) | ||
| 27 | + if _is_ascendc_backend(): | ||
| 28 | + self.npu_deterministic_level = torch_npu.npu._get_deterministic_level() | ||
| 29 | + | ||
| 30 | + FxGraphHashDetails.__init__ = fx_graph_hash_details_init_with_npu_deterministic_level | ||
| 31 | + AOTAutogradCacheDetails.__init__ = aot_autograd_cache_details_init_with_npu_deterministic_level | ||
| 32 | + FxGraphHashDetails._npu_deterministic_level_patched = True | ||
| @@ -0,0 +1,27 @@ | |||
| 1 | +def install_npu_deterministic_level_guard() -> bool: | ||
| 2 | + """Guard the NPU deterministic_level to trigger recompilation when level changes.""" | ||
| 3 | + import torch._guards as _guards | ||
| 4 | + from torch._dynamo.guards import get_verbose_code_parts | ||
| 5 | + from torch._dynamo.source import GlobalStateSource | ||
| 6 | + import torch_npu | ||
| 7 | + | ||
| 8 | + tc = _guards.TracingContext.try_get() | ||
| 9 | + if tc is None: | ||
| 10 | + return False | ||
| 11 | + | ||
| 12 | + captured_level = torch_npu.npu._get_deterministic_level() | ||
| 13 | + | ||
| 14 | + def _create_guard_fn(builder, guard): | ||
| 15 | + code = [f"torch_npu.npu._get_deterministic_level() == {captured_level}"] | ||
| 16 | + | ||
| 17 | + def check_fn(_): | ||
| 18 | + return torch_npu.npu._get_deterministic_level() == captured_level | ||
| 19 | + | ||
| 20 | + builder.guard_manager.root.add_lambda_guard( | ||
| 21 | + check_fn, | ||
| 22 | + get_verbose_code_parts(code, guard), | ||
| 23 | + guard.user_stack, | ||
| 24 | + ) | ||
| 25 | + | ||
| 26 | + tc.guards_context.dynamo_guards.add(GlobalStateSource().make_guard(_create_guard_fn)) | ||
| 27 | + return True | ||
| @@ -1,6 +1,11 @@ | |||
| 1 | -__all__ = ["limit_core_num"] | 1 | +__all__ = ["limit_core_num", "deterministic"] |
| 2 | 2 | ||
| 3 | 3 | ||
| 4 | def limit_core_num(op_aicore_num: int, op_vectorcore_num: int, stream=None): | 4 | def limit_core_num(op_aicore_num: int, op_vectorcore_num: int, stream=None): |
| 5 | from torch_npu.dynamo.npugraph_ex import scope | 5 | from torch_npu.dynamo.npugraph_ex import scope |
| 6 | return scope.limit_core_num(op_aicore_num, op_vectorcore_num, stream=stream) | 6 | return scope.limit_core_num(op_aicore_num, op_vectorcore_num, stream=stream) |
| 7 | + | ||
| 8 | + | ||
| 9 | +def deterministic(level: int): | ||
| 10 | + from torch_npu.dynamo.npugraph_ex import scope | ||
| 11 | + return scope.deterministic(level) | ||
| @@ -150,6 +150,12 @@ class _NpuBackendScope: | |||
| 150 | self._old_env = os.environ.get("TORCHINDUCTOR_NPU_BACKEND") | 150 | self._old_env = os.environ.get("TORCHINDUCTOR_NPU_BACKEND") |
| 151 | os.environ["TORCHINDUCTOR_NPU_BACKEND"] = self.backend | 151 | os.environ["TORCHINDUCTOR_NPU_BACKEND"] = self.backend |
| 152 | register_inductor_npu() | 152 | register_inductor_npu() |
| 153 | + if self.backend == "ascendc": | ||
| 154 | + from torch_npu._inductor.deterministic_cache import ( | ||
| 155 | + patch_npu_deterministic_level_cache_keys, | ||
| 156 | + ) | ||
| 157 | + | ||
| 158 | + patch_npu_deterministic_level_cache_keys() | ||
| 153 | return self | 159 | return self |
| 154 | 160 | ||
| 155 | def __exit__(self, exc_type, exc, tb): | 161 | def __exit__(self, exc_type, exc, tb): |
| @@ -173,6 +179,12 @@ def patch_inductor_wrapper(): | |||
| 173 | def new_call(self, model_, inputs_): | 179 | def new_call(self, model_, inputs_): |
| 174 | backend = _resolve_npu_backend_from_wrapper(self) | 180 | backend = _resolve_npu_backend_from_wrapper(self) |
| 175 | with _NpuBackendScope(backend): | 181 | with _NpuBackendScope(backend): |
| 182 | + if backend == "ascendc": | ||
| 183 | + from torch_npu.dynamo._deterministic_guard import ( | ||
| 184 | + install_npu_deterministic_level_guard, | ||
| 185 | + ) | ||
| 186 | + | ||
| 187 | + install_npu_deterministic_level_guard() | ||
 | |||
| 176 | return src_call(self, model_, inputs_) | 188 | return src_call(self, model_, inputs_) |
| 177 | 189 | ||
| 178 | def new_get_config_copy(self) -> dict[str, Any]: | 190 | def new_get_config_copy(self) -> dict[str, Any]: |
🟠 High Priority
在
patch_inductor_wrapper的new_call中（第 182-187 行），install_npu_deterministic_level_guard()在src_call（即原始_TorchCompileInductorWrapper.__call__）之前调用。install_npu_deterministic_level_guard()内部通过_guards.TracingContext.try_get()获取当前 trace 上下文，并在其上安装 guard。但在调用src_call之前，Dynamo 的 TracingContext 尚未创建，因此try_get()返回None，函数直接返回False，guard 永远不会被安装到任何 trace 中。后果：生产环境中，
torch.compile使用 AscendC 后端时，deterministic_level的 Dynamo guard 机制完全失效——切换 level 不会触发 Dynamo 重新 tracing（虽然 inductor 级别的 cache key patching 仍可能通过不同的 cache key 触发重新编译，但 Dynamo 层面的 re-trace 不会发生）。测试中可工作的原因：测试用自定义
backend函数调用install_npu_deterministic_level_guard()，而该backend是在 Dynamo tracing 过程中（TracingContext 已存在）被回调的，因此 guard 能成功安装。但生产路径（_TorchCompileInductorWrapper.__call__）中调用时机过早。建议：
install_npu_deterministic_level_guard()必须在 TracingContext 存在时调用（即在 Dynamo tracing 过程中）。建议将 guard 安装逻辑移到 Dynamo 的编译回调链路中——例如通过 patching Dynamo 的InstructionTranslator初始化过程、或在 inductor 编译的适当 hook 点（如compile_fx内部）安装 guard。一个可行方向是在_TorchCompileInductorWrapper的编译路径中，找到 TracingContext 已被创建之后的点再调用 guard 安装函数，而非在__call__入口处调用。