mindscience.distributed.mappings.reduce_scatter_to_sequence

mindscience.distributed.mappings.reduce_scatter_to_sequence(x, group)[source]

Performs reduce-scatter across sequence partitions along the first dimension.

Parameters
  • x (Tensor) – Input tensor to be reduced and scattered across sequence partitions.

  • group (Union[CommGroup, CommGroupBase]) – Communication group for the operation.

Returns

Reduced and scattered tensor for the current rank.