This document is relevant for: Trn2, Trn3

nki.collectives.all_gather#

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

Perform an all-gather 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 gathered 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-gather 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 gather

  • dsts – List of output tensors to store results

  • replica_group – ReplicaGroup defining rank groups for the collective

  • collective_dim – Dimension along which output tensors are concatenated. Currently only 0 is supported for HBM tensors. For SBUF tensors, 0 or 1 is supported as SBUF collectives currently only operate on 2D tensors with a single free dimension.

  • 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