espnet2.gan_codec.shared.quantizer.modules.distrib.all_reduce

source

espnet2.gan_codec.shared.quantizer.modules.distrib.all_reduce(tensor: ~torch.Tensor, op=<RedOpType.SUM: 0>)

Perform a collective operation to reduce the tensor across all processes.

This function uses PyTorch’s distributed backend to perform an all-reduce operation on the provided tensor. The operation can be specified using the op argument, which defaults to summing the values across all processes.

• Parameters:
  • tensor (torch.Tensor) – The tensor to be reduced across all processes.
  • op (torch.distributed.ReduceOp , optional) – The reduction operation to perform. Default is torch.distributed.ReduceOp.SUM.
• Returns: The reduced tensor, updated in-place.
• Return type: torch.Tensor
• Raises:RuntimeError – If called when not in a distributed environment.

Examples

>>> import torch
>>> from your_module import all_reduce
>>> torch.distributed.init_process_group(backend='nccl')
>>> tensor = torch.tensor([1.0, 2.0, 3.0], device='cuda')
>>> all_reduce(tensor)
>>> print(tensor)  # Outputs the sum of the tensor across all processes.

NOTE

This function should only be called when the distributed process group is initialized. If the environment is not set up for distributed processing, a RuntimeError will be raised.