已合并
fix_2_activate_drivers #42512
qiaoyaodan创建于 7月23日
fix_2_activate_drivers #42512
已合并
共 2 个文件变更+50-0
| @@ -422,5 +422,43 @@ class TestTransferToNpu(TestCase): | |||
| 422 | kwargs_output = mock_function(**kwargs_input) | 422 | kwargs_output = mock_function(**kwargs_input) |
| 423 | self.assertEqual(kwargs_output, expected_kwargs_output) | 423 | self.assertEqual(kwargs_output, expected_kwargs_output) |
| 424 | 424 | ||
| 425 | + def test_transfer_to_npu_works_without_triton(self): | ||
| 426 | + # 验证：当前环境未安装 triton / triton_ascend 时， | ||
| 427 | + # _patch_triton_nvidia_driver() 的 try/except 容错能保证 | ||
| 428 | + # transfer_to_npu 整个文件依旧可被 import 且核心功能（cuda -> npu）正常。 | ||
| 429 | + import importlib.util | ||
| 430 | + | ||
| 431 | + # 1. 确认当前环境确实没有 triton 和 triton_ascend | ||
| 432 | + triton_spec = importlib.util.find_spec("triton") | ||
| 433 | + triton_ascend_spec = importlib.util.find_spec("triton_ascend") | ||
| 434 | + if triton_spec is not None or triton_ascend_spec is not None: | ||
| 435 | + import unittest | ||
| 436 | + raise unittest.SkipTest( | ||
| 437 | + "triton or triton_ascend installed, skip" | ||
| 438 | + ) | ||
| 439 | + | ||
| 440 | + # 2. 验证 _patch_triton_nvidia_driver 可被调用且不抛异常 | ||
| 441 | + # （transfer_to_npu 在 import 时已调用过一次，这里再显式调用一次验证幂等性） | ||
| 442 | + transfer_to_npu._patch_triton_nvidia_driver() | ||
| 443 | + | ||
| 444 | + # 3. 验证 transfer_to_npu 的核心 patch 已生效（说明 _init() 完整跑完，没有因 triton 缺失而中断） | ||
| 445 | + # 3.1 cuda.is_available 被替换为 npu.is_available | ||
| 446 | + self.assertTrue(torch.cuda.is_available()) | ||
| 447 | + | ||
| 448 | + # 3.2 tensor.cuda() 实际创建在 npu 上 | ||
| 449 | + t = torch.tensor([1.0, 2.0, 3.0]).cuda() | ||
| 450 | + self.assertEqual(t.device.type, "npu") | ||
| 451 | + | ||
| 452 | + # 3.3 eager 计算正常 | ||
| 453 | + y = t + t | ||
| 454 | + self.assertEqual(y.device.type, "npu") | ||
| 455 | + self.assertTrue(torch.allclose(y, torch.tensor([2.0, 4.0, 6.0], device="npu"))) | ||
| 456 | + | ||
| 457 | + # 3.4 nn.Module.cuda() 转换正常 | ||
| 458 | + linear = torch.nn.Linear(3, 1).cuda() | ||
| 459 | + self.assertEqual(next(linear.parameters()).device.type, "npu") | ||
| 460 | + out = linear(t) | ||
| 461 | + self.assertEqual(out.device.type, "npu") | ||
| 462 | + | ||
| 425 | if __name__ == "__main__": | 463 | if __name__ == "__main__": |
| 426 | run_tests() | 464 | run_tests() |
| @@ -362,6 +362,16 @@ def _patch_cuda(): | |||
| 362 | _apply_patches(patches) | 362 | _apply_patches(patches) |
| 363 | 363 | ||
| 364 | 364 | ||
| 365 | +def _patch_triton_nvidia_driver(): | ||
| 366 | + try: | ||
| 367 | + from triton.backends.nvidia.driver import CudaDriver | ||
| 368 | + # Triton uses is_active() (not is_available()) to determine active drivers. | ||
| 369 | + # See triton/runtime/driver.py: active_drivers = [x.driver for x in backends.values() if x.driver.is_active()] | ||
| 370 | + CudaDriver.is_active = staticmethod(lambda: False) | ||
| 371 | + except (ImportError, AttributeError): | ||
| 372 | + pass | ||
| 373 | + | ||
| 374 | + | ||
| 365 | def _patch_profiler(): | 375 | def _patch_profiler(): |
| 366 | patches = [ | 376 | patches = [ |
| 367 | ['profiler.profile', torch_npu.profiler.profile], | 377 | ['profiler.profile', torch_npu.profiler.profile], |
| @@ -529,6 +539,8 @@ def _init(): | |||
| 529 | 539 | ||
| 530 | _do_wrapper_libraries_func(_load_json_file(config_path)) | 540 | _do_wrapper_libraries_func(_load_json_file(config_path)) |
| 531 | 541 | ||
| 542 | + _patch_triton_nvidia_driver() | ||
| 543 | + | ||
| 532 | setattr(torch.utils._triton, 'has_triton', _patch_has_triton) | 544 | setattr(torch.utils._triton, 'has_triton', _patch_has_triton) |
| 533 | setattr(torch._dynamo.utils, 'has_triton', _patch_has_triton) | 545 | setattr(torch._dynamo.utils, 'has_triton', _patch_has_triton) |
| 534 | setattr(torch._inductor.runtime.autotune_cache, 'has_triton', _patch_has_triton) | 546 | setattr(torch._inductor.runtime.autotune_cache, 'has_triton', _patch_has_triton) |