Gang

class fairseq2.gang.Gang(rank, size, device)[source]

Bases: ABC

Represents a set of processes that work collectively.

Parameters:

  • rank (int) – The rank of this process in the gang.

  • size (int) – The number of processes that are part of the gang.

  • device (device) – The associated device.

abstract all_gather(output_tensor, input_tensor)[source]

Gather tensors from all ranks and put them in a single tensor.

Parameters:

  • output_tensor (Tensor) – The output tensor to accomodate tensor elements from all ranks.

  • input_tensor (Tensor) – The tensor to be gathered from the current rank.

abstract all_reduce(tensor, op)[source]

Reduce the tensor across all ranks.

Parameters:

  • tensor (Tensor) – The input and output tensor of the operation.

  • op (ReduceOperation) – The element-wise reduce operation.

abstract as_process_group()[source]

Return this gang as a process group.

Return type:

ProcessGroup

abstract barrier()[source]

Synchronize all ranks.