nemo_rl.weight_sync.membership#

Rank layout for a refit communicator rebuilt over the surviving generation shards.

Kept as pure arithmetic, separate from the Ray dispatch that applies it, because this is the part that has to be exactly right and the part that cannot be exercised on a workstation: reproducing a shard loss needs at least three GPUs (one trainer and two generation shards, so that losing one still leaves a fleet). The dispatch is a handful of .remote() calls; the rank layout is where an off-by-one silently corrupts a refit.

Module Contents#

Classes#

RefitMembership

Where each surviving generation shard sits in the rebuilt communicator.

Functions#

plan_refit_membership

Lay out a refit communicator containing only the surviving generation shards.

API#

class nemo_rl.weight_sync.membership.RefitMembership#

Where each surviving generation shard sits in the rebuilt communicator.

.. attribute:: world_size

total ranks, training plus surviving generation.

.. attribute:: train_world_size

unchanged; trainers are never excluded, which is what keeps the broadcast root (rank 0) stable across a rebuild.

.. attribute:: shard_prefixes

surviving shard index -> its rank prefix within the generation block. Insertion-ordered by shard index.

.. attribute:: workers_per_shard

Ray workers backing one shard (tp x pp).

world_size: int#

None

train_world_size: int#

None

shard_prefixes: dict[int, int]#

None

workers_per_shard: int#

None

property surviving_shards: list[int]#
exception nemo_rl.weight_sync.membership.NoSurvivingShards#

Bases: RuntimeError

Every generation shard is gone, so there is nothing to rebuild onto.

Initialization

Initialize self. See help(type(self)) for accurate signature.

nemo_rl.weight_sync.membership.plan_refit_membership(
*,
surviving_shards: collections.abc.Sequence[int],
dp_size: int,
total_gen_workers: int,
train_world_size: int,
) → nemo_rl.weight_sync.membership.RefitMembership#

Lay out a refit communicator containing only the surviving generation shards.

Prefixes are reassigned so the surviving ranks are contiguous from the start of the generation block, rather than leaving a hole where the dead shard was. That is not cosmetic: the nccl_reshard transport builds its destination device mesh as torch.arange(rank_offset, rank_offset + num_gpus), so a gap would silently misalign every parameter’s placements. It also matches what shrink does to a live communicator, which keeps the two paths describing the same world.

Parameters:

  • surviving_shards – shard indices still able to take part. Order is ignored; the result is sorted, so a caller that hands them over in discovery order still gets a deterministic layout.

  • dp_size – generation data-parallel size, i.e. the shard count.

  • total_gen_workers – Ray workers across the whole generation fleet.

  • train_world_size – training ranks, all of which stay in the communicator.

Raises:

  • NoSurvivingShards – if nothing survives.

  • ValueError – on an inconsistent topology or an unknown shard index.