mmengine.dist.all_reduce_params¶

mmengine.dist.all_reduce_params(params, coalesce=True, bucket_size_mb=-1, op='sum', group=None)[source]¶

All-reduce parameters.

Parameters:

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.

Return type:

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