Skip to content
Snippets Groups Projects
Unverified Commit 8864bd88 authored by Austin Welch's avatar Austin Welch Committed by GitHub
Browse files

[Feats]: Add smddp dist backend option (#579)


* Add smddp dist backend option

* [Dev]: Upgrade pre commit hooks (#576)

* Upgrade the versions of pre-commit-hooks

* update zh-cn.yaml

* [Docs] Fix the docstring of model sub-package (#573)

* [Doc]: Update config.md (#562)

* Update config.md

* Update config.md

* [Doc] delete the error comment  in docs (#514)

Co-authored-by: default avatarZaida Zhou <58739961+zhouzaida@users.noreply.github.com>
Co-authored-by: default avatarZhengfei-0311 <78833899+Zhengfei-0311@users.noreply.github.com>
Co-authored-by: default avatarvansin <msnode@163.com>
parent e73c4bf1
No related branches found
No related tags found
No related merge requests found
......@@ -94,6 +94,15 @@ def _init_dist_mpi(backend, **kwargs) -> None:
'nccl', 'gloo' and 'mpi'. Defaults to 'nccl'.
**kwargs: keyword arguments are passed to ``init_process_group``.
"""
if backend == 'smddp':
try:
import smdistributed.dataparallel.torch.torch_smddp # noqa: F401
except ModuleNotFoundError as e:
raise ModuleNotFoundError(
'Please use an Amazon SageMaker DLC to access smdistributed: '
'https://github.com/aws/deep-learning-containers/blob/master'
'/available_images.md#sagemaker-framework-containers'
'-sm-support-only') from e
local_rank = int(os.environ['OMPI_COMM_WORLD_LOCAL_RANK'])
torch.cuda.set_device(local_rank)
if 'MASTER_PORT' not in os.environ:
......@@ -433,6 +442,8 @@ def get_comm_device(group: Optional[ProcessGroup] = None) -> torch.device:
elif backend == 'cncl':
import torch_mlu # noqa: F401
return torch.device('mlu', torch.mlu.current_device())
elif backend == 'smddp':
return torch.device('cuda', torch.cuda.current_device())
else:
# GLOO and MPI backends use cpu device by default
return torch.device('cpu')
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment