mmengine.dist.dist 源代码
# Copyright (c) OpenMMLab. All rights reserved.
import os.path as osp
import pickle
import shutil
import tempfile
from collections import OrderedDict
from typing import Any, Dict, Generator, List, Optional, Tuple, Union
import numpy as np
import torch
from torch import Tensor
from torch import distributed as torch_dist
from torch._utils import (_flatten_dense_tensors, _take_tensors,
_unflatten_dense_tensors)
from torch.distributed import ProcessGroup
from itertools import zip_longest, chain
import mmengine
from .utils import (get_world_size, get_rank, get_backend, get_dist_info,
get_default_group, barrier, get_data_device,
get_comm_device, cast_data_device)
from mmengine.utils import digit_version
from mmengine.utils.dl_utils import TORCH_VERSION
from mmengine.device import is_npu_available
def _get_reduce_op(name: str) -> torch_dist.ReduceOp:
op_mappings = {
'sum': torch_dist.ReduceOp.SUM,
'product': torch_dist.ReduceOp.PRODUCT,
'min': torch_dist.ReduceOp.MIN,
'max': torch_dist.ReduceOp.MAX,
'band': torch_dist.ReduceOp.BAND,
'bor': torch_dist.ReduceOp.BOR,
'bxor': torch_dist.ReduceOp.BXOR,
}
if name.lower() not in op_mappings:
raise ValueError(
f'reduce op should be one of {op_mappings.keys()}, bug got {name}')
return op_mappings[name.lower()]
[文档]def all_reduce(data: Tensor,
op: str = 'sum',
group: Optional[ProcessGroup] = None) -> None:
"""Reduces the tensor data across all machines in such a way that all get
the final result.
After the call ``data`` is going to be bitwise identical in all
processes.
Note:
Calling ``all_reduce`` in non-distributed environment does nothing.
Args:
data (Tensor): Input and output of the collective. The function
operates in-place.
op (str): Operation to reduce data. Defaults to 'sum'. Optional values
are 'sum', 'mean' and 'produce', 'min', 'max', 'band', 'bor' and
'bxor'.
group (ProcessGroup, optional): The process group to work on. If None,
the default process group will be used. Defaults to None.
Examples:
>>> import torch
>>> import mmengine.dist as dist
>>> # non-distributed environment
>>> data = torch.arange(2, dtype=torch.int64)
>>> dist.all_reduce(data)
>>> data
tensor([0, 1])
>>> # distributed environment
>>> # We have 2 process groups, 2 ranks.
>>> data = torch.arange(2, dtype=torch.int64) + 1 + 2 * rank
>>> data
tensor([1, 2]) # Rank 0
tensor([3, 4]) # Rank 1
>>> dist.all_reduce(data, op=dist.ReduceOp.SUM)
>>> data
tensor([4, 6]) # Rank 0
tensor([4, 6]) # Rank 1
"""
world_size = get_world_size(group)
if world_size > 1:
if group is None:
group = get_default_group()
input_device = get_data_device(data)
backend_device = get_comm_device(group)
data_on_device = cast_data_device(data, backend_device)
# pytorch does not support 'mean' operation so we fall back to support
# it with 'sum' operation.
if op.lower() == 'mean':
torch_dist.all_reduce(data_on_device, _get_reduce_op('sum'), group)
# use true_divide to handle torch1.6.0 throws an RuntimeError when
# the type of `data_on_device` is int64
data_on_device = torch.true_divide(data_on_device, world_size)
else:
torch_dist.all_reduce(data_on_device, _get_reduce_op(op), group)
cast_data_device(data_on_device, input_device, out=data)
[文档]def all_gather(data: Tensor,
group: Optional[ProcessGroup] = None) -> List[Tensor]:
"""Gather data from the whole group in a list.
Note:
Calling ``all_gather`` in non-distributed environment does nothing
and just returns a list containing :attr:`data` itself.
Note:
Unlike PyTorch ``torch.distributed.all_gather``, :meth:`all_gather` in
MMEngine does not pass in an empty list ``gather_list`` and returns
the ``gather_list`` directly, which is more convenient. The difference
between their interfaces is as below:
- MMEngine: all_gather(data, group) -> gather_list
- PyTorch: all_gather(gather_list, data, group) -> None
Args:
data (Tensor): Tensor to be gathered.
group (ProcessGroup, optional): The process group to work on. If None,
the default process group will be used. Defaults to None.
Returns:
list[Tensor]: Return a list containing data from the whole group if
in distributed environment, otherwise a list only containing
:attr:`data` itself.
Examples:
>>> import torch
>>> import mmengine.dist as dist
>>> # non-distributed environment
>>> data = torch.arange(2, dtype=torch.int64)
>>> data
tensor([0, 1])
>>> output = dist.all_gather(data)
>>> output
[tensor([0, 1])]
>>> # distributed environment
>>> # We have 2 process groups, 2 ranks.
>>> data = torch.arange(2, dtype=torch.int64) + 1 + 2 * rank
>>> data
tensor([1, 2]) # Rank 0
tensor([3, 4]) # Rank 1
>>> output = dist.all_gather(data)
>>> output
[tensor([1, 2]), tensor([3, 4])] # Rank 0
[tensor([1, 2]), tensor([3, 4])] # Rank 1
"""
world_size = get_world_size(group)
if world_size == 1:
return [data]
if group is None:
group = get_default_group()
input_device = get_data_device(data)
backend_device = get_comm_device(group)
data_on_device = cast_data_device(data, backend_device)
gather_list = [
torch.empty_like(data, device=backend_device)
for _ in range(world_size)
]
torch_dist.all_gather(gather_list, data_on_device, group)
return cast_data_device(gather_list, input_device) # type: ignore
[文档]def gather(data: Tensor,
dst: int = 0,
group: Optional[ProcessGroup] = None) -> List[Optional[Tensor]]:
"""Gather data from the whole group to ``dst`` process.
Note:
Calling ``gather`` in non-distributed environment dose nothing
and just returns a list containing :attr:`data` itself.
Note:
``NCCL`` backend does not support ``gather``.
Note:
Unlike PyTorch ``torch.distributed.gather``, :meth:`gather` in
MMEngine does not pass in an empty list ``gather_list`` and returns
the ``gather_list`` directly, which is more convenient. The difference
between their interfaces is as below:
- MMEngine: gather(data, dst, group) -> gather_list
- PyTorch: gather(data, gather_list, dst, group) -> None
Args:
data (Tensor): Tensor to be gathered. CUDA tensor is not supported.
dst (int): Destination rank. Defaults to 0.
group (ProcessGroup, optional): The process group to work on. If None,
the default process group will be used. Defaults to None.
Returns:
list[Tensor]: ``dst`` process will get a list of tensor gathering from
the whole group. Other process will get a empty list. If in
non-distributed environment, just return a list containing
:attr:`data` itself.
Examples:
>>> import torch
>>> import mmengine.dist as dist
>>> # non-distributed environment
>>> data = torch.arange(2, dtype=torch.int64)
>>> data
tensor([0, 1])
>>> output = dist.gather(data)
>>> output
[tensor([0, 1])]
>>> # distributed environment
>>> # We have 2 process groups, 2 ranks.
>>> data = torch.arange(2, dtype=torch.int64) + 1 + 2 * rank
>>> data
tensor([1, 2]) # Rank 0
tensor([3, 4]) # Rank 1
>>> output = dist.gather(data)
>>> output
[tensor([1, 2]), tensor([3, 4])] # Rank 0
[] # Rank 1
"""
world_size = get_world_size(group)
if world_size == 1:
return [data]
if group is None:
group = get_default_group()
input_device = get_data_device(data)
backend_device = get_comm_device(group)
if get_rank(group) == dst:
gather_list = [
torch.empty_like(data, device=backend_device)
for _ in range(world_size)
]
else:
gather_list = []
torch_dist.gather(data, gather_list, dst, group)
if get_rank(group) == dst:
return cast_data_device(gather_list, input_device) # type: ignore
else:
return gather_list
[文档]def broadcast(data: Tensor,
src: int = 0,
group: Optional[ProcessGroup] = None) -> None:
"""Broadcast the data from ``src`` process to the whole group.
``data`` must have the same number of elements in all processes
participating in the collective.
Note:
Calling ``broadcast`` in non-distributed environment does nothing.
Args:
data (Tensor): Data to be sent if ``src`` is the rank of current
process, and data to be used to save received data otherwise.
src (int): Source rank. Defaults to 0.
group (ProcessGroup, optional): The process group to work on. If None,
the default process group will be used. Defaults to None.
Examples:
>>> import torch
>>> import mmengine.dist as dist
>>> # non-distributed environment
>>> data = torch.arange(2, dtype=torch.int64)
>>> data
tensor([0, 1])
>>> dist.broadcast(data)
>>> data
tensor([0, 1])
>>> # distributed environment
>>> # We have 2 process groups, 2 ranks.
>>> data = torch.arange(2, dtype=torch.int64) + 1 + 2 * rank
>>> data
tensor([1, 2]) # Rank 0
tensor([3, 4]) # Rank 1
>>> dist.broadcast(data)
>>> data
tensor([1, 2]) # Rank 0
tensor([1, 2]) # Rank 1
"""
if get_world_size(group) > 1:
if group is None:
group = get_default_group()
input_device = get_data_device(data)
backend_device = get_comm_device(group)
data_on_device = cast_data_device(data, backend_device)
# broadcast requires tensor is contiguous
data_on_device = data_on_device.contiguous() # type: ignore
torch_dist.broadcast(data_on_device, src, group)
if get_rank(group) != src:
cast_data_device(data_on_device, input_device, data)
[文档]def sync_random_seed(group: Optional[ProcessGroup] = None) -> int:
"""Synchronize a random seed to all processes.
In distributed sampling, different ranks should sample non-overlapped
data in the dataset. Therefore, this function is used to make sure that
each rank shuffles the data indices in the same order based
on the same seed. Then different ranks could use different indices
to select non-overlapped data from the same data list.
Args:
group (ProcessGroup, optional): The process group to work on. If None,
the default process group will be used. Defaults to None.
Returns:
int: Random seed.
Examples:
>>> import torch
>>> import mmengine.dist as dist
>>> # non-distributed environment
>>> seed = dist.sync_random_seed()
>>> seed # which a random number
587791752
>>> distributed environment
>>> # We have 2 process groups, 2 ranks.
>>> seed = dist.sync_random_seed()
>>> seed
587791752 # Rank 0
587791752 # Rank 1
"""
seed = np.random.randint(2**31)
if get_world_size(group) == 1:
return seed
if group is None:
group = get_default_group()
backend_device = get_comm_device(group)
if get_rank(group) == 0:
random_num = torch.tensor(seed, dtype=torch.int32).to(backend_device)
else:
random_num = torch.tensor(0, dtype=torch.int32).to(backend_device)
torch_dist.broadcast(random_num, src=0, group=group)
return random_num.item()
def _object_to_tensor(obj: Any) -> Tuple[Tensor, Tensor]:
"""Serialize picklable python object to tensor."""
byte_storage = torch.ByteStorage.from_buffer(pickle.dumps(obj))
# Do not replace `torch.ByteTensor` or `torch.LongTensor` with torch.tensor
# and specifying dtype. Otherwise, it will cause 100X slowdown.
# See: https://github.com/pytorch/pytorch/issues/65696
byte_tensor = torch.ByteTensor(byte_storage)
local_size = torch.LongTensor([byte_tensor.numel()])
return byte_tensor, local_size
def _tensor_to_object(tensor: Tensor, tensor_size: int) -> Any:
"""Deserialize tensor to picklable python object."""
buf = tensor.cpu().numpy().tobytes()[:tensor_size]
return pickle.loads(buf)
def _broadcast_object_list(object_list: List[Any],
src: int = 0,
group: Optional[ProcessGroup] = None) -> None:
"""Broadcast picklable objects in ``object_list`` to the whole group.
Similar to :func:`broadcast`, but Python objects can be passed in. Note
that all objects in ``object_list`` must be picklable in order to be
broadcasted.
"""
if torch_dist.distributed_c10d._rank_not_in_group(group):
return
my_rank = get_rank()
# Serialize object_list elements to tensors on src rank.
if my_rank == src:
tensor_list, size_list = zip(
*[_object_to_tensor(obj) for obj in object_list])
object_sizes_tensor = torch.cat(size_list)
else:
object_sizes_tensor = torch.empty(len(object_list), dtype=torch.long)
# Current device selection.
# To preserve backwards compatibility, ``device`` is ``None`` by default.
# in which case we run current logic of device selection, i.e.
# ``current_device`` is CUDA if backend is NCCL otherwise CPU device. In
# the case it is not ``None`` we move the size and object tensors to be
# broadcasted to this device.
group_backend = get_backend(group)
is_nccl_backend = group_backend == torch_dist.Backend.NCCL
current_device = torch.device('cpu')
is_hccl_backend = group_backend == 'hccl'
is_cncl_backend = group_backend == 'cncl'
is_mccl_backend = group_backend == 'mccl'
if is_hccl_backend:
current_device = torch.device('npu', torch.npu.current_device())
object_sizes_tensor = object_sizes_tensor.to(current_device)
elif is_cncl_backend:
current_device = torch.device('mlu', torch.mlu.current_device())
object_sizes_tensor = object_sizes_tensor.to(current_device)
elif is_mccl_backend:
current_device = torch.device('musa', torch.musa.current_device())
object_sizes_tensor = object_sizes_tensor.to(current_device)
elif is_nccl_backend:
# See note about using torch.cuda.current_device() here in
# docstring. We cannot simply use my_rank since rank == device is
# not necessarily true.
current_device = torch.device('cuda', torch.cuda.current_device())
object_sizes_tensor = object_sizes_tensor.to(current_device)
# Broadcast object sizes
torch_dist.broadcast(object_sizes_tensor, src=src, group=group)
# Concatenate and broadcast serialized object tensors
if my_rank == src:
object_tensor = torch.cat(tensor_list)
else:
object_tensor = torch.empty(
torch.sum(object_sizes_tensor).int().item(),
dtype=torch.uint8,
)
if is_nccl_backend or is_hccl_backend or is_cncl_backend:
object_tensor = object_tensor.to(current_device)
torch_dist.broadcast(object_tensor, src=src, group=group)
# Deserialize objects using their stored sizes.
offset = 0
if my_rank != src:
for i, obj_size in enumerate(object_sizes_tensor):
obj_view = object_tensor[offset:offset + obj_size]
obj_view = obj_view.type(torch.uint8)
if obj_view.device != torch.device('cpu'):
obj_view = obj_view.cpu()
offset += obj_size
object_list[i] = _tensor_to_object(obj_view, obj_size)
[文档]def broadcast_object_list(data: List[Any],
src: int = 0,
group: Optional[ProcessGroup] = None) -> None:
"""Broadcasts picklable objects in ``object_list`` to the whole group.
Similar to :func:`broadcast`, but Python objects can be passed in. Note
that all objects in ``object_list`` must be picklable in order to be
broadcasted.
Note:
Calling ``broadcast_object_list`` in non-distributed environment does
nothing.
Args:
data (List[Any]): List of input objects to broadcast.
Each object must be picklable. Only objects on the ``src`` rank
will be broadcast, but each rank must provide lists of equal sizes.
src (int): Source rank from which to broadcast ``object_list``.
group: (ProcessGroup, optional): The process group to work on. If None,
the default process group will be used. Default is ``None``.
device (``torch.device``, optional): If not None, the objects are
serialized and converted to tensors which are moved to the
``device`` before broadcasting. Default is ``None``.
Note:
For NCCL-based process groups, internal tensor representations of
objects must be moved to the GPU device before communication starts.
In this case, the used device is given by
``torch.cuda.current_device()`` and it is the user's responsibility to
ensure that this is correctly set so that each rank has an individual
GPU, via ``torch.cuda.set_device()``.
Examples:
>>> import torch
>>> import mmengine.dist as dist
>>> # non-distributed environment
>>> data = ['foo', 12, {1: 2}]
>>> dist.broadcast_object_list(data)
>>> data
['foo', 12, {1: 2}]
>>> # distributed environment
>>> # We have 2 process groups, 2 ranks.
>>> if dist.get_rank() == 0:
>>> # Assumes world_size of 3.
>>> data = ["foo", 12, {1: 2}] # any picklable object
>>> else:
>>> data = [None, None, None]
>>> dist.broadcast_object_list(data)
>>> data
["foo", 12, {1: 2}] # Rank 0
["foo", 12, {1: 2}] # Rank 1
"""
assert isinstance(data, list)
if get_world_size(group) > 1:
if group is None:
group = get_default_group()
if digit_version(TORCH_VERSION) >= digit_version(
'1.8.0') and not is_npu_available():
torch_dist.broadcast_object_list(data, src, group)
else:
_broadcast_object_list(data, src, group)
[文档]def all_reduce_dict(data: Dict[str, Tensor],
op: str = 'sum',
group: Optional[ProcessGroup] = None) -> None:
"""Reduces the dict across all machines in such a way that all get the
final result.
The code is modified from https://github.com/Megvii-
BaseDetection/YOLOX/blob/main/yolox/utils/allreduce_norm.py.
Args:
data (dict[str, Tensor]): Data to be reduced.
op (str): Operation to reduce data. Defaults to 'sum'. Optional values
are 'sum', 'mean' and 'produce', 'min', 'max', 'band', 'bor' and
'bxor'.
group (ProcessGroup, optional): The process group to work on. If None,
the default process group will be used. Defaults to None.
Examples:
>>> import torch
>>> import mmengine.dist as dist
>>> # non-distributed environment
>>> data = {
'key1': torch.arange(2, dtype=torch.int64),
'key2': torch.arange(3, dtype=torch.int64)
}
>>> dist.all_reduce_dict(data)
>>> data
{'key1': tensor([0, 1]), 'key2': tensor([0, 1, 2])}
>>> # distributed environment
>>> # We have 2 process groups, 2 ranks.
>>> data = {
'key1': torch.arange(2, dtype=torch.int64),
'key2': torch.arange(3, dtype=torch.int64)
}
>>> dist.all_reduce_dict(data)
>>> data
{'key1': tensor([0, 2]), 'key2': tensor([0, 2, 4])} # Rank 0
{'key1': tensor([0, 2]), 'key2': tensor([0, 2, 4])} # Rank 1
"""
assert isinstance(data, dict)
world_size = get_world_size(group)
if world_size > 1:
if group is None:
group = get_default_group()
# ensure keys are consistent across processes
keys = sorted(data.keys())
tensor_shapes = [data[k].shape for k in keys]
tensor_sizes = [data[k].numel() for k in keys]
if digit_version(TORCH_VERSION) == digit_version('1.5.0'):
# `torch.cat` in torch1.5 can not concatenate different types so
# we fallback to convert them all to float type.
flatten_tensor = torch.cat(
[data[k].flatten().float() for k in keys])
else:
flatten_tensor = torch.cat([data[k].flatten() for k in keys])
all_reduce(flatten_tensor, op=op, group=group)
split_tensors = [
x.reshape(shape) for x, shape in zip(
torch.split(flatten_tensor, tensor_sizes), tensor_shapes)
]
for k, v in zip(keys, split_tensors):
data[k] = v
def _all_gather_object(object_list: List[Any],
obj: Any,
group: Optional[ProcessGroup] = None) -> None:
"""Gather picklable objects from the whole group into a list.
Similar to :func:`all_gather`, but Python objects can be passed in.
Note that the object must be picklable in order to be gathered.
Args:
object_list (list[Any]): Output list. It should be correctly sized as
the size of the group for this collective and will contain the
output.
object (Any): Pickable Python object to be broadcast from current
process.
group (ProcessGroup, optional): The process group to work on. If None,
the default process group will be used. Defaults to None.
Returns:
None. If the calling rank is part of this group, the output of the
collective will be populated into the input ``object_list``. If the
calling rank is not part of the group, the passed in ``object_list``
will be unmodified.
"""
if torch_dist.distributed_c10d._rank_not_in_group(group):
return
input_tensor, local_size = _object_to_tensor(obj)
group_backend = get_backend(group)
current_device = torch.device('cpu')
is_nccl_backend = group_backend == torch_dist.Backend.NCCL
is_mccl_backend = group_backend == 'mccl'
if is_nccl_backend:
# See note about using torch.cuda.current_device() here in docstring.
# We cannot simply use my_rank since rank == device is not necessarily
# true.
current_device = torch.device('cuda', torch.cuda.current_device())
input_tensor = input_tensor.to(current_device)
local_size = local_size.to(current_device)
elif is_mccl_backend:
# See note about using torch.musa.current_device() here in docstring.
# We cannot simply use my_rank since rank == device is not necessarily
# true.
current_device = torch.device('musa', torch.musa.current_device())
input_tensor = input_tensor.to(current_device)
local_size = local_size.to(current_device)
# Gather all local sizes. This is so that we can find the max size, and
# index until the correct size when deserializing the tensors.
group_size = get_world_size(group=group)
object_sizes_tensor = torch.zeros(
group_size, dtype=torch.long, device=current_device)
object_size_list = [
object_sizes_tensor[i].unsqueeze(dim=0) for i in range(group_size)
]
# Allgather tensor sizes
torch_dist.all_gather(object_size_list, local_size, group=group)
max_object_size = int(max(object_size_list).item())
# Resize tensor to max size across all ranks.
input_tensor.resize_(max_object_size)
coalesced_output_tensor = torch.empty(
max_object_size * group_size, dtype=torch.uint8, device=current_device)
# Output tensors are nonoverlapping views of coalesced_output_tensor
output_tensors = [
coalesced_output_tensor[max_object_size * i:max_object_size * (i + 1)]
for i in range(group_size)
]
torch_dist.all_gather(output_tensors, input_tensor, group=group)
# Deserialize outputs back to object.
for i, tensor in enumerate(output_tensors):
tensor = tensor.type(torch.uint8)
if tensor.device != torch.device('cpu'):
tensor = tensor.cpu()
tensor_size = object_size_list[i]
object_list[i] = _tensor_to_object(tensor, tensor_size)
[文档]def all_gather_object(data: Any,
group: Optional[ProcessGroup] = None) -> List[Any]:
"""Gather picklable objects from the whole group into a list. Similar to
:func:`all_gather`, but Python objects can be passed in. Note that the
object must be picklable in order to be gathered.
Note:
Calling ``all_gather_object`` in non-distributed environment does
nothing and just returns a list containing :attr:`data` itself.
Note:
Unlike PyTorch ``torch.distributed.all_gather_object``,
:meth:`all_gather_object` in MMEngine does not pass in an empty list
``gather_list`` and returns the ``gather_list`` directly, which is
more convenient. The difference between their interfaces is as below:
- MMEngine: all_gather_object(data, group) -> gather_list
- PyTorch: all_gather_object(gather_list, data, group) -> None
Args:
data (Any): Pickable Python object to be broadcast from current
process.
group (ProcessGroup, optional): The process group to work on. If None,
the default process group will be used. Defaults to None.
Returns:
list[Tensor]: Return a list containing data from the whole group if
in distributed environment, otherwise a list only containing
:attr:`data` itself.
Note:
For NCCL-based process groups, internal tensor representations
of objects must be moved to the GPU device before communication starts.
In this case, the used device is given by
``torch.cuda.current_device()`` and it is the user's responsibility to
ensure that this is correctly set so that each rank has an individual
GPU, via ``torch.cuda.set_device()``.
Examples:
>>> import torch
>>> import mmengine.dist as dist
>>> # non-distributed environment
>>> data = ['foo', 12, {1: 2}] # any picklable object
>>> gather_objects = dist.all_gather_object(data[dist.get_rank()])
>>> output
['foo']
>>> # distributed environment
>>> # We have 3 process groups, 3 ranks.
>>> output = dist.all_gather_object(data[dist.get_rank()])
>>> output
['foo', 12, {1: 2}] # Rank 0
['foo', 12, {1: 2}] # Rank 1
['foo', 12, {1: 2}] # Rank 2
"""
world_size = get_world_size(group)
if world_size == 1:
return [data]
if group is None:
group = get_default_group()
gather_list = [None] * world_size
if digit_version(TORCH_VERSION) >= digit_version('1.8.0'):
torch_dist.all_gather_object(gather_list, data, group)
else:
_all_gather_object(gather_list, data, group)
return gather_list
def _validate_output_list_for_rank(my_rank: int, dst: int,
gather_list: Optional[list]) -> None:
"""Validate whether ``gather_list`` is None in non-dst ranks."""
if dst == my_rank:
if not gather_list:
raise ValueError(
'Argument ``gather_list`` must be specified on destination '
'rank.')
elif gather_list:
raise ValueError('Argument ``gather_list`` must NOT be specified '
'on non-destination ranks.')
def _gather_object(obj: Any,
object_gather_list=None,
dst: int = 0,
group: Optional[ProcessGroup] = None) -> None:
"""Gathers picklable objects from the whole group in a single process.
Similar to :func:`gather`, but Python objects can be passed in. Note that
the object must be picklable in order to be gathered.
Args:
obj (Any): Input object. Must be picklable.
object_gather_list (list[Any], optional): Output list. On the ``dst``
rank, it should be correctly sized as the size of the group for
this collective and will contain the output. Must be ``None`` on
non-dst ranks. Defaults to None.
dst (int): Destination rank. Defaults to 0.
group: (ProcessGroup, optional): The process group to work on. If None,
the default process group will be used. Defaults to None.
"""
if torch_dist.distributed_c10d._rank_not_in_group(group):
return
# Ensure object_gather_list is specified appopriately.
my_rank = get_rank()
_validate_output_list_for_rank(my_rank, dst, object_gather_list)
input_tensor, local_size = _object_to_tensor(obj)
group_backend = get_backend(group)
current_device = torch.device('cpu')
is_nccl_backend = group_backend == torch_dist.Backend.NCCL
is_mccl_backend = group_backend == 'mccl'
if is_nccl_backend:
current_device = torch.device('cuda', torch.cuda.current_device())
input_tensor = input_tensor.to(current_device)
local_size = local_size.to(current_device)
elif is_mccl_backend:
current_device = torch.device('musa', torch.musa.current_device())
input_tensor = input_tensor.to(current_device)
local_size = local_size.to(current_device)
# Gather all local sizes. This is so that we can find the max size, and
# index until the correct size when deserializing the tensors.
group_size = get_world_size(group=group)
object_sizes_tensor = torch.zeros(
group_size, dtype=torch.long, device=current_device)
object_size_list = [
object_sizes_tensor[i].unsqueeze(dim=0) for i in range(group_size)
]
# Allgather tensor sizes. An all-gather is needed here despite this being a
# gather, since each rank needs to broadcast a tensor of the same (maximal)
# size.
torch_dist.all_gather(object_size_list, local_size, group=group)
max_object_size = int(max(object_size_list).item())
# Resize tensor to max size across all ranks.
input_tensor.resize_(max_object_size)
# Avoid populating output tensors if the result won't be gathered on this
# rank.
if my_rank == dst:
coalesced_output_tensor = torch.empty(
max_object_size * group_size,
dtype=torch.uint8,
device=current_device)
# Output tensors are nonoverlapping views of coalesced_output_tensor
output_tensors = [
coalesced_output_tensor[max_object_size * i:max_object_size *
(i + 1)] for i in range(group_size)
]
# All ranks call gather with equal-sized tensors.
torch_dist.gather(
input_tensor,
gather_list=output_tensors if my_rank == dst else None,
dst=dst,
group=group,
)
if my_rank != dst:
return
for i, tensor in enumerate(output_tensors):
tensor = tensor.type(torch.uint8)
tensor_size = object_size_list[i]
object_gather_list[i] = _tensor_to_object(tensor, tensor_size)
[文档]def gather_object(data: Any,
dst: int = 0,
group: Optional[ProcessGroup] = None) -> Optional[List[Any]]:
"""Gathers picklable objects from the whole group in a single process.
Similar to :func:`gather`, but Python objects can be passed in. Note that
the object must be picklable in order to be gathered.
Note:
``NCCL backend`` does not support ``gather_object``.
Note:
Unlike PyTorch ``torch.distributed.gather_object``,
:meth:`gather_object` in MMEngine does not pass in an empty list
``gather_list`` and returns the ``gather_list`` directly, which is
more convenient. The difference between their interfaces is as below:
- MMEngine: gather_object(data, dst, group) -> gather_list
- PyTorch: gather_object(data, gather_list, data, group) -> None
Args:
data (Any): Input object. Must be picklable.
dst (int): Destination rank. Defaults to 0.
group: (ProcessGroup, optional): The process group to work on. If None,
the default process group will be used. Defaults to None.
Returns:
list[Any]. On the ``dst`` rank, return ``gather_list`` which contains
the output of the collective.
Examples:
>>> import torch
>>> import mmengine.dist as dist
>>> # non-distributed environment
>>> data = ['foo', 12, {1: 2}] # any picklable object
>>> gather_objects = dist.gather_object(data[dist.get_rank()])
>>> output
['foo']
>>> # distributed environment
>>> # We have 3 process groups, 3 ranks.
>>> dist.gather_object(gather_objects[dist.get_rank()], dst=0)
>>> output
['foo', 12, {1: 2}] # Rank 0
None # Rank 1
None # Rank 2
"""
world_size = get_world_size(group)
if world_size == 1:
return [data]
if group is None:
group = get_default_group()
gather_list = [None] * world_size if get_rank(group) == dst else None
if digit_version(TORCH_VERSION) >= digit_version('1.8.0'):
torch_dist.gather_object(data, gather_list, dst, group)
else:
_gather_object(data, gather_list, dst, group)
return gather_list
[文档]def collect_results(results: list,
size: int,
device: str = 'cpu',
tmpdir: Optional[str] = None) -> Optional[list]:
"""Collected results in distributed environments.
Args:
results (list[object]): Result list containing result parts to be
collected. Each item of ``result_part`` should be a picklable
object.
size (int): Size of the results, commonly equal to length of
the results.
device (str): Device name. Optional values are 'cpu', 'gpu' or 'npu'.
tmpdir (str | None): Temporal directory for collected results to
store. If set to None, it will create a temporal directory for it.
``tmpdir`` should be None when device is 'gpu' or 'npu'.
Defaults to None.
Returns:
list or None: The collected results.
Examples:
>>> # distributed environment
>>> # We have 2 process groups, 2 ranks.
>>> import mmengine.dist as dist
>>> if dist.get_rank() == 0:
data = ['foo', {1: 2}]
else:
data = [24, {'a': 'b'}]
>>> size = 4
>>> output = dist.collect_results(data, size, device='cpu')
>>> output
['foo', 24, {1: 2}, {'a': 'b'}] # rank 0
None # rank 1
"""
if device not in ['gpu', 'cpu', 'npu']:
raise NotImplementedError(
f"device must be 'cpu' , 'gpu' or 'npu', but got {device}")
if device == 'gpu' or device == 'npu':
assert tmpdir is None, f'tmpdir should be None when device is {device}'
return _collect_results_device(results, size)
else:
return collect_results_cpu(results, size, tmpdir)
[文档]def collect_results_cpu(result_part: list,
size: int,
tmpdir: Optional[str] = None) -> Optional[list]:
"""Collect results under cpu mode.
On cpu mode, this function will save the results on different gpus to
``tmpdir`` and collect them by the rank 0 worker.
Args:
result_part (list): Result list containing result parts
to be collected. Each item of ``result_part`` should be a picklable
object.
size (int): Size of the results, commonly equal to length of
the results.
tmpdir (str | None): Temporal directory for collected results to
store. If set to None, it will create a random temporal directory
for it. Defaults to None.
Returns:
list or None: The collected results.
Examples:
>>> # distributed environment
>>> # We have 2 process groups, 2 ranks.
>>> import mmengine.dist as dist
>>> if dist.get_rank() == 0:
data = ['foo', {1: 2}]
else:
data = [24, {'a': 'b'}]
>>> size = 4
>>> output = dist.collect_results_cpu(data, size)
>>> output
['foo', 24, {1: 2}, {'a': 'b'}] # rank 0
None # rank 1
"""
rank, world_size = get_dist_info()
if world_size == 1:
return result_part[:size]
# create a tmp dir if it is not specified
if tmpdir is None:
MAX_LEN = 512
# 32 is whitespace
dir_tensor = torch.full((MAX_LEN, ), 32, dtype=torch.uint8)
if rank == 0:
mmengine.mkdir_or_exist('.dist_test')
tmpdir = tempfile.mkdtemp(dir='.dist_test')
tmpdir = torch.tensor(
bytearray(tmpdir.encode()), dtype=torch.uint8)
dir_tensor[:len(tmpdir)] = tmpdir
broadcast(dir_tensor, 0)
tmpdir = dir_tensor.numpy().tobytes().decode().rstrip()
else:
mmengine.mkdir_or_exist(tmpdir)
# dump the part result to the dir
with open(osp.join(tmpdir, f'part_{rank}.pkl'), 'wb') as f: # type: ignore
pickle.dump(result_part, f, protocol=2)
barrier()
# collect all parts
if rank != 0:
return None
else:
# load results of all parts from tmp dir
part_list = []
for i in range(world_size):
path = osp.join(tmpdir, f'part_{i}.pkl') # type: ignore
if not osp.exists(path):
raise FileNotFoundError(
f'{tmpdir} is not an shared directory for '
f'rank {i}, please make sure {tmpdir} is a shared '
'directory for all ranks!')
with open(path, 'rb') as f:
part_list.append(pickle.load(f))
# sort the results
ordered_results = []
zipped_results = zip_longest(*part_list)
ordered_results = [
i for i in chain.from_iterable(zipped_results) if i is not None
]
# the dataloader may pad some samples
ordered_results = ordered_results[:size]
# remove tmp dir
shutil.rmtree(tmpdir) # type: ignore
return ordered_results
def _collect_results_device(result_part: list, size: int) -> Optional[list]:
"""Collect results under gpu or npu mode."""
rank, world_size = get_dist_info()
if world_size == 1:
return result_part[:size]
# gather all result part. Note that NCCL does not support gather so use
# all_gather_object instead.
part_list = all_gather_object(result_part)
if rank == 0:
# sort the results
ordered_results = []
zipped_results = zip_longest(*part_list)
ordered_results = [
i for i in chain.from_iterable(zipped_results) if i is not None
]
# the dataloader may pad some samples
ordered_results = ordered_results[:size]
return ordered_results
else:
return None
[文档]def collect_results_gpu(result_part: list, size: int) -> Optional[list]:
"""Collect results under gpu mode.
On gpu mode, this function will encode results to gpu tensors and use gpu
communication for results collection.
Args:
result_part (list[object]): Result list containing result parts
to be collected. Each item of ``result_part`` should be a picklable
object.
size (int): Size of the results, commonly equal to length of
the results.
Returns:
list or None: The collected results.
Examples:
>>> # distributed environment
>>> # We have 2 process groups, 2 ranks.
>>> import mmengine.dist as dist
>>> if dist.get_rank() == 0:
data = ['foo', {1: 2}]
else:
data = [24, {'a': 'b'}]
>>> size = 4
>>> output = dist.collect_results_gpu(data, size)
>>> output
['foo', 24, {1: 2}, {'a': 'b'}] # rank 0
None # rank 1
"""
return _collect_results_device(result_part, size)
def _all_reduce_coalesced(tensors: List[torch.Tensor],
bucket_size_mb: int = -1,
op: str = 'sum',
group: Optional[ProcessGroup] = None) -> None:
"""All-reduce a sequence of tensors as a whole.
Args:
tensors (List[torch.Tensor]): A sequence of tensors to be
all-reduced.
bucket_size_mb (int): The limit of each chunk in megabytes
for grouping tensors into chunks. Defaults to -1.
op (str): Operation to reduce data. Defaults to 'sum'. Optional values
are 'sum', 'mean' and 'produce', 'min', 'max', 'band', 'bor' and
'bxor'.
group (ProcessGroup, optional): The process group to work on. If None,
the default process group will be used. Defaults to None.
"""
if bucket_size_mb > 0:
bucket_size_bytes = bucket_size_mb * 1024 * 1024
buckets = _take_tensors(tensors, bucket_size_bytes)
else:
buckets = OrderedDict()
for tensor in tensors:
tp = tensor.type()
if tp not in buckets:
buckets[tp] = []
buckets[tp].append(tensor)
buckets = buckets.values()
for bucket in buckets:
flat_tensors = _flatten_dense_tensors(bucket)
all_reduce(flat_tensors, op=op, group=group)
for tensor, synced in zip(
bucket, _unflatten_dense_tensors(flat_tensors, bucket)):
tensor.copy_(synced)
[文档]def all_reduce_params(params: Union[List, Generator[torch.Tensor, None, None]],
coalesce: bool = True,
bucket_size_mb: int = -1,
op: str = 'sum',
group: Optional[ProcessGroup] = None) -> None:
"""All-reduce parameters.
Args:
params (List or Generator[torch.Tensor, None, None]): List of
parameters or buffers of a model.
coalesce (bool, optional): Whether to reduce parameters as a whole.
Defaults to True.
bucket_size_mb (int, optional): Size of bucket, the unit is MB.
Defaults to -1.
op (str): Operation to reduce data. Defaults to 'sum'. Optional values
are 'sum', 'mean' and 'produce', 'min', 'max', 'band', 'bor' and
'bxor'.
group (ProcessGroup, optional): The process group to work on. If None,
the default process group will be used. Defaults to None.
Examples:
>>> import torch
>>> import mmengine.dist as dist
>>> # non-distributed environment
>>> data = [torch.arange(2), torch.arange(3)]
>>> dist.all_reduce_params(data)
>>> data
[tensor([0, 1]), tensor([0, 1, 2])]
>>> # distributed environment
>>> # We have 2 process groups, 2 ranks.
>>> if dist.get_rank() == 0:
... data = [torch.tensor([1, 2]), torch.tensor([3, 4])]
... else:
... data = [torch.tensor([2, 3]), torch.tensor([4, 5])]
>>> dist.all_reduce_params(data)
>>> data
[torch.tensor([3, 5]), torch.tensor([7, 9])]
"""
world_size = get_world_size(group)
if world_size == 1:
return
params_data = [param.data for param in params]
if coalesce:
_all_reduce_coalesced(params_data, bucket_size_mb, op=op, group=group)
else:
for tensor in params_data:
all_reduce(tensor, op=op, group=group)