| @@ -359,7 +359,7 @@ class TestTransferToNpu(TestCase): | |||
| 359 | 359 | ||
| 360 | def test_torch_utils_cpp_extension_include_paths(self): | 360 | def test_torch_utils_cpp_extension_include_paths(self): |
| 361 | torch.utils.cpp_extension.include_paths(device_type='cuda') | 361 | torch.utils.cpp_extension.include_paths(device_type='cuda') |
| 362 | - | 362 | + |
| 363 | def test_init_process_group(self): | 363 | def test_init_process_group(self): |
| 364 | MASTER_ADDR = "127.0.0.1" | 364 | MASTER_ADDR = "127.0.0.1" |
| 365 | MASTER_PORT = "29500" | 365 | MASTER_PORT = "29500" |
| @@ -370,19 +370,19 @@ class TestTransferToNpu(TestCase): | |||
| 370 | os.environ['MASTER_PORT'] = MASTER_PORT | 370 | os.environ['MASTER_PORT'] = MASTER_PORT |
| 371 | os.environ['RANK'] = str(RANK) | 371 | os.environ['RANK'] = str(RANK) |
| 372 | os.environ['WORLD_SIZE'] = str(WORLD_SIZE) | 372 | os.environ['WORLD_SIZE'] = str(WORLD_SIZE) |
| 373 | - | 373 | + |
| 374 | try: | 374 | try: |
| 375 | 375 | ||
| 376 | torch.distributed.init_process_group( | 376 | torch.distributed.init_process_group( |
| 377 | - backend='nccl', | 377 | + backend='cuda:nccl', |
| 378 | init_method=f"tcp://{MASTER_ADDR}:{MASTER_PORT}", | 378 | init_method=f"tcp://{MASTER_ADDR}:{MASTER_PORT}", |
| 379 | world_size=WORLD_SIZE, | 379 | world_size=WORLD_SIZE, |
| 380 | rank=RANK, | 380 | rank=RANK, |
| 381 | device_id=torch.device(f"cuda:{RANK}") | 381 | device_id=torch.device(f"cuda:{RANK}") |
| 382 | ) | 382 | ) |
| 383 | - self.assertEqual(torch.distributed.get_backend(), 'hccl') | 383 | + self.assertEqual(torch.distributed.get_backend(), 'npu:hccl') |
| 384 | torch.distributed.barrier() | 384 | torch.distributed.barrier() |
| 385 | - | 385 | + |
| 386 | finally: | 386 | finally: |
| 387 | if torch.distributed.is_initialized(): | 387 | if torch.distributed.is_initialized(): |
| 388 | torch.distributed.destroy_process_group() | 388 | torch.distributed.destroy_process_group() |
| @@ -398,7 +398,7 @@ class TestTransferToNpu(TestCase): | |||
| 398 | 398 | ||
| 399 | def test_host_empty_cache_is_patched(self): | 399 | def test_host_empty_cache_is_patched(self): |
| 400 | self.assertEqual(torch._C._host_emptyCache, torch_npu._C._npu_hostEmptyCache) | 400 | self.assertEqual(torch._C._host_emptyCache, torch_npu._C._npu_hostEmptyCache) |
| 401 | - | 401 | + |
| 402 | def test_update_cuda_default_generators(self): | 402 | def test_update_cuda_default_generators(self): |
| 403 | torch.randn(1).npu() | 403 | torch.randn(1).npu() |
| 404 | self.assertEqual(torch.cuda.default_generators, torch_npu.npu.default_generators) | 404 | self.assertEqual(torch.cuda.default_generators, torch_npu.npu.default_generators) |
| @@ -243,13 +243,21 @@ def _wrapper_hccl(fn): | |||
| 243 | if args: | 243 | if args: |
| 244 | args_new = list(args) | 244 | args_new = list(args) |
| 245 | for idx, arg in enumerate(args_new): | 245 | for idx, arg in enumerate(args_new): |
| 246 | - if type(arg) == str and 'nccl' in arg: | 246 | + if type(arg) is str: |
| 247 | - args_new[idx] = arg.replace('nccl', 'hccl') | 247 | + if 'nccl' in arg: |
| 248 | + arg = arg.replace('nccl', 'hccl') | ||
| 249 | + if 'cuda' in arg: | ||
| 250 | + arg = arg.replace('cuda', 'npu') | ||
| 251 | + args_new[idx] = arg | ||
| 248 | args = args_new | 252 | args = args_new |
| 249 | if kwargs: | 253 | if kwargs: |
| 250 | backend = kwargs.get('backend', None) | 254 | backend = kwargs.get('backend', None) |
| 251 | - if type(backend) == str and 'nccl' in backend: | 255 | + if type(backend) is str: |
| 252 | - kwargs['backend'] = backend.replace('nccl', 'hccl') | 256 | + if 'nccl' in backend: |
| 257 | + backend = backend.replace('nccl', 'hccl') | ||
| 258 | + if 'cuda' in backend: | ||
| 259 | + backend = backend.replace('cuda', 'npu') | ||
| 260 | + kwargs['backend'] = backend | ||
🔵 Low Priority 在 建议:可选优化:将 ![]() ![]() 🔵 Low Priority
建议:逻辑正确,无需修改。两个独立的 if(非 elif)是正确的设计,因为部分 backend 字符串可能同时包含 'nccl' 和 'cuda'。 ![]() ![]() | |||
| 253 | return fn(*args, **kwargs) | 261 | return fn(*args, **kwargs) |
| 254 | 262 | ||
| 255 | return decorated | 263 | return decorated |
| @@ -337,7 +345,7 @@ def _patch_OverlappingCpuLoader_init_(self, resolve_fun: Callable, stream: Optio | |||
| 337 | 345 | ||
| 338 | 346 | ||
| 339 | def _patch_cuda(): | 347 | def _patch_cuda(): |
| 340 | - patchs = [ | 348 | + patches = [ |
| 341 | ['cuda', torch_npu.npu], ['cuda.amp', torch_npu.npu.amp], | 349 | ['cuda', torch_npu.npu], ['cuda.amp', torch_npu.npu.amp], |
| 342 | ['cuda.random', torch_npu.npu.random], | 350 | ['cuda.random', torch_npu.npu.random], |
| 343 | ['cuda.amp.autocast_mode', torch_npu.npu.amp.autocast_mode], | 351 | ['cuda.amp.autocast_mode', torch_npu.npu.amp.autocast_mode], |
| @@ -346,11 +354,11 @@ def _patch_cuda(): | |||
| 346 | ] | 354 | ] |
| 347 | 355 | ||
| 348 | from torch_npu._init.patches.monkey_patches import _apply_patches | 356 | from torch_npu._init.patches.monkey_patches import _apply_patches |
| 349 | - _apply_patches(patchs) | 357 | + _apply_patches(patches) |
| 350 | 358 | ||
| 351 | 359 | ||
| 352 | def _patch_profiler(): | 360 | def _patch_profiler(): |
| 353 | - patchs = [ | 361 | + patches = [ |
| 354 | ['profiler.profile', torch_npu.profiler.profile], | 362 | ['profiler.profile', torch_npu.profiler.profile], |
| 355 | ['profiler.schedule', torch_npu.profiler.schedule], | 363 | ['profiler.schedule', torch_npu.profiler.schedule], |
| 356 | ['profiler.tensorboard_trace_handler', torch_npu.profiler.tensorboard_trace_handler], | 364 | ['profiler.tensorboard_trace_handler', torch_npu.profiler.tensorboard_trace_handler], |
| @@ -360,7 +368,7 @@ def _patch_profiler(): | |||
| 360 | ] | 368 | ] |
| 361 | 369 | ||
| 362 | from torch_npu._init.patches.monkey_patches import _apply_patches | 370 | from torch_npu._init.patches.monkey_patches import _apply_patches |
| 363 | - _apply_patches(patchs) | 371 | + _apply_patches(patches) |
🔵 Low Priority 在 建议:无需修改,该修正正确且安全。 ![]() ![]() | |||
| 364 | 372 | ||
| 365 | 373 | ||
| 366 | def _warning_fn(msg, rank0=True): | 374 | def _warning_fn(msg, rank0=True): |


🔵 Low Priority
_wrapper_hccl的 args 处理路径中,对于所有字符串类型的参数,新代码无条件执行args_new[idx] = arg(第251行),即使该值未发生任何替换。旧代码仅在包含'nccl'时才赋值(此时arg已通过初始list(args)拷贝了原值)。虽然字符串不可变且最终值一致(原地重新赋值相同对象),但这与 kwargs 路径一样引入了无实际效果的操作。建议仅在发生替换时才写入,保持代码意图清晰。建议:可选优化:仅在
arg实际发生替换(即'nccl' in arg或'cuda' in arg为 True)时才执行args_new[idx] = arg。可以引入一个局部变量标记是否发生替换,或在条件内直接赋值。