File size: 4,822 Bytes
2216aae |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 |
"""
This is an experimental new API for PyTorch Distributed. This is actively in development and subject to change or deletion entirely.
This is intended as a proving ground for more flexible and object oriented distributed APIs.
"""
from collections.abc import Generator
from contextlib import contextmanager
from datetime import timedelta
from typing import Protocol, Union
import torch
from torch._C._distributed_c10d import (
_current_process_group,
_set_process_group,
ProcessGroup,
ReduceOp,
Store,
)
from torch.distributed.rendezvous import rendezvous
_BACKENDS: dict[str, "ProcessGroupFactory"] = {}
__all__ = [
"ProcessGroup",
"ReduceOp",
"ProcessGroupFactory",
"register_backend",
"new_group",
"current_process_group",
"process_group",
]
class ProcessGroupFactory(Protocol):
"""Protocol for process group factories."""
def __call__(
self,
store: Store,
rank: int,
world_size: int,
timeout: timedelta,
device: torch.device,
**kwargs: object,
) -> ProcessGroup: ...
def register_backend(name: str, func: ProcessGroupFactory) -> None:
"""
Register a new process group backend.
Args:
name: The name of the backend.
func: The function to create the process group.
"""
if name in _BACKENDS:
raise ValueError(f"Backend {name} already registered")
_BACKENDS[name] = func
def _gloo_factory(
store: Store,
rank: int,
world_size: int,
timeout: timedelta,
device: torch.device,
**kwargs: object,
) -> ProcessGroup:
from torch.distributed import ProcessGroupGloo
assert len(kwargs) == 0, "Gloo backend received unexpected kwargs"
backend_class = ProcessGroupGloo(store, rank, world_size, timeout)
backend_class._set_sequence_number_for_group()
pg = ProcessGroup(store, rank, world_size)
pg._set_default_backend(ProcessGroup.BackendType.GLOO)
# register devices
pg._register_backend(device, ProcessGroup.BackendType.GLOO, backend_class)
pg._register_backend(
torch.device("cpu"), ProcessGroup.BackendType.GLOO, backend_class
)
if torch.cuda.is_available():
pg._register_backend(
torch.device("cuda"), ProcessGroup.BackendType.GLOO, backend_class
)
return pg
def _nccl_factory(
store: Store,
rank: int,
world_size: int,
timeout: timedelta,
device: torch.device,
**kwargs: object,
) -> ProcessGroup:
from torch.distributed import ProcessGroupNCCL
opts = ProcessGroupNCCL.Options()
opts._timeout = timeout
for k, v in kwargs.items():
if not hasattr(opts, k):
raise KeyError(f"Unknown option {k}")
setattr(opts, k, v)
backend_class = ProcessGroupNCCL(store, rank, world_size, opts)
backend_class._set_sequence_number_for_group()
backend_class.eager_connect_single_device(device)
pg = ProcessGroup(store, rank, world_size)
pg._set_default_backend(ProcessGroup.BackendType.NCCL)
pg._register_backend(device, ProcessGroup.BackendType.NCCL, backend_class)
return pg
register_backend("gloo", _gloo_factory)
register_backend("nccl", _nccl_factory)
def new_group(
backend: str,
timeout: timedelta,
device: Union[str, torch.device],
**kwargs: object,
) -> ProcessGroup:
"""
Create a new process group with the given backend and options. This group is
independent and will not be globally registered and thus not usable via the
standard torch.distributed.* APIs.
Args:
backend: The backend to use for the process group.
timeout: The timeout for collective operations.
device: The device to use for the process group.
**kwargs: All remaining arguments are passed to the backend constructor.
See the backend specific documentation for details.
Returns:
A new process group.
"""
if backend not in _BACKENDS:
raise ValueError(f"Backend {backend} not registered")
device = torch.device(device)
store, rank, world_size = next(iter(rendezvous("env://")))
store.set_timeout(timeout)
return _BACKENDS[backend](store, rank, world_size, timeout, device, **kwargs)
def current_process_group() -> ProcessGroup:
"""
Get the current process group. Thread local method.
Returns:
The current process group.
"""
return _current_process_group()
@contextmanager
def process_group(pg: ProcessGroup) -> Generator[None, None, None]:
"""
Context manager for process groups. Thread local method.
Args:
pg: The process group to use.
"""
prev_pg = current_process_group()
_set_process_group(pg)
try:
yield
finally:
_set_process_group(prev_pg)
|