已合并
【inductor】fix thread pool when inductor is imported #30477
kkjocker创建于 2月3日
【inductor】fix thread pool when inductor is imported #30477
已合并
共 3 个文件变更+7-1
| @@ -11,7 +11,7 @@ class TestAssertScalar(TestUtils): | |||
| 11 | torch.ops.aten._assert_scalar(condition, "sumq should be non-negative") | 11 | torch.ops.aten._assert_scalar(condition, "sumq should be non-negative") |
| 12 | return x * 2 | 12 | return x * 2 |
| 13 | 13 | ||
| 14 | - @parametrize('shape', [(1024, 32), (8, 16, 256)]) | 14 | + @parametrize('shape', [(512, 64), (4, 32, 128)]) |
| 15 | @parametrize('dtype', ['float32', 'float16']) | 15 | @parametrize('dtype', ['float32', 'float16']) |
| 16 | def test_assert_scalar_pass(self, shape, dtype): | 16 | def test_assert_scalar_pass(self, shape, dtype): |
| 17 | input_element = self._generate_tensor(shape, dtype, floatPOSIFLAG=1) | 17 | input_element = self._generate_tensor(shape, dtype, floatPOSIFLAG=1) |
| @@ -13,6 +13,7 @@ from functools import wraps | |||
| 13 | # Disable autoloading before running 'import torch' to avoid circular dependencies | 13 | # Disable autoloading before running 'import torch' to avoid circular dependencies |
| 14 | ORG_AUTOLOAD = os.getenv("TORCH_DEVICE_BACKEND_AUTOLOAD", "1") | 14 | ORG_AUTOLOAD = os.getenv("TORCH_DEVICE_BACKEND_AUTOLOAD", "1") |
| 15 | os.environ["TORCH_DEVICE_BACKEND_AUTOLOAD"] = "0" | 15 | os.environ["TORCH_DEVICE_BACKEND_AUTOLOAD"] = "0" |
| 16 | +os.environ["TORCH_WARM_POOL"] = "0" | ||
| 16 | 17 | ||
| 17 | import torch | 18 | import torch |
| 18 | from torch.distributed.fsdp import sharded_grad_scaler | 19 | from torch.distributed.fsdp import sharded_grad_scaler |
| @@ -1,4 +1,9 @@ | |||
| 1 | import os | 1 | import os |
| 2 | +ORG_AUTOLOAD = os.getenv("TORCH_DEVICE_BACKEND_AUTOLOAD", "1") | ||
| 3 | +os.environ["TORCH_DEVICE_BACKEND_AUTOLOAD"] = "0" | ||
| 4 | +from torch._inductor.async_compile import AsyncCompile | ||
| 5 | +AsyncCompile.warm_pool() | ||
| 6 | +os.environ["TORCH_DEVICE_BACKEND_AUTOLOAD"] = ORG_AUTOLOAD | ||
| 2 | 7 | ||
| 3 | if os.getenv('TORCHINDUCTOR_NPU_BACKEND', 'default') == 'mlir': | 8 | if os.getenv('TORCHINDUCTOR_NPU_BACKEND', 'default') == 'mlir': |
| 4 | try: | 9 | try: |