all_gather#

composer.utils.dist.all_gather(tensor, group=None)[source]#

Collects a Tensor from each rank.

See also

torch.distributed.all_gather()

Parameters

  • tensor (Tensor) โ€“ Tensor from each rank to be gathered.

  • group (ProcessGroup, optional) โ€“ The process group to work on. If None, the default process group will be used. Default is None.

Returns

Sequence[Tensor] โ€“ A sequence of tensors indexed by rank.