`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#

`desired_membership`	The membership the refit communicator should have, given who is absent.
`should_rebuild`	Whether the refit communicator has to be rebuilt before the next transfer.
`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.desired_membership( *, absent_shards: collections.abc.Sequence[int], dp_size: int, total_gen_workers: int, train_world_size: int, ) → nemo_rl.weight_sync.membership.RefitMembership#

The membership the refit communicator should have, given who is absent.

Phrased as “what it should be” rather than “what changed”, so the same call answers both directions: a shard leaving and a restarted shard rejoining are both just a different desired set. Comparing this against what was actually built is what lets a recovered engine be re-admitted – keyed off the absent set alone, an empty absent set is indistinguishable from “nothing to do”, and a restarted shard would stay excluded forever.

nemo_rl.weight_sync.membership.should_rebuild( *, desired: nemo_rl.weight_sync.membership.RefitMembership, built: nemo_rl.weight_sync.membership.RefitMembership, absent_shards: collections.abc.Sequence[int], force: bool, ) → bool#

Whether the refit communicator has to be rebuilt before the next transfer.

Here rather than in each synchronizer because both hardened transports need the same answer and had the same ten lines, byte for byte. Same reason desired_membership is here: the arithmetic is what has to be exactly right, and two copies of a rule is the bug class this stack keeps hitting.

Compared against what was built, not against “is anything absent”. Keyed off the absent set alone this returns False the moment a restarted shard comes back, leaving it permanently excluded from a communicator it should rejoin.

force is how the recovery path says the communicator is GONE rather than merely unchanged: after an abort the membership is identical and the communicator is dead, so skipping would retry over nothing.

It only overrides the skip when something IS absent. Forcing a rebuild with an empty absent set would produce a communicator that still contains the rank that just went silent, and the retry would hang on it exactly as the first attempt did. That case – a frozen-but-alive rank, which never becomes absent – must fall through to False so the caller reports “no generation shard could be identified as absent” and stops.

Parameters:

desired – what the membership should be, given who is absent now.
built – what the live communicator was actually built over. Callers treat an unrecorded membership as the full fleet: init_communicator builds over everything, so “not recorded” is not “unknown”, and calling it a difference would rebuild pointlessly on the first refit of every run.
absent_shards – shards currently absent, as the caller sees them.
force – the caller knows the communicator is gone rather than stale.

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.