mirrors/mmengine
1
0
Fork 0
mirror of https://github.com/open-mmlab/mmengine.git synced 2025-06-03 21:54:44 +08:00
Code Issues Projects Releases Wiki Activity

[Fix] Support multi-node distributed training with NPU backend (#1459)

Browse Source
This commit is contained in:
lanzeshun 2023-12-26 16:14:45 +08:00 committed by GitHub
parent 671f3bcdf4
commit 8e6fb12b1f
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 3 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
mmengine/dist/utils.py vendored
View File

@ -99,9 +99,10 @@ def _init_dist_pytorch(backend, init_backend='torch', **kwargs) -> None:
**kwargs: keyword arguments are passed to ``init_process_group``. **kwargs: keyword arguments are passed to ``init_process_group``.
""" """
rank = int(os.environ['RANK']) rank = int(os.environ['RANK'])
# LOCAL_RANK is set by `torch.distributed.launch` since PyTorch 1.1
local_rank = int(os.environ['LOCAL_RANK'])
if is_mlu_available(): if is_mlu_available():
import torch_mlu # noqa: F401 import torch_mlu # noqa: F401
local_rank = int(os.environ['LOCAL_RANK'])
torch.mlu.set_device(local_rank) torch.mlu.set_device(local_rank)
torch_dist.init_process_group( torch_dist.init_process_group(
backend='cncl', backend='cncl',
@ -110,15 +111,13 @@ def _init_dist_pytorch(backend, init_backend='torch', **kwargs) -> None:
**kwargs) **kwargs)
elif is_npu_available(): elif is_npu_available():
import torch_npu # noqa: F401 import torch_npu # noqa: F401
torch.npu.set_device(rank) torch.npu.set_device(local_rank)
torch_dist.init_process_group( torch_dist.init_process_group(
backend='hccl', backend='hccl',
rank=rank, rank=rank,
world_size=int(os.environ['WORLD_SIZE']), world_size=int(os.environ['WORLD_SIZE']),
**kwargs) **kwargs)
else: else:
# LOCAL_RANK is set by `torch.distributed.launch` since PyTorch 1.1
local_rank = int(os.environ['LOCAL_RANK'])
torch.cuda.set_device(local_rank) torch.cuda.set_device(local_rank)
if init_backend == 'torch': if init_backend == 'torch':