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

Support multi-node distributed training with MLU backend (#1266)

Browse Source
This commit is contained in:
黄启元 2023-07-26 10:32:53 +08:00 committed by GitHub
parent 68360e7ce8
commit 78205c3254
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
mmengine/dist/utils.py vendored
View File

@ -101,7 +101,8 @@ def _init_dist_pytorch(backend, init_backend='torch', **kwargs) -> None:
rank = int(os.environ['RANK']) rank = int(os.environ['RANK'])
if is_mlu_available(): if is_mlu_available():
import torch_mlu # noqa: F401 import torch_mlu # noqa: F401
torch.mlu.set_device(rank) local_rank = int(os.environ['LOCAL_RANK'])
torch.mlu.set_device(local_rank)
torch_dist.init_process_group( torch_dist.init_process_group(
backend='cncl', backend='cncl',
rank=rank, rank=rank,