已合并
Rewrite cuda to npu in backend string of init_process_group #39195
yinglinwei创建于 6月24日
Rewrite cuda to npu in backend string of init_process_group #39195
已合并
共 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_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) |
| @@ -245,13 +245,21 @@ def _wrapper_hccl(fn): | |||
| 245 | if args: | 245 | if args: |
| 246 | args_new = list(args) | 246 | args_new = list(args) |
| 247 | for idx, arg in enumerate(args_new): | 247 | for idx, arg in enumerate(args_new): |
| 248 | - if type(arg) is str and 'nccl' in arg: | 248 | + if type(arg) is str: |
| 249 | - args_new[idx] = arg.replace('nccl', 'hccl') | 249 | + if 'nccl' in arg: |
| 250 | + arg = arg.replace('nccl', 'hccl') | ||
| 251 | + if 'cuda' in arg: | ||
| 252 | + arg = arg.replace('cuda', 'npu') | ||
| 253 | + args_new[idx] = arg | ||
| 250 | args = args_new | 254 | args = args_new |
| 251 | if kwargs: | 255 | if kwargs: |
| 252 | backend = kwargs.get('backend', None) | 256 | backend = kwargs.get('backend', None) |
| 253 | - if type(backend) is str and 'nccl' in backend: | 257 | + if type(backend) is str: |
| 254 | - kwargs['backend'] = backend.replace('nccl', 'hccl') | 258 | + if 'nccl' in backend: |
| 259 | + backend = backend.replace('nccl', 'hccl') | ||
| 260 | + if 'cuda' in backend: | ||
| 261 | + backend = backend.replace('cuda', 'npu') | ||
| 262 | + kwargs['backend'] = backend | ||
| 255 | return fn(*args, **kwargs) | 263 | return fn(*args, **kwargs) |
| 256 | 264 | ||
| 257 | return decorated | 265 | return decorated |
| @@ -339,7 +347,7 @@ def _patch_OverlappingCpuLoader_init_(self, resolve_fun: Callable, stream: Optio | |||
| 339 | 347 | ||
| 340 | 348 | ||
| 341 | def _patch_cuda(): | 349 | def _patch_cuda(): |
| 342 | - patchs = [ | 350 | + patches = [ |
| 343 | ['cuda', torch_npu.npu], ['cuda.amp', torch_npu.npu.amp], | 351 | ['cuda', torch_npu.npu], ['cuda.amp', torch_npu.npu.amp], |
| 344 | ['cuda.random', torch_npu.npu.random], | 352 | ['cuda.random', torch_npu.npu.random], |
| 345 | ['cuda.amp.autocast_mode', torch_npu.npu.amp.autocast_mode], | 353 | ['cuda.amp.autocast_mode', torch_npu.npu.amp.autocast_mode], |
| @@ -348,11 +356,11 @@ def _patch_cuda(): | |||
| 348 | ] | 356 | ] |
| 349 | 357 | ||
| 350 | from torch_npu._init.patches.monkey_patches import _apply_patches | 358 | from torch_npu._init.patches.monkey_patches import _apply_patches |
| 351 | - _apply_patches(patchs) | 359 | + _apply_patches(patches) |
| 352 | 360 | ||
| 353 | 361 | ||
| 354 | def _patch_profiler(): | 362 | def _patch_profiler(): |
| 355 | - patchs = [ | 363 | + patches = [ |
| 356 | ['profiler.profile', torch_npu.profiler.profile], | 364 | ['profiler.profile', torch_npu.profiler.profile], |
| 357 | ['profiler.schedule', torch_npu.profiler.schedule], | 365 | ['profiler.schedule', torch_npu.profiler.schedule], |
| 358 | ['profiler.tensorboard_trace_handler', torch_npu.profiler.tensorboard_trace_handler], | 366 | ['profiler.tensorboard_trace_handler', torch_npu.profiler.tensorboard_trace_handler], |
| @@ -362,7 +370,7 @@ def _patch_profiler(): | |||
| 362 | ] | 370 | ] |
| 363 | 371 | ||
| 364 | from torch_npu._init.patches.monkey_patches import _apply_patches | 372 | from torch_npu._init.patches.monkey_patches import _apply_patches |
| 365 | - _apply_patches(patchs) | 373 | + _apply_patches(patches) |
| 366 | 374 | ||
| 367 | 375 | ||
| 368 | def _warning_fn(msg, rank0=True): | 376 | def _warning_fn(msg, rank0=True): |