已合并
Rewrite cuda to npu in backend string of init_process_group #39244
Rewrite cuda to npu in backend string of init_process_group #39244
已合并
yinglinwei创建于 6月25日
2 个文件变更+22-14
@@ -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_PORT370 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
atomgit-bot
atomgit-botatomgit-bot6月25日

🔵 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。可以引入一个局部变量标记是否发生替换,或在条件内直接赋值。

likedislike
yinglinwei
6月25日 评论:
248 args = args_new252 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
atomgit-botatomgit-bot
atomgit-botatomgit-bot6月25日

🔵 Low Priority

_wrapper_hccl 的 kwargs 处理路径中,对于任意字符串类型的 backend 值,新代码始终执行 kwargs['backend'] = backend(第260行),即使该值未发生任何替换(如 backend='gloo')。旧代码仅在 backend 包含 'nccl' 时才写入 kwargs['backend']。虽然结果值相同(只是原地重新赋值),但语义上有不必要的字典写入操作。这不是缺陷,但表明代码结构引入了无实际效果的赋值。

建议:可选优化:将 kwargs['backend'] = backend 的赋值移到 if 块内部,仅在实际发生替换时才写入。例如使用一个局部标志位追踪是否发生了替换,或保持旧代码风格:仅在满足替换条件时更新字典。

likedislike
yinglinwei
6月25日 评论:
atomgit-botatomgit-bot6月25日

🔵 Low Priority

_wrapper_hccl 的 args 处理循环中(第246-251行),对每个字符串类型参数依次检查 'nccl' in arg'cuda' in arg。当参数同时包含两个子串时(如 'cuda:nccl'),两次替换均会被执行,结果正确。但当参数为 'nccl' 时,先 replace('nccl', 'hccl') 得到 'hccl',再检查 'cuda' in 'hccl' 为 False;当参数为 'cuda:gloo'(无 nccl)时,'nccl' 检查为 False 跳过,'cuda' 检查为 True 执行替换得到 'npu:gloo'。两个 if 语句是顺序独立执行的(非 elif),逻辑正确且符合预期。kwargs 路径(第255-260行)同理。无缺陷。

建议:逻辑正确,无需修改。两个独立的 if(非 elif)是正确的设计,因为部分 backend 字符串可能同时包含 'nccl' 和 'cuda'。

likedislike
yinglinwei
6月25日 评论:
253 return fn(*args, **kwargs)261 return fn(*args, **kwargs)
254 262 
255 return decorated263 return decorated
@@ -337,7 +345,7 @@ def _patch_OverlappingCpuLoader_init_(self, resolve_fun: Callable, stream: Optio
337 345 
338 346 
339def _patch_cuda():347def _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_patches356 from torch_npu._init.patches.monkey_patches import _apply_patches
349- _apply_patches(patchs)357+ _apply_patches(patches)
350 358 
351 359 
352def _patch_profiler():360def _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_patches370 from torch_npu._init.patches.monkey_patches import _apply_patches
363- _apply_patches(patchs)371+ _apply_patches(patches)
atomgit-bot
atomgit-botatomgit-bot6月25日

🔵 Low Priority

_patch_cuda 函数(第348行)和 _patch_profiler 函数(第361行)中,变量名从 patchs 修正为 patches。这是一个纯拼写修正,行为无变化。

建议:无需修改,该修正正确且安全。

likedislike
yinglinwei
6月25日 评论:
364 372 
365 373 
366def _warning_fn(msg, rank0=True):374def _warning_fn(msg, rank0=True):