diff --git a/docs/source/distributed.md b/docs/source/distributed.md index 1fac9170d8f3..95820f8244c5 100644 --- a/docs/source/distributed.md +++ b/docs/source/distributed.md @@ -224,6 +224,10 @@ inconsistent 'UUID' assignment across ranks, and to prevent races during initial .. autofunction:: is_torchelastic_launched ``` +```{eval-rst} +.. autofunction:: get_default_backend_for_device +``` + ______________________________________________________________________ Currently three initialization methods are supported: diff --git a/torch/distributed/distributed_c10d.py b/torch/distributed/distributed_c10d.py index 866658515a74..2bc99a51cd64 100644 --- a/torch/distributed/distributed_c10d.py +++ b/torch/distributed/distributed_c10d.py @@ -1379,7 +1379,7 @@ def get_default_backend_for_device(device: Union[str, torch.device]) -> str: Return the default backend for the given device. Args: - Union[str, torch.device]: The device to get the default backend for. + device (Union[str, torch.device]): The device to get the default backend for. Returns: The default backend for the given device as a lower case string. @@ -1584,6 +1584,8 @@ def init_process_group( process must have exclusive access to every GPU it uses, as sharing GPUs between processes can result in deadlock or NCCL invalid usage. ``ucc`` backend is experimental. + Default backend for the device can be queried with + :func:`get_default_backend_for_device`. init_method (str, optional): URL specifying how to initialize the process group. Default is "env://" if no ``init_method`` or ``store`` is specified.