This document is relevant for: Trn2, Trn3

nki.collectives.ReplicaGroup

class nki.collectives.ReplicaGroup

The set of ranks a collective is executed on, as a list of rank-lists.

The collective executes independently within each rank-list, and executes on all rank-lists in parallel. Each rank in the replica-group must participate in only one rank-list.

e.g. with [[0, 1, 2, 3], [4, 5, 6, 7]], the collective runs independently within ranks 0-3 and within ranks 4-7, in parallel.

Methods

from_partition_op

Build a ReplicaGroup by partitioning ranks [0, world_size).

static from_partition_op(world_size: int, op: rg_partition_op, operand: int) → ReplicaGroup

Build a ReplicaGroup by partitioning ranks [0, world_size).

Ranks sharing the same rank_list_id = op(rank, operand) form one rank-list.

Examples (8 ranks):

>>> ReplicaGroup.from_partition_op(8, rg_partition_op.bitwise_and, 0)
ReplicaGroup([[0, 1, 2, 3, 4, 5, 6, 7]])       # all ranks, one rank-list
>>> ReplicaGroup.from_partition_op(8, rg_partition_op.div, 4)
ReplicaGroup([[0, 1, 2, 3], [4, 5, 6, 7]])     # contiguous
>>> ReplicaGroup.from_partition_op(8, rg_partition_op.mod, 4)
ReplicaGroup([[0, 4], [1, 5], [2, 6], [3, 7]]) # strided

This document is relevant for: Trn2, Trn3