已合并
[v2.10.0] restore deferred Triton backend loading #44718
黄桂军创建于 18 天前
[v2.10.0] restore deferred Triton backend loading #44718
已合并
共 2 个文件变更+36-41
| @@ -473,10 +473,11 @@ class TorchCompileTriggerTests(unittest.TestCase): | |||
| 473 | """ | 473 | """ |
| 474 | ) | 474 | ) |
| 475 | 475 | ||
| 476 | - # Verify shape handling is installed after selecting the requested NPU backend. | 476 | + # Verify shape handling is installed only after the selected backend scope. |
| 477 | def test_shape_handling_initializes_after_backend_selection(self): | 477 | def test_shape_handling_initializes_after_backend_selection(self): |
| 478 | self.run_in_subprocess( | 478 | self.run_in_subprocess( |
| 479 | """ | 479 | """ |
| 480 | + import os | ||
| 480 | import types | 481 | import types |
| 481 | from unittest import mock | 482 | from unittest import mock |
| 482 | 483 | ||
| @@ -484,11 +485,6 @@ class TorchCompileTriggerTests(unittest.TestCase): | |||
| 484 | import torch_npu | 485 | import torch_npu |
| 485 | from torch_npu.utils import _dynamo | 486 | from torch_npu.utils import _dynamo |
| 486 | 487 | ||
| 487 | - def fake_setup(actual_options): | ||
| 488 | - actual_options = dict(actual_options) | ||
| 489 | - events.append(("setup", actual_options)) | ||
| 490 | - return actual_options["npu_backend"] | ||
| 491 | - | ||
| 492 | options = { | 488 | options = { |
| 493 | "npu_backend": "mlir", | 489 | "npu_backend": "mlir", |
| 494 | "enable_shape_handling": True, | 490 | "enable_shape_handling": True, |
| @@ -496,7 +492,9 @@ class TorchCompileTriggerTests(unittest.TestCase): | |||
| 496 | events = [] | 492 | events = [] |
| 497 | 493 | ||
| 498 | def scope_register(): | 494 | def scope_register(): |
| 499 | - events.append(("scope_register", None)) | 495 | + events.append( |
| 496 | + ("scope_register", os.environ.get("TORCHINDUCTOR_NPU_BACKEND")) | ||
| 497 | + ) | ||
| 500 | 498 | ||
| 501 | fake_inductor = types.SimpleNamespace( | 499 | fake_inductor = types.SimpleNamespace( |
| 502 | patch_shape_handling=lambda: events.append( | 500 | patch_shape_handling=lambda: events.append( |
| @@ -504,7 +502,9 @@ class TorchCompileTriggerTests(unittest.TestCase): | |||
| 504 | ) | 502 | ) |
| 505 | ) | 503 | ) |
| 506 | with mock.patch.object( | 504 | with mock.patch.object( |
| 507 | - _dynamo, "_setup_inductor_for_compile", fake_setup | 505 | + _dynamo, "_lazy_dynamo_setup", lambda: None |
| 506 | + ), mock.patch.object( | ||
| 507 | + _dynamo, "_lazy_inductor_setup", lambda: None | ||
| 508 | ), mock.patch.object( | 508 | ), mock.patch.object( |
| 509 | _dynamo, "register_inductor_npu", scope_register | 509 | _dynamo, "register_inductor_npu", scope_register |
| 510 | ), mock.patch.object( | 510 | ), mock.patch.object( |
| @@ -514,11 +514,7 @@ class TorchCompileTriggerTests(unittest.TestCase): | |||
| 514 | 514 | ||
| 515 | assert wrapper.config["npu_backend"] == "mlir" | 515 | assert wrapper.config["npu_backend"] == "mlir" |
| 516 | assert wrapper.config["enable_shape_handling"] is True | 516 | assert wrapper.config["enable_shape_handling"] is True |
| 517 | - assert events == [ | 517 | + assert events == [("scope_register", "mlir")], events |
| 518 | - ("setup", options), | ||
| 519 | - ("shape_handling", None), | ||
| 520 | - ("scope_register", None), | ||
| 521 | - ], events | ||
| 522 | """ | 518 | """ |
| 523 | ) | 519 | ) |
| 524 | 520 | ||
| @@ -668,6 +664,27 @@ class TorchCompileTriggerTests(unittest.TestCase): | |||
| 668 | """ | 664 | """ |
| 669 | ) | 665 | ) |
| 670 | 666 | ||
| 667 | + # Creating an Inductor wrapper must not load the Triton backend yet. | ||
| 668 | + def test_inductor_backend_load_is_deferred_until_first_call(self): | ||
| 669 | + self.run_in_subprocess( | ||
| 670 | + """ | ||
| 671 | + import sys | ||
| 672 | + import torch | ||
| 673 | + import torch_npu | ||
| 674 | + from torch_npu.utils import _dynamo | ||
| 675 | + | ||
| 676 | + torch.compile( | ||
| 677 | + lambda x: x + 1, | ||
| 678 | + backend="inductor", | ||
| 679 | + options={"enable_shape_handling": True}, | ||
| 680 | + ) | ||
| 681 | + | ||
| 682 | + assert _dynamo._lazy_dynamo_setup.has_run | ||
| 683 | + assert not _dynamo._lazy_inductor_setup.has_run | ||
| 684 | + assert "torch_npu._inductor" not in sys.modules | ||
| 685 | + """ | ||
| 686 | + ) | ||
| 687 | + | ||
| 671 | # Verify lazy setup completes before compile backend lookup. | 688 | # Verify lazy setup completes before compile backend lookup. |
| 672 | def test_compile_triggers_setup_before_backend_lookup(self): | 689 | def test_compile_triggers_setup_before_backend_lookup(self): |
| 673 | self.run_in_subprocess( | 690 | self.run_in_subprocess( |
| @@ -168,6 +168,7 @@ class _NpuBackendScope: | |||
| 168 | try: | 168 | try: |
| 169 | os.environ["TORCHINDUCTOR_NPU_BACKEND"] = self.backend | 169 | os.environ["TORCHINDUCTOR_NPU_BACKEND"] = self.backend |
| 170 | register_inductor_npu() | 170 | register_inductor_npu() |
| 171 | + _lazy_inductor_setup() | ||
| 171 | if self.backend == "ascendc": | 172 | if self.backend == "ascendc": |
| 172 | from torch_npu._inductor.deterministic_cache import ( | 173 | from torch_npu._inductor.deterministic_cache import ( |
| 173 | patch_npu_deterministic_level_cache_keys, | 174 | patch_npu_deterministic_level_cache_keys, |
| @@ -209,10 +210,8 @@ def patch_inductor_wrapper(): | |||
| 209 | if shape_handling_requested: | 210 | if shape_handling_requested: |
| 210 | if getattr(self, "_npu_defer_shape_handling", False): | 211 | if getattr(self, "_npu_defer_shape_handling", False): |
| 211 | self._npu_shape_handling_requested = True | 212 | self._npu_shape_handling_requested = True |
| 212 | - return | 213 | + # Shape handling is installed in new_call, after the selected |
| 213 | - if not is_inductor_npu_initialized(): | 214 | + # backend scope has loaded the matching NPU Inductor backend. |
| 214 | - register_inductor_npu() | ||
| 215 | - torch_npu._inductor.patch_shape_handling() | ||
| 216 | 215 | ||
| 217 | def new_get_config_copy(self) -> dict[str, Any]: | 216 | def new_get_config_copy(self) -> dict[str, Any]: |
| 218 | ori_dict = src_get_config_copy(self) | 217 | ori_dict = src_get_config_copy(self) |
| @@ -250,13 +249,10 @@ def patch_inductor_wrapper(): | |||
| 250 | self._npu_shape_handling_requested = False | 249 | self._npu_shape_handling_requested = False |
| 251 | try: | 250 | try: |
| 252 | src_init(self, mode, options, dynamic) | 251 | src_init(self, mode, options, dynamic) |
| 253 | - shape_handling_requested = self._npu_shape_handling_requested | ||
| 254 | finally: | 252 | finally: |
| 255 | del self._npu_defer_shape_handling | 253 | del self._npu_defer_shape_handling |
| 256 | del self._npu_shape_handling_requested | 254 | del self._npu_shape_handling_requested |
| 257 | - _setup_inductor_for_compile(self.config) | 255 | + _lazy_dynamo_setup() |
| 258 | - if shape_handling_requested: | ||
| 259 | - torch_npu._inductor.patch_shape_handling() | ||
| 260 | backend = _resolve_npu_backend_from_wrapper(self) | 256 | backend = _resolve_npu_backend_from_wrapper(self) |
| 261 | if backend=="mlir": | 257 | if backend=="mlir": |
| 262 | with _NpuBackendScope(backend): | 258 | with _NpuBackendScope(backend): |
| @@ -270,6 +266,8 @@ def patch_inductor_wrapper(): | |||
| 270 | def new_call(self, model_, inputs_): | 266 | def new_call(self, model_, inputs_): |
| 271 | backend = _resolve_npu_backend_from_wrapper(self) | 267 | backend = _resolve_npu_backend_from_wrapper(self) |
| 272 | with _NpuBackendScope(backend): | 268 | with _NpuBackendScope(backend): |
| 269 | + if self.config.get("enable_shape_handling", False): | ||
| 270 | + torch_npu._inductor.patch_shape_handling() | ||
| 273 | if backend == "ascendc": | 271 | if backend == "ascendc": |
| 274 | from torch_npu.dynamo._deterministic_guard import ( | 272 | from torch_npu.dynamo._deterministic_guard import ( |
| 275 | install_npu_deterministic_level_guard, | 273 | install_npu_deterministic_level_guard, |
| @@ -710,26 +708,6 @@ def _lazy_inductor_setup(): | |||
| 710 | _inject_inductor_npu_backend_config() | 708 | _inject_inductor_npu_backend_config() |
| 711 | 709 | ||
| 712 | 710 | ||
| 713 | -def _setup_inductor_for_compile(options=None): | ||
| 714 | - """Initialize the NPU Inductor backend selected for this compile call.""" | ||
| 715 | - _lazy_dynamo_setup() | ||
| 716 | - | ||
| 717 | - option_backend = options.get("npu_backend") if isinstance(options, dict) else None | ||
| 718 | - selected_backend = _resolve_npu_backend(option_backend) | ||
| 719 | - | ||
| 720 | - old_backend = os.environ.get("TORCHINDUCTOR_NPU_BACKEND") | ||
| 721 | - if selected_backend not in (None, "", "default"): | ||
| 722 | - os.environ["TORCHINDUCTOR_NPU_BACKEND"] = selected_backend | ||
| 723 | - try: | ||
| 724 | - _lazy_inductor_setup() | ||
| 725 | - finally: | ||
| 726 | - if old_backend is None: | ||
| 727 | - os.environ.pop("TORCHINDUCTOR_NPU_BACKEND", None) | ||
| 728 | - else: | ||
| 729 | - os.environ["TORCHINDUCTOR_NPU_BACKEND"] = old_backend | ||
| 730 | - return selected_backend | ||
| 731 | - | ||
| 732 | - | ||
| 733 | 711 | ||
| 734 | def install_npugraph_mark_step_trigger(): | 712 | def install_npugraph_mark_step_trigger(): |
| 735 | """Expose the public NPUGraph step API without importing compiler internals.""" | 713 | """Expose the public NPUGraph step API without importing compiler internals.""" |