已合并
[feat] deterministic_level: add graph-mode guard and EX scope API #41762
Lyric创建于 7月15日
[feat] deterministic_level: add graph-mode guard and EX scope API #41762
已合并
共 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() | ||
| @@ -2297,6 +2297,9 @@ | |||
| 2297 | "torch_npu.npu.npugraph_ex.scope.limit_core_num": { | 2297 | "torch_npu.npu.npugraph_ex.scope.limit_core_num": { |
| 2298 | "signature": "(op_aicore_num: int, op_vectorcore_num: int, stream=None)" | 2298 | "signature": "(op_aicore_num: int, op_vectorcore_num: int, stream=None)" |
| 2299 | }, | 2299 | }, |
| 2300 | + "torch_npu.npu.npugraph_ex.scope.deterministic": { | ||
| 2301 | + "signature": "(level: int)" | ||
| 2302 | + }, | ||
| 2300 | "torch_npu.npu.npugraph_ex.compile_fx": { | 2303 | "torch_npu.npu.npugraph_ex.compile_fx": { |
| 2301 | "signature": "(gm, example_inputs=None, options=None)" | 2304 | "signature": "(gm, example_inputs=None, options=None)" |
| 2302 | }, | 2305 | }, |
| @@ -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) | ||
| @@ -153,6 +153,12 @@ class _NpuBackendScope: | |||
| 153 | self._old_env = os.environ.get("TORCHINDUCTOR_NPU_BACKEND") | 153 | self._old_env = os.environ.get("TORCHINDUCTOR_NPU_BACKEND") |
| 154 | os.environ["TORCHINDUCTOR_NPU_BACKEND"] = self.backend | 154 | os.environ["TORCHINDUCTOR_NPU_BACKEND"] = self.backend |
| 155 | register_inductor_npu() | 155 | register_inductor_npu() |
| 156 | + if self.backend == "ascendc": | ||
| 157 | + from torch_npu._inductor.deterministic_cache import ( | ||
| 158 | + patch_npu_deterministic_level_cache_keys, | ||
| 159 | + ) | ||
| 160 | + | ||
| 161 | + patch_npu_deterministic_level_cache_keys() | ||
| 156 | return self | 162 | return self |
| 157 | 163 | ||
| 158 | def __exit__(self, exc_type, exc, tb): | 164 | def __exit__(self, exc_type, exc, tb): |
| @@ -207,6 +213,12 @@ def patch_inductor_wrapper(): | |||
| 207 | def new_call(self, model_, inputs_): | 213 | def new_call(self, model_, inputs_): |
| 208 | backend = _resolve_npu_backend_from_wrapper(self) | 214 | backend = _resolve_npu_backend_from_wrapper(self) |
| 209 | with _NpuBackendScope(backend): | 215 | with _NpuBackendScope(backend): |
| 216 | + if backend == "ascendc": | ||
| 217 | + from torch_npu.dynamo._deterministic_guard import ( | ||
| 218 | + install_npu_deterministic_level_guard, | ||
| 219 | + ) | ||
| 220 | + | ||
| 221 | + install_npu_deterministic_level_guard() | ||
| 210 | return src_call(self, model_, inputs_) | 222 | return src_call(self, model_, inputs_) |
| 211 | 223 | ||
| 212 | _TorchCompileInductorWrapper.__call__ = new_call | 224 | _TorchCompileInductorWrapper.__call__ = new_call |