colossalai.communication.collective

colossalai.communication.collective.all_gather(tensor, dim, parallel_mode, async_op=False)

Gathers all tensors from the parallel group and concatenates them in a specific dimension.

Parameters

  • tensor (torch.Tensor) – Tensor to be gathered

  • dim (int) – The dimension concatenating in

  • parallel_mode (colossalai.context.ParallelMode) – Parallel group mode used in this communication

  • async_op (bool, optional) – Whether operations are asynchronous

Returns

The tensor generated by all-gather

Return type

torch.Tensor

colossalai.communication.collective.reduce_scatter(tensor, dim, parallel_mode, op=<ReduceOp.SUM: 0>, async_op=False)

Reduces all tensors then scatters it in a specific dimension to all members in the parallel group.

Parameters

  • tensor (torch.Tensor) – Tensor to be reduced and scattered

  • dim (int) – The dimension scattering in

  • parallel_mode (colossalai.context.ParallelMode) – Parallel group mode used in this communication

  • op (ReduceOp, optional) – The type of reduce operation

  • async_op (bool, optional) – Whether operations are asynchronous

Returns

The tensor generated by reduce-scatter

Return type

Tensor