This document is relevant for: Trn2, Trn3

nki.collectives.all_reduce#

nki.collectives.all_reduce(srcs: List[NkiTensor], dsts: List[NkiTensor], replica_group: ReplicaGroup, op, priority: Optional[int] = None, name: Optional[str] = None) → NkiInstruction[source]#

Perform an all-reduce on the given replica group and input/output tensors.

The srcs and dsts parameters accept lists of tensors to support coalesced collective communication, which allows multiple tensors to be reduced in a single collective operation for improved efficiency.

Tensors can reside on either HBM or SBUF, but not a mix of the two. Coalesced collective communication (multiple tensors) is only supported when tensors are on HBM.

On HBM, the tensors within srcs must be either all private_hbm or all shared_hbm, and likewise within dsts, but srcs and dsts need not match. When LNC degree is 1, the srcs must be on private_hbm; only the dsts may be shared, and only when NKI_LNC1_SHARED_HBM=1 is set.

A shared_hbm destination is a single allocation that both NeuronCores on an HBM stack address, rather than one buffer per core. An all-reduce produces the same result on every rank, so when the LNC degree is 1 and both ranks of a replica group share an HBM, the collective can write the result once and both ranks read it. A private_hbm destination gives each rank its own copy of identical data, costing twice the capacity and bandwidth.

Parameters:

  • srcs – List of input tensors to reduce

  • dsts – List of output tensors to store results

  • replica_group – ReplicaGroup defining rank groups for the collective

  • op – The reduction operation to perform (nl.add, nl.minimum, or nl.maximum)

  • priority – DMA quality-of-service priority level 0-3 where lower is higher priority (NeuronCore-v4+ only)

  • name – (optional) name for the instruction.

This document is relevant for: Trn2, Trn3