Shortcuts

mmengine.dist.all_reduce_params

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

All-reduce 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.

返回类型:

None

示例

>>> 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])]
Read the Docs v: latest
Versions
latest
stable
v0.10.3
v0.10.2
v0.10.1
v0.10.0
v0.9.1
v0.9.0
v0.8.5
v0.8.4
v0.8.3
v0.8.2
v0.8.1
v0.8.0
v0.7.4
v0.7.3
v0.7.2
v0.7.1
v0.7.0
v0.6.0
v0.5.0
v0.4.0
v0.3.0
v0.2.0
Downloads
epub
On Read the Docs
Project Home
Builds

Free document hosting provided by Read the Docs.