This document is relevant for: Trn2, Trn3

nki.collectives#

The nki.collectives module provides APIs for multi-core collective communication operations such as all-reduce and all-gather across NeuronCores.

NKI Collectives#

Collective operations for multi-rank communication.

`all_reduce`	Perform an all-reduce on the given replica group and input/output tensors.
`all_gather`	Perform an all-gather on the given replica group and input/output tensors.
`reduce_scatter`	Perform a reduce-scatter on the given replica group and input/output tensors.
`all_to_all`	Perform an all-to-all on the given replica group and input/output tensors.
`all_to_all_v`	Perform a variable-length all-to-all on the given replica group and input/output tensors.
`collective_permute`	Send and receive data between ranks based on explicitly defined source-target pairs.
`collective_permute_implicit`	Send and receive data between ranks in a ring, where sources and destinations are implicitly determined by the ring structure during runtime.
`collective_permute_implicit_reduce`	Perform an implicit collective permute with reduction in a ring, where sources and destinations are implicitly determined by the ring structure during runtime.
`collective_permute_implicit_current_processing_rank_id`	Returns the rank ID of the data to be processed in the current ring iteration.
`rank_id`	Get the rank ID of the current rank.

Defines a group of ranks that participate in a collective operation.

This document is relevant for: Trn2, Trn3