mmengine.dist.all_reduce

mmengine.dist.all_reduce(data, op='sum', group=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.

备注

Calling all_reduce in non-distributed environment does nothing.

参数:

• 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.

返回类型:

None

示例

>>> 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