已合并
fix: preserve deferred Triton backend loading (v2.9.0) #44735
黄桂军创建于 6 天前
fix: preserve deferred Triton backend loading (v2.9.0) #44735
已合并
共 2 个文件变更+32-41
| @@ -528,10 +528,11 @@ class TorchCompileTriggerTests(unittest.TestCase): | |||
| 528 | """ | 528 | """ |
| 529 | ) | 529 | ) |
| 530 | 530 | ||
| 531 | - # Verify shape handling is installed after selecting the requested NPU backend. | 531 | + # Verify shape handling is installed only after the selected backend scope. |
| 532 | def test_shape_handling_initializes_after_backend_selection(self): | 532 | def test_shape_handling_initializes_after_backend_selection(self): |
| 533 | self.run_in_subprocess( | 533 | self.run_in_subprocess( |
| 534 | """ | 534 | """ |
| 535 | + import os | ||
| 535 | import types | 536 | import types |
| 536 | from unittest import mock | 537 | from unittest import mock |
| 537 | 538 | ||
| @@ -539,11 +540,6 @@ class TorchCompileTriggerTests(unittest.TestCase): | |||
| 539 | import torch_npu | 540 | import torch_npu |
| 540 | from torch_npu.utils import _dynamo | 541 | from torch_npu.utils import _dynamo |
| 541 | 542 | ||
| 542 | - def fake_setup(actual_options): | ||
| 543 | - actual_options = dict(actual_options) | ||
| 544 | - events.append(("setup", actual_options)) | ||
| 545 | - return actual_options["npu_backend"] | ||
| 546 | - | ||
| 547 | options = { | 543 | options = { |
| 548 | "npu_backend": "mlir", | 544 | "npu_backend": "mlir", |
| 549 | "enable_shape_handling": True, | 545 | "enable_shape_handling": True, |
| @@ -551,7 +547,9 @@ class TorchCompileTriggerTests(unittest.TestCase): | |||
| 551 | events = [] | 547 | events = [] |
| 552 | 548 | ||
| 553 | def scope_register(): | 549 | def scope_register(): |
| 554 | - events.append(("scope_register", None)) | 550 | + events.append( |
| 551 | + ("scope_register", os.environ.get("TORCHINDUCTOR_NPU_BACKEND")) | ||
| 552 | + ) | ||
| 555 | 553 | ||
| 556 | fake_inductor = types.SimpleNamespace( | 554 | fake_inductor = types.SimpleNamespace( |
| 557 | patch_shape_handling=lambda: events.append( | 555 | patch_shape_handling=lambda: events.append( |
| @@ -559,7 +557,9 @@ class TorchCompileTriggerTests(unittest.TestCase): | |||
| 559 | ) | 557 | ) |
| 560 | ) | 558 | ) |
| 561 | with mock.patch.object( | 559 | with mock.patch.object( |
| 562 | - _dynamo, "_setup_inductor_for_compile", fake_setup | 560 | + _dynamo, "_lazy_dynamo_setup" |
| 561 | + ), mock.patch.object( | ||
| 562 | + _dynamo, "_lazy_inductor_setup" | ||
| 563 | ), mock.patch.object( | 563 | ), mock.patch.object( |
| 564 | _dynamo, "register_inductor_npu", scope_register | 564 | _dynamo, "register_inductor_npu", scope_register |
| 565 | ), mock.patch.object( | 565 | ), mock.patch.object( |
| @@ -569,11 +569,7 @@ class TorchCompileTriggerTests(unittest.TestCase): | |||
| 569 | 569 | ||
| 570 | assert wrapper.config["npu_backend"] == "mlir" | 570 | assert wrapper.config["npu_backend"] == "mlir" |
| 571 | assert wrapper.config["enable_shape_handling"] is True | 571 | assert wrapper.config["enable_shape_handling"] is True |
| 572 | - assert events == [ | 572 | + assert events == [("scope_register", "mlir")], events |
| 573 | - ("setup", options), | ||
| 574 | - ("shape_handling", None), | ||
| 575 | - ("scope_register", None), | ||
| 576 | - ], events | ||
| 577 | """ | 573 | """ |
| 578 | ) | 574 | ) |
| 579 | 575 | ||
| @@ -723,6 +719,23 @@ class TorchCompileTriggerTests(unittest.TestCase): | |||
| 723 | """ | 719 | """ |
| 724 | ) | 720 | ) |
| 725 | 721 | ||
| 722 | + # Verify creating an Inductor wrapper does not load the backend prematurely. | ||
| 723 | + def test_inductor_backend_load_is_deferred_until_first_call(self): | ||
| 724 | + self.run_in_subprocess( | ||
| 725 | + """ | ||
| 726 | + import sys | ||
| 727 | + import torch | ||
| 728 | + import torch_npu | ||
| 729 | + from torch_npu.utils import _dynamo | ||
| 730 | + | ||
| 731 | + torch.compile(lambda x: x + 1, backend="inductor") | ||
| 732 | + | ||
| 733 | + assert _dynamo._lazy_dynamo_setup.has_run | ||
| 734 | + assert not _dynamo._lazy_inductor_setup.has_run | ||
| 735 | + assert "torch_npu._inductor" not in sys.modules | ||
| 736 | + """ | ||
| 737 | + ) | ||
| 738 | + | ||
| 726 | # Verify lazy setup completes before compile backend lookup. | 739 | # Verify lazy setup completes before compile backend lookup. |
| 727 | def test_compile_triggers_setup_before_backend_lookup(self): | 740 | def test_compile_triggers_setup_before_backend_lookup(self): |
| 728 | self.run_in_subprocess( | 741 | self.run_in_subprocess( |
| @@ -167,6 +167,7 @@ class _NpuBackendScope: | |||
| 167 | try: | 167 | try: |
| 168 | os.environ["TORCHINDUCTOR_NPU_BACKEND"] = self.backend | 168 | os.environ["TORCHINDUCTOR_NPU_BACKEND"] = self.backend |
| 169 | register_inductor_npu() | 169 | register_inductor_npu() |
| 170 | + _lazy_inductor_setup() | ||
| 170 | if self.backend == "ascendc": | 171 | if self.backend == "ascendc": |
| 171 | from torch_npu._inductor.deterministic_cache import ( | 172 | from torch_npu._inductor.deterministic_cache import ( |
| 172 | patch_npu_deterministic_level_cache_keys, | 173 | patch_npu_deterministic_level_cache_keys, |
| @@ -504,10 +505,8 @@ def patch_inductor_wrapper(): | |||
| 504 | if shape_handling_requested: | 505 | if shape_handling_requested: |
| 505 | if getattr(self, "_npu_defer_shape_handling", False): | 506 | if getattr(self, "_npu_defer_shape_handling", False): |
| 506 | self._npu_shape_handling_requested = True | 507 | self._npu_shape_handling_requested = True |
| 507 | - return | 508 | + # Shape handling is installed in new_call, after the selected |
| 508 | - if not is_inductor_npu_initialized(): | 509 | + # backend scope has loaded the matching NPU Inductor backend. |
| 509 | - register_inductor_npu() | ||
| 510 | - torch_npu._inductor.patch_shape_handling() | ||
| 511 | 510 | ||
| 512 | def new_get_config_copy(self) -> dict[str, Any]: | 511 | def new_get_config_copy(self) -> dict[str, Any]: |
| 513 | ori_dict = src_get_config_copy(self) | 512 | ori_dict = src_get_config_copy(self) |
| @@ -545,13 +544,10 @@ def patch_inductor_wrapper(): | |||
| 545 | self._npu_shape_handling_requested = False | 544 | self._npu_shape_handling_requested = False |
| 546 | try: | 545 | try: |
| 547 | src_init(self, mode, options, dynamic) | 546 | src_init(self, mode, options, dynamic) |
| 548 | - shape_handling_requested = self._npu_shape_handling_requested | ||
| 549 | finally: | 547 | finally: |
| 550 | del self._npu_defer_shape_handling | 548 | del self._npu_defer_shape_handling |
| 551 | del self._npu_shape_handling_requested | 549 | del self._npu_shape_handling_requested |
| 552 | - _setup_inductor_for_compile(self.config) | 550 | + _lazy_dynamo_setup() |
| 553 | - if shape_handling_requested: | ||
| 554 | - torch_npu._inductor.patch_shape_handling() | ||
| 555 | backend = _resolve_npu_backend_from_wrapper(self) | 551 | backend = _resolve_npu_backend_from_wrapper(self) |
| 556 | if backend=="mlir": | 552 | if backend=="mlir": |
| 557 | with _NpuBackendScope(backend): | 553 | with _NpuBackendScope(backend): |
| @@ -565,6 +561,8 @@ def patch_inductor_wrapper(): | |||
| 565 | def new_call(self, model_, inputs_): | 561 | def new_call(self, model_, inputs_): |
| 566 | backend = _resolve_npu_backend_from_wrapper(self) | 562 | backend = _resolve_npu_backend_from_wrapper(self) |
| 567 | with _NpuBackendScope(backend): | 563 | with _NpuBackendScope(backend): |
| 564 | + if self.config.get("enable_shape_handling", False): | ||
| 565 | + torch_npu._inductor.patch_shape_handling() | ||
| 568 | if backend == "ascendc": | 566 | if backend == "ascendc": |
| 569 | from torch_npu.dynamo._deterministic_guard import ( | 567 | from torch_npu.dynamo._deterministic_guard import ( |
| 570 | install_npu_deterministic_level_guard, | 568 | install_npu_deterministic_level_guard, |
| @@ -654,26 +652,6 @@ def _lazy_inductor_setup(): | |||
| 654 | _inject_inductor_npu_backend_config() | 652 | _inject_inductor_npu_backend_config() |
| 655 | 653 | ||
| 656 | 654 | ||
| 657 | -def _setup_inductor_for_compile(options=None): | ||
| 658 | - """Initialize the NPU Inductor backend selected for this compile call.""" | ||
| 659 | - _lazy_dynamo_setup() | ||
| 660 | - | ||
| 661 | - option_backend = options.get("npu_backend") if isinstance(options, dict) else None | ||
| 662 | - selected_backend = _resolve_npu_backend(option_backend) | ||
| 663 | - | ||
| 664 | - old_backend = os.environ.get("TORCHINDUCTOR_NPU_BACKEND") | ||
| 665 | - if selected_backend not in (None, "", "default"): | ||
| 666 | - os.environ["TORCHINDUCTOR_NPU_BACKEND"] = selected_backend | ||
| 667 | - try: | ||
| 668 | - _lazy_inductor_setup() | ||
| 669 | - finally: | ||
| 670 | - if old_backend is None: | ||
| 671 | - os.environ.pop("TORCHINDUCTOR_NPU_BACKEND", None) | ||
| 672 | - else: | ||
| 673 | - os.environ["TORCHINDUCTOR_NPU_BACKEND"] = old_backend | ||
| 674 | - return selected_backend | ||
| 675 | - | ||
| 676 | - | ||
| 677 | 655 | ||
| 678 | def install_npugraph_mark_step_trigger(): | 656 | def install_npugraph_mark_step_trigger(): |
| 679 | """Expose the public NPUGraph step API without importing compiler internals.""" | 657 | """Expose the public NPUGraph step API without importing compiler internals.""" |