How a dead generation shard stops killing the run

PR #3591 · part 3/4 of #3454 · all links pinned at ab0ffe2f

Six pages, one question each — this one: what happens, in order.

#3590 gave the fleet a ledger: every shard has a state, and a dead one stops receiving traffic. That is enough to keep generating. It is not enough to keep training, because the weight refit is a collective — and a collective does not care that you stopped talking to a rank. Every remaining rank blocks inside NCCL waiting for one that will never arrive: no exception, no CPU burn, and Ray still reporting every actor healthy. This PR notices that from inside the stuck worker, breaks it, and rebuilds the communicator over whoever is left.

The wedge: one rank leaves, everybody waits The refit group is built once at setup over every training and inference rank. World = 2 + 3 = 5. rank 0 · trainer broadcast root waiting… rank 1 · trainer waiting… rank 2 · gen shard 0 process gone ✕ never calls broadcast rank 3 · gen shard 1 waiting… rank 4 · gen shard 2 waiting… every rank must call it — packed_tensor.py:111 The broadcast never returns. No exception. No error. No CPU burn. Ray reports every actor healthy. Nothing upstream can tell this apart from a refit that is merely slow — which is why the timeout has to live inside the stuck process, not in the driver.
Before this PR the run stayed here. One recorded case sat for 33 minutes before anyone noticed.
One refit, one dead shard, in order 2 trainer ranks · 3 generation shards · shard 0 dies · refit_timeout_s: 30 2 trainer ranks shard 0 shards 1 & 2 NEW IN #3591 1 Refit starts. Every rank enters the collective. Built once at setup over all training + inference ranks. World = 2 + 3 = 5. in collective in collective in collective 2 Shard 0's process dies. Mid-collective — it never reaches the other side of the broadcast. dead ✕ 3 Everyone else blocks inside NCCL. This is the wedge. No exception. No CPU burn. Ray still reports every actor healthy, so nothing upstream notices. Before this PR the run sat here — job 6258553 for 33 minutes. blocked forever blocked forever gone 4 At +30 s a watchdog thread inside each stuck worker gives up and aborts. Armed around the refit, in-process — because the stuck thread cannot report on itself. RefitAbortWatchdog · refit_watchdog.py:57 abort() released in 0.15 s abort() released in 0.15 s 5 abort, not destroy — and that is forced, not a preference. NCCL documents destroy as a collective every rank must call, which is exactly what a dead rank cannot do. The freed ranks raise RefitAborted so the controller can tell why. StatelessProcessGroup.abort · stateless_process_group.py:115 6 The controller asks the ledger who is gone, then plans the new rank layout. absent = {0}. Survivors are compacted to a contiguous prefix, not left with a hole. plan_refit_membership · membership.py:55 before 0 tr 1 tr 2 sh0 3 sh1 4 sh2 world 5 after 0 tr 1 tr 2 sh1 3 sh2 ✕ world 4 7 The transport rebuilds its communicator over the survivors — and re-addresses dispatch. Rebuilding is only half of it: refit dispatch walks the worker group, so without the second half it keeps calling shard 0's dead Ray actor and the next sync fails with RayActorError. reconcile_communicator · weight_sync/interfaces.py:125 8 The refit is retried on the smaller fleet. Training continues. Trainers are never excluded, so rank 0 stays a trainer and the broadcast root is stable. training generating gone Steps 6–7 do not wait for a death event. They run at the top of every weight sync — the one moment the refit group is provably idle and every rank is synchronized, which matters because the operations that change membership are themselves collectives. Doing it every time converges after a missed or reordered health update, instead of needing replay. Green = working · red = blocked or dead · blue = added by this PR
The whole feature is steps 4–7. Steps 1–3 are what already happened; step 8 is what #3590 could not reach.

The four things this PR had to invent

Three of these are new files or new interface methods; the fourth is a rule about the ledger from #3590 that is easy to get backwards.

AbstractionWhat it isWhy it had to exist
RefitAbortWatchdog
new file		 A thread armed around the refit inside each worker. If the deadline passes it calls abort() on every group it was handed, and records that it fired. The stuck thread is inside NCCL and cannot report on itself, and the controller cannot see the difference between a slow refit and a wedged one. Somebody local has to hold a clock.
abort()
new method		 Tears down the communicator with ncclCommAbort and drops the rendezvous store. Safe to call twice, and safe on a group that never built a communicator. destroy is a collective every rank must call. A dead rank cannot call it. abort works whether or not the peers are alive, which makes it the only safe choice on a path that exists to handle dead peers.
plan_refit_membership
new file, pure		 Takes the set of surviving shards and returns the whole rank layout: new world size, and each survivor's new offset, compacted to a contiguous prefix with no holes. The reshard destination mesh is arange(offset, offset + n), so a gap does not error — it silently misaligns every parameter. This is the part that must be exactly right and the part you cannot exercise without ≥3 GPUs, so it was split out as plain Python with no Ray in it.
reconcile_communicator
new interface method		 One hook per weight-sync transport. Returns True rebuilt, False nothing was absent, None this transport owns no membership. Only two of the seven transports own an NCCL world. The other five inherit a no-op, so the controller can say "this transport cannot recover that way" instead of guessing. Adding an eighth NCCL transport costs one method.
The rule that is easy to get backwards. "Absent" is deliberately not the complement of "serving". The obvious implementation reads the serving set, and it is a bug: SUSPECT (failing probes, not yet condemned) and STALE (reloaded, holding old weights) are both withheld from traffic while their processes are alive, and both join a refit normally. Using the serving set would abort a run on one failed probe — and would abort it precisely when a STALE shard is waiting to be refit, which is the recovery, not the failure. Hence absent_shards() over {DEAD, RESTARTING, RETIRED} only.

Does it work?

The abort itself is measured. On 2×A6000, with a peer SIGKILLed mid-broadcast, a survivor blocked in the collective was released 0.15 s after another thread called abort(). That is the one number on this page that came from hardware, and it pins the only part that was genuinely uncertain — that ncclCommAbort frees a thread already inside a collective.

End to end, nobody has watched it recover in CI. The recovery harness needs 3 GPUs (2 generation shards + 1 trainer) and both CI lanes have 2, so on this PR's own green run every one of the seven registered lines printed [recovery] SKIP: … found 2. and exited — seven lines in 0.79 s on H100 and 0.48 s on GB200, both reported PASSED. The feature was validated by hand on ≥3-GPU hardware — the 7-digit job IDs in the harness comments are an internal scheduler's, not GitHub Actions' — so the evidence exists; it just is not in CI, and a green lane currently says nothing about shard-death recovery either way.

One gap the review found, because it is the same failure this page is about. The Python reshard path splits the communicator into per-replica-group children. NCCL gives a split child its own abort flag unless splitShare is set, and it defaults to 0 (init.cc:2196), while ncclCommAbort sets only the parent's (:2915). So step 4 does not release a rank blocked on a child, and the run wedges exactly as it did before — the watchdog fires and nothing happens. Narrow (it needs the Python transfer path and a multi-member replica group), but it is step 3 coming back after step 4 was supposed to end it. Raised on the PR with a two-line fix: hand the children to the same cleanup before aborting the parent.