feat(sc): keep training when a generation shard dies - #3591
Open
terrykong wants to merge 29 commits into
Open
Conversation
Elastic recovery rebuilds the refit communicator whenever the generation
fleet's membership changes, so init_collective goes from running once per
job to running once per recovery. Two things had to change before that is
safe.
Add StatelessProcessGroup.abort(). It is idempotent, safe on a group whose
communicator was never built, and drops its reference *before* calling
abort so a failed release cannot leave broadcast() pointing at a dead
communicator. abort(), not destroy(): NCCL documents destroy as an
intra-node collective that every rank must call or it hangs, which is
precisely what a rank whose process has died cannot do. Verified on
2xA6000 -- with a peer SIGKILLed mid-broadcast, a survivor blocked in the
collective was released 0.15s after another thread called abort().
Release the previous group on both sides of the refit before rebuilding.
Both init_collective implementations previously overwrote
self.model_update_group outright, stranding the old NCCL communicator and
its TCPStore. That is invisible in a one-shot job, which is why it
survived until membership became dynamic, and unbounded once recovery can
repeat.
model_update_group is now declared on both classes instead of springing
into existence on first assignment, so a rebuild can test for a previous
group without probing for the attribute. That also removes a
pyrefly ignore[implicitly-defined-attribute] on the vLLM side.
broadcast() now raises with a diagnostic instead of an AttributeError
when the group has no communicator -- the failure a rebuild bug produces.
Verification: 8 new process-group tests; 30/30 in the vllm lane including
2 new init_collective tests; ruff clean; pyrefly errors drop 9 -> 7 (the
7 remaining are missing optional imports, fastokens and awscrt, in files
this commit does not touch). test_vllm_generation.py was run with and
without this change and is identical at 9 failed / 38 passed / 3 skipped
-- those failures are a pre-existing vLLM/dynamo engine-init issue on
A6000 ('QKVParallelLinear' object has no attribute 'workspace'),
structurally upstream of anything this commit changes.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Signed-off-by: asolergibert <asolergibert@nvidia.com>
A NCCL broadcast needs every rank in the communicator to take part, so
when a generation rank dies the refit blocks forever inside NCCL: no
exception, no progress, and Ray still reporting every actor healthy. That
silent wedge is the failure this effort exists to remove. This commit
converts it into a precise error; rebuilding over the survivors, which
turns the stop into a recovery, is the next step.
Add WeightSynchronizer.reconcile_communicator(absent_shards). It is
non-abstract and defaults to a no-op, so the transports that own no NCCL
world of their own -- IPC, HTTP, checkpoint-engine -- are unaffected. The
two NCCL transports refuse the refit when a rank is missing.
Call it from the top of _sync_weights. Reconciling on a schedule rather
than on a death event is idempotent and converges after a missed or
reordered health update, and that point is the only one where the refit
group is provably idle and every rank is synchronized -- which matters
because the operations that change membership are themselves collectives.
"Absent" is deliberately not the complement of "serving". A SUSPECT shard
is failing probes but not yet condemned, and a STALE shard has reloaded
and holds old weights; both are withheld from traffic, and both processes
are alive and join a refit normally. Reading the serving set would abort
a run on a single probe blip, and would abort it precisely when a STALE
shard is waiting to be refit -- which is the recovery, not the failure.
GenerationFleetMonitor.absent_shards() therefore uses its own state set,
{DEAD, RESTARTING, RETIRED}.
The two transports raise different messages on purpose. The plain
broadcast could in principle drop a receiver, but nccl_reshard cannot:
prepare_nccl_reshard_refit_info derives each parameter's destination
placements from gen_world_size, so resizing without regenerating the plan
would leave survivors holding slices nobody wrote -- silent corruption,
worse than stopping.
Inert by default: async_rl.fleet_health.enabled is false, so there is no
monitor, no notion of a shard being gone, and the transport keeps the
membership it was built with.
Verification: 18 new tests, including the cases pinning SUSPECT and STALE
as present and DEAD and RESTARTING as absent; 1259 passed / 7 skipped
across algorithms, experience, weight_sync, single_controller,
distributed, fleet health and the router; ruff clean; pyrefly unchanged
at the same 7 pre-existing missing-import errors in files this commit
does not touch. Three existing _sync_weights tests build their controller
by hand and needed _fleet_monitor added.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Signed-off-by: asolergibert <asolergibert@nvidia.com>
…ards Turns the previous stop into a recovery. When a generation shard dies the refit communicator still contains its ranks, so the broadcast blocks forever inside NCCL. P3a.2 detected that and failed loudly; this rebuilds the communicator without the dead ranks so training continues on what is left. Rebuild rather than shrink. The pinned NCCL runtime exports ncclCommShrink but not ncclCommGrow, so a shrunk world could never take a recovered engine back -- one mechanism that works in both directions beats two that each work in one. It is also what nccl_reshard will need, since that transport must regenerate its refit plan rather than resize. The rank arithmetic is a pure function in weight_sync/membership.py, separate from the Ray dispatch that applies it, because it is the part that has to be exactly right and the part that cannot be exercised here: observing a real shard loss needs at least three GPUs, so that losing one still leaves a fleet. An off-by-one does not crash, it points a receiver at the wrong slice of the broadcast. Survivors are compacted to contiguous prefixes rather than leaving a hole where the dead shard was. Not cosmetic: the nccl_reshard destination mesh is torch.arange(offset, offset + num_gpus), so a gap would misalign every parameter's placements. It also matches what shrink does to a live communicator, so both paths describe the same world. VllmGeneration.rebuild_collective addresses the surviving DP leaders directly instead of going through run_all_workers_multiple_data, which walks every worker in the group and would therefore dispatch to the shard we are rebuilding *because* it is gone. Only leaders are called; each collective_rpcs into its own TP/PP workers. Trainers are never excluded, so rank 0 stays a trainer and the broadcast root is stable across a rebuild. Each rebuild takes a fresh port, since the previous world's rendezvous store may still be bound, and StatelessProcessGroup.abort() now releases that store so repeated recoveries do not accumulate one per recovery. nccl_reshard still refuses, with a message pointing at the transport that does recover. Regenerating its plan is the next step. Adds tests/functional/grpo_sc_generation_shard_recovery.sh: kills one of two generation shards mid-run and asserts the job completes all steps, that the rebuild actually happened, and that the metrics are intact -- completion alone would also be satisfied by a run that never noticed the death. It needs >= 3 GPUs and self-skips below that rather than passing vacuously, and is registered in the SingleController lane in full mode. Verification: 36 new unit tests (18 on the rank layout alone); 801 passed across weight_sync, single_controller, distributed, experience, fleet health and the router; ruff clean; pyrefly back to the same 7 pre-existing missing-import errors with membership.py added to scope. End-to-end recovery is NOT verified here -- it needs the >= 3 GPU functional test. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: asolergibert <asolergibert@nvidia.com>
… refit dispatch Two things, because they are the same defect seen from two sides. FIXES A GAP IN THE PREVIOUS COMMIT. Rebuilding the communicator is only half a recovery: update_weights_from_collective and nccl_reshard_refit dispatch through run_all_workers_single_data, which walks the whole worker group. After a shard was lost they kept calling its dead Ray actor, so the next refit failed with RayActorError and the run still died -- just later, and with a less obvious cause than the hang it replaced. Both now address the surviving DP leaders, recorded by set_refit_membership. Before any loss the membership is unset and every leader is addressed, which is the whole life of a run that never loses a shard. nccl_reshard now recovers instead of refusing. It needs more than the plain broadcast: the shared model_update_group and the per-PP-stage bulk groups both have to be rebuilt, and the refit plan has to be regenerated, because prepare_nccl_reshard_refit_info derives each parameter's destination placements from the inference world size. Reusing a plan built for the old fleet does not error -- a stale mesh is still a valid mesh -- it just has survivors writing the slices the dead shard owned and leaving their own unwritten. init_communicator and the rebuild now share one _build(membership), so there is a single copy of that arithmetic. Two copies is exactly how the communicators and the plan would drift apart, and the drift is silent. It also means every normal run exercises the rebuild path rather than leaving it as a rarely-taken branch. Membership is recorded before the build, not after: step 3 distributes the regenerated plan through prepare_nccl_reshard_refit_info, which consults it, so setting it afterwards sent the new plan to the shard the rebuild had just excluded. Caught by the tests, not by review. RefitMembershipChanged is removed. Nothing raises it now that both transports recover; NoSurvivingShards covers the terminal case where the whole fleet is gone. The recovery functional test gains REFIT_TRANSPORT, and the CI lane runs it on both transports. Its rebuild assertion now matches either transport's log line -- as written it only matched the collective one, so a reshard regression would have passed silently. Verification: 12 new tests on reshard rebuild and refit dispatch; 180 passing in weight_sync, 626 across weight_sync, single_controller, distributed, fleet health and the router; ruff clean; pyrefly unchanged at 7 pre-existing missing-import errors. End-to-end recovery remains unverified here and needs the >= 3 GPU functional test. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: asolergibert <asolergibert@nvidia.com>
Two portability bugs, both of which only bite on the >= 3 GPU machines
where this test is the only place the recovery scenario can run at all.
1. Megatron asserts global_batch_size % (micro_batch_size *
data_parallel_size) == 0. The config inherits 512 / 4 from
grpo_math_1B.yaml and tp=pp=cp=1, so dp is just the training GPU
count -- and the test took every GPU the host had:
3 GPUs -> dp=1 512 % 4 = 0 ok
4 GPUs -> dp=2 512 % 8 = 0 ok
5 GPUs -> dp=3 512 % 12 = 8 assertion failure
8 GPUs -> dp=6 512 % 24 = 8 assertion failure
16 GPUs -> dp=14 512 % 56 = 8 assertion failure
8 is the usual CI runner size, so this would have failed on the first
real CI run with an error saying nothing about shard recovery. Round
the training ranks down to a power of two -- 512 and 4 are both powers
of two, so any power-of-two dp divides -- and claim only the GPUs
actually used rather than the whole node.
2. `pgrep -f VllmAsyncGenerationWorker` matched three things per shard:
the Ray actor, the per-worker venv python child, and the bash launcher
that execs it. Measured live on a single-shard run, the substring
matched 2-3 processes where there was exactly one actor.
So GEN_PIDS[0], the lowest pid, was as likely to be a child as an
actor, and the guard `(( ${#GEN_PIDS[@]} < 2 ))` was satisfied by the
children alone -- it would pass with ZERO actors present. Killing a
child leaves both shards serving, the run completes exactly as it
would have anyway, and the test reports a pass having never exercised
recovery. A false pass, which is worse than a false failure here.
Match the actor structurally instead, the same way the chaos test now
does, require exactly GEN_GPUS of them, print the victim, and verify
it actually died.
Cannot be run end-to-end on a 2-GPU box, so: the actor matcher is
verified against all seven process titles observed in a live run and
against live pid counts, and the sizing arithmetic against every host
size from 2 to 16 GPUs.
Signed-off-by: asolergibert <asolergibert@nvidia.com>
collections.abc sorts before dataclasses. Same invisible-locally rule as the previous two branches -- ruff's `I` selection only runs via the second ruff hook in .pre-commit-config.yaml. Signed-off-by: asolergibert <asolergibert@nvidia.com>
Job 5861743 on a 4-GPU GB200 cluster failed all five kill-based functional tests for one reason: the harness found zero generation actors. [recovery] FAIL: expected exactly 2 generation actors, found 0 [chaos] FAIL: job died before the kill (both idle and serving) At train step 3, with generation demonstrably working. Both chaos variants failing identically is the tell -- idle and serving are different patterns, so this is not a state-timing race, nothing matched at all. Both tests picked their victim by Ray's process TITLE (ray::VllmAsyncGenerationWorker). That works on the development workstation and matches nothing on that cluster. Titles are a runtime implementation detail; the GCS actor table is the runtime's own record of which pid is which actor, so ask that instead. Not ray.util.state.list_actors: it goes through the dashboard HTTP state server and init_ray starts Ray with include_dashboard=False, so it raises ServerUnavailable. ray._private.state.actors() reads GCS directly. Verified against a live cluster from a separate process -- it returns exactly the actor pids. The old diagnostics were useless here: both greped for "ray::" and so printed nothing at precisely the moment the ray:: assumption was itself wrong. They now dump unfiltered process listings. LIMITATION, stated rather than hidden. Ray only exposes the running method via the dashboard state API, so idle-vs-serving still needs titles. Where titles work the distinction is preserved (verified: one actor shows ray::...Worker, another ray::...Worker.<method>). Where they do not, chaos now says so and offers VICTIM_STATE=any, which still asserts a bounded attributable failure but stops distinguishing the detection path from the in-flight-RPC path. NOTE ON PLACEMENT: the chaos change belongs on the containment branch, where that test is introduced. It is here temporarily so branches 3 and 4 can be re-run while the containment PR's CI is in flight; it should be moved down once that finishes. Signed-off-by: asolergibert <asolergibert@nvidia.com>
Job 5866390 still failed all five kill-based tests, but the two harnesses
failed differently and only one of them was informative.
recovery: "found 0" and the diagnostic printed NOTHING -- no error, no
pids. So the helper attached to Ray fine and matched zero
actors, and there was no way to tell "no actors" from "wrong
field name" from "wrong state string".
chaos: "could not attach to a running Ray cluster" -- but that came
after "job died before the kill", so Ray was legitimately gone
by then. Consistent, not a second bug.
The helper now always dumps the GCS actor table to stderr -- class name,
state, pid, name, one line each -- before reporting matches. One run will
say what that cluster actually calls its generation actors, instead of a
third round of guessing. It is also less eager to filter: State is no
longer required to be ALIVE, since a RESTARTING shard is still a real
process, and over-filtering is what produced the silent empty result.
Only Pid 0 is excluded, because there is nothing to kill.
Also fixes a performance bug this introduced. The chaos loop called the
helper every 0.2s, and each call is a full ray.init/shutdown of about
three seconds, so VICTIM_WAIT_S=300 meant roughly 60 checks rather than
1500 and spent the budget on Ray handshakes. Discovery now runs once with
its own retry and budget (ACTOR_WAIT_S), then only the titles are polled
-- the actor set is stable once generation is up, only the state changes.
Separating the budgets also separates the failure messages: "no actors
found" and "actors never reached the state" were previously the same line.
The recovery query gets a bounded retry for the same reason: a single
query races the GCS write.
Signed-off-by: asolergibert <asolergibert@nvidia.com>
Job 5886540 failed all five kill-based tests and the logs could not say why, because two things were missing. 1. NEITHER failure path dumped the training log. recovery printed only "expected exactly 2 generation actors, found 0", which reads as a harness bug -- but the ps dump in the same block listed ONLY harness processes (the two sruns, `uv run ...recovery.sh`, `bash ...recovery.sh`). The training job the harness had just reported as pid=523789 was absent, and its command line contains `policy.generation` six times, so it would have matched had it been alive. The run had ended; the harness had no idea and said something misleading instead. Both paths now tail the training log. 2. The recovery retry loop never rechecked the job. Each attempt is a full ray.init/shutdown, so ten of them is roughly thirty seconds -- and on this hardware `dp` runs its whole 2-step job in about 2.5 minutes, nearly all startup, so the remaining steps finish well inside that window. The loop now checks the job every attempt and reports "the run ended before a shard could be killed", with the hint to raise MAX_STEPS or lower KILL_AFTER_STEP, rather than blaming actor discovery. Also makes the attach failure diagnosable. `address="auto"` is resolved from a file in Ray's temp dir, so the helper now reports RAY_ADDRESS, RAY_TMPDIR, TMPDIR, which temp dir it consulted, and what is in it. "Could not find any running Ray instance" currently cannot distinguish "the cluster is gone" from "the cluster is up but its session is elsewhere", and those need opposite fixes. What is still unexplained: during the chaos discovery loop the training job was demonstrably alive and the helper still could not attach. That is the open question, and these diagnostics exist to answer it on the next run rather than by a fourth guess. Signed-off-by: asolergibert <asolergibert@nvidia.com>
…teps
The training logs settled both open questions.
WHY THE JOB VANISHED. It did not crash -- it finished. The chaos log ends
with "SC run complete: {'train_steps': 50, 'trainer_version': 50}". On
GB200 a step takes seconds (the dp test runs its whole job in about 2.5
minutes, nearly all startup), so the run outlasts nothing. Recovery's 12
steps completed while the harness was still looking for a shard to kill.
MAX_STEPS is now 24, so the run is still going well after the kill --
otherwise "it completed" proves nothing about recovery.
WHY THE HELPER COULD NOT ATTACH. `address="auto"` resolves through
/tmp/ray/ray_current_cluster, and that marker is not written for a
driver-managed cluster -- verified: it is absent here too, while
`ray.init()` is live. The session directory does record the GCS port, so
the address can be rebuilt from session_latest/gcs_server_port_* plus the
recorded node ip. Verified against a live cluster: the rebuilt address is
byte-identical to ray.get_runtime_context().gcs_address.
The helper now tries RAY_ADDRESS, then "auto", then the rebuilt address,
and only then reports failure -- with RAY_ADDRESS/RAY_TMPDIR/TMPDIR, the
temp dir it consulted and its contents, and what the rebuild produced.
Signed-off-by: asolergibert <asolergibert@nvidia.com>
Job 5888097 got further -- the diagnostics now say the right things -- but the one message that would explain the failure is still being thrown away. What the run showed. The address fallback works: it rebuilt '10.109.19.148:21829' from session_latest, which is a real cluster address. The recovery test correctly reported "the run ended before a shard could be killed" instead of blaming actor discovery. And the training logs confirm both jobs ran to completion -- 50/50 and 24/24 steps, the recovery run taking 8.5 minutes, so the runs are no longer too short. What is still unexplained is why discovery finds nothing while the job is alive. Both failure dumps run AFTER the job has gone, so they can only ever report "cluster not found" -- which is true but useless. The in-loop attempts are the informative ones and their stderr went to /dev/null. It now goes to $EXP_DIR/actors.err, and the failure path prints the last 30 lines of it, labelled as the attempt made while the job was alive and distinguished from the post-mortem call. Three rounds of this have been undiagnosable for want of that one redirect. Suspected but unconfirmed: each attempt now costs two ray.init calls (auto, then the rebuilt address), so ten attempts can outlast the run's remaining steps. The captured stderr will show whether that is what is happening or whether the connect fails for another reason. Signed-off-by: asolergibert <asolergibert@nvidia.com>
…ath that fires Job 5890572 failed for two reasons, both in the harness. 1. The discovery loop outlived the job it was meant to interrupt. Step timings in the log are 20.6s, not the "seconds" assumed earlier -- that figure came from the dp test, which uses a far smaller config. So the 24-step run lasts about 8 minutes and had roughly 7 left after step 3. Ten discovery attempts, each costing two ray.init connect timeouts of about 20s, consume all of it. Each attempt is now bounded by ACTOR_QUERY_TIMEOUT_S (20s) so a failing query fails fast. 2. The stderr dump added last round was on the wrong branch. The run took the "run ended before a shard could be killed" path, which only tailed the training log; the actors.err dump sat in the "expected exactly N actors" path, which did not fire. So the file was written and never shown. Both paths print it now. MAX_STEPS stays at 24: with 20s steps that is 8 minutes, which is ample runway once discovery stops eating it. Signed-off-by: asolergibert <asolergibert@nvidia.com>
The actors.err file did not exist after job 5890572, even though the run took the path that writes it. Two rounds have now been spent routing this diagnostic through a separate artifact, and neither produced a file that survived to be read. So stop using one. The harness log is already captured and already collected; put the helper's stderr straight into it. stdout (the pids) goes to a temp file, stderr into a shell variable, and the first two failing attempts print it inline. Capped at two so a ten-attempt loop does not flood the log, and shown again on the failure path. Verified end to end against a live Ray cluster: the pids are still parsed correctly from stdout (2 actors found, matching ray.get) while stderr is captured separately for reporting. This is a diagnostic-plumbing fix, not a fix for the underlying problem. Discovery finding nothing while the job is alive is still unexplained -- but the next run will finally say why, in the log you already have. Signed-off-by: asolergibert <asolergibert@nvidia.com>
Root cause, finally. The harness detects progress by grepping RUN_LOG, and RUN_LOG was empty until the run ended. Job 5892910, with timestamps: 10:40:38 "train step 3/24" produced 10:48:15 EngineCore shutdown -- the run finishes 10:48:21 discovery attempt 1 Discovery attempt ONE ran six seconds after the job had already gone. The wait-for-step-3 loop polls every 5s, so it should have fired at 10:40:43; it did not fire for 7 minutes 43 seconds, by which point there was nothing left to kill. The actor's print already uses flush=True, but that only reaches the DRIVER: Ray forwards actor output there, and the driver's own stdout is a redirected file, which Python block-buffers. Demonstrated directly -- a process writing to a redirected file shows 0 bytes on disk two seconds in and only appears at exit, while the same process with PYTHONUNBUFFERED=1 is visible immediately. This explains every kill-based failure on this cluster, and why they all looked like actor-discovery problems: discovery was fine, it was just never started until the cluster it was meant to query had shut down. It also explains why the earlier process-title and Ray-attach theories each looked plausible and each failed to fix anything. Applies to both harnesses; the chaos test greps the same way. Signed-off-by: asolergibert <asolergibert@nvidia.com>
Job 5893807 is the first run where the harness worked end to end: step 3 detected, both actors found via the GCS, one killed cleanly. The run then wedged, and the log could not say where. It also disproved the diagnosis I had drawn from the previous run. The ordering is: 711 _sync_weights: sync done step 3's refit COMPLETED 724 train step 3/24 step 3 finished cleanly 726 fleet: shard 0 healthy -> suspect 734 fleet: shard 0 suspect -> dead 742 rollout stall -- 0 rollouts in flight, step 4/24 So the shard died AFTER the refit completed, and step 4's _sync_weights never started. There is no hung collective: the wedge is on the rollout side, and "0 rollouts in flight" says the pump stopped without saying whether it is parked on _rollout_permitted, on the _buffer_capacity semaphore, or in the sampler. Those need different fixes. (The surviving shard's EngineDeadError at 11:39:00 is not the cause either -- that is the last timestamp in the log, i.e. teardown when the harness SIGKILLed the run.) The chaos test has dumped stacks on a wedge for a while; this one did not, which is why a full cycle was spent guessing. It now dumps the SingleControllerActor with --locals -- the locals are the point, since the frame alone does not distinguish which of the three parks it is -- plus the policy and generation workers, and the fleet/refit history. No production change. The abort-and-rebuild fix is not implemented, because the evidence no longer supports the hang it was meant to break. Signed-off-by: asolergibert <asolergibert@nvidia.com>
…builds
The py-spy dump from job 5898311 settles where the recovery test wedges,
and it is the refit collective after all:
ray::MegatronPolicyWorker.broadcast_weights_for_collective
synchronize (torch/cuda/streams.py:108)
packed_broadcast_producer (nemo_rl/utils/packed_tensor.py:109)
broadcast_weights_for_collective (megatron_policy_worker.py:2180)
Both policy workers, actively stuck on a CUDA stream sync that cannot
complete because the dead generation rank never joins. The surviving vLLM
worker is idle, not in update_weights_from_collective. (An earlier reading
of the log ordering suggested the refit had completed and the wedge was on
the rollout side. The stacks disprove that.)
What the logs still cannot say is WHY the membership was not reconciled
first, and that is the difference between two opposite fixes:
absent == {} reconcile ran before fleet health recorded the death, so
it correctly did nothing -- the bug is the race, and the
fix is to re-check immediately before the broadcast.
absent == {0} reconcile saw the dead shard and reconcile_communicator
still returned False -- the bug is in the rebuild path.
`rebuilt = 0` is consistent with both, because the print fired only on a
rebuild. It now reports the absent set and the decision every time.
Not yet fixed, and worth recording: an abort-and-rebuild cannot be
delivered as designed. Breaking a hung collective means calling abort() on
the worker's process group, but the worker actors are created without
max_concurrency, so a blocked actor has no thread to service that RPC. Any
abort path has to make those actors concurrent first, or bound the
collective from inside the worker. Which of those is right depends on the
absent set above.
Signed-off-by: asolergibert <asolergibert@nvidia.com>
Elastic recovery worked only when the death landed OUTSIDE the refit
window. Job 5899746 recovered cleanly and completed 24/24 steps; job
5898311, same code, died inside the window and wedged for 1801s. py-spy
showed both policy workers blocked in
packed_broadcast_producer -> cuda stream synchronize, waiting on a
generation rank that no longer existed. The refit is about 10% of
wall-clock (weight_sync 1.93s of an 18.94s step), so this is a routine
event, not a corner case.
The existing reconcile could not cover it. It runs at the top of
_sync_weights, which is right for what it does -- the group is idle there
and every rank is synchronized, and the membership operations are
themselves collectives -- but it is a pre-check. Once the broadcast is
running, the only code that could repair membership sits at the top of the
function that is already stuck, and _rollout_permitted stays cleared, so
the rollout pump parks too. P0.6 recorded this as constraint C1 and
deferred it here.
THE ABORT HAS TO COME FROM INSIDE THE WORKER. The controller cannot
deliver it: the collective blocks the worker actor's event loop and the
worker actors carry no max_concurrency, so an abort RPC would queue behind
the operation it is meant to interrupt. RefitAbortWatchdog therefore runs
a thread in the worker itself, which is the arrangement the design's NCCL
spike actually validated (a survivor released 0.15s after another thread
called abort()).
Two spike findings shape the API, and both are easy to get wrong:
- An aborted collective returns WITHOUT raising, so try/except cannot
detect it. The caller must ask whether the abort fired; hence the
explicit `fired` flag rather than exception handling.
- The destination buffers hold PARTIAL weights afterwards. Survivors are
marked via the new mark_weights_partial(), which moves them to STALE --
alive, refits normally, does not serve. Reusing STALE rather than
inventing a state matters, because DEAD -> HEALTHY is already
unreachable and STALE ignores successful probes, so the only route back
into the serving set is a completed refit.
_sync_weights then reconciles, rebuilds, and retries exactly once. A second
abort is a real fault, and retrying forever would recreate the wedge this
removes.
Also reconciles immediately before the collective, not only at the top:
between the two there is calibration and two to_thread hops, and a death
recorded in that gap would otherwise be ignored until the next step.
Inert by default. refit_timeout_s defaults to None, which starts no thread
and leaves every path byte-identical to before.
Tests: 11 on the watchdog (disarmed cases start no thread; fired survives
the clean return; a failing abort does not escape; an exception in the
guarded block still disarms; 50 refits do not leak threads) and 4 on the
new state transition. 848 passed across the affected suites.
Not covered, and stated so a reviewer does not have to find it: a death
during the REBUILD itself, which is also a collective. Smaller window, same
watchdog could extend to it, out of scope here.
Signed-off-by: asolergibert <asolergibert@nvidia.com>
The recovery test kills at a step boundary, and the refit is only about 10% of wall-clock, so it reaches that window by chance. That is not a theoretical concern: identical code both passed (job 5899746, recovered and completed 24/24) and wedged for 1801s (job 5898311) on consecutive runs, purely on where the kill landed. A fix validated by that coin flip is not validated. KILL_DURING_REFIT=true waits for a refit to begin and kills immediately. The trigger is the "refit membership absent=" line, which is now printed on every sync right before the collective, so it marks the start of the window reliably rather than by timing guesswork. Registered as a fourth recovery variant. It is the only test that exercises the abort-and-rebuild path deterministically; the other three reach it occasionally at best. Ordering note for whoever reads this stack: this test is deliberately in the same PR as the fix rather than after it. P3a.3 shipped a defect that P3a.4 had to repair precisely because its suite never covered "and then a refit happens" -- the fix looked right and nothing exercised it. Signed-off-by: asolergibert <asolergibert@nvidia.com>
… test The abort watchdog shipped unreachable. CollectiveWeightSynchronizer took a refit_timeout_s argument, but there was no config field, the factory did not forward one, and no test set it -- so the watchdog could never arm and the whole path was dead code. Caught by asking which functional test actually enables the feature, which is a better question than it sounds. Adds async_rl.fleet_health.refit_timeout_s (Optional[PositiveFloat], default None), forwards it through create_weight_synchronizer, and reads it in the SingleController setup -- the only path with a fleet monitor to recover with. Enabled in the recovery functional test at 60s, in EVERY variant rather than only KILL_DURING_REFIT. A kill at a step boundary can still land inside the refit by chance -- that is precisely how job 5898311 wedged -- so leaving it off elsewhere would let the flaky case stay flaky. 60s against a healthy refit of ~1.9s for a 1.5B model on GB200. Chosen far above rather than near it: firing early aborts a run that was merely slow, firing late only means a wedge lasts a little longer before it breaks. Still None by default, so nothing outside the test changes. 369 passed; ruff, isort and format clean. Signed-off-by: asolergibert <asolergibert@nvidia.com>
The refit deadline reached only half the code. vllm_generation picks between two methods by config: update_weights_from_collective_async if async_engine <- what runs update_weights_from_collective otherwise <- what I changed Every SC functional test sets async_engine=true, so the very first refit died in Ray's argument validation: TypeError: got an unexpected keyword argument 'refit_timeout_s' Both recovery variants failed at step 1 of job 5924693 -- before reaching any of the behaviour under test. Two fixes, not one: 1. The async entrypoint takes refit_timeout_s and threads it through collective_rpc into the worker extension, which already had it. 2. That method wraps its body in `except Exception: return False`, which would have swallowed RefitAborted and reported a plain failure. The controller would then never see the signal telling it to rebuild and retry -- the run would just die, which is the outcome this whole change exists to replace. RefitAborted is now re-raised ahead of the blanket handler. The second would not have shown up in this job even with the signature fixed. It only appears when an abort actually fires, so it would have looked like "the fix does not work" on some later run. Adds a signature assertion covering both entrypoints. Nothing behavioural could catch this: the call crosses a Ray actor boundary, where signatures are validated at dispatch rather than at import, and the fakes in these suites do not model that. Verified the test fails (2 failed) with the signature reverted and passes (13) with it restored -- a test for a bug this cheap is worthless if it does not actually detect the bug. 630 passed, 3 skipped across the affected suites; lint clean. Signed-off-by: asolergibert <asolergibert@nvidia.com>
Job 5925668 ran both recovery variants and both died. The watchdog half worked -- the trainer and the surviving vLLM engine each aborted at 60s, and the 1800s wedge was gone -- but nothing after the abort did. Four defects, one per stage of the repair: 1. The handler asked the health monitor which shards were absent. That verdict is paced by probe rounds while the abort is an event, so the corpse was still SUSPECT, absent came back empty, and nothing was rebuilt. The retry then died on "StatelessProcessGroup has no communicator: the group was aborted and not rebuilt". The repair now probes on its own thread first, turning a race into a lookup, and a RayActorError is treated as conclusive rather than as one more ambiguous failure to count. 2. mark_weights_partial walks the *serving* shards, and a shard that has only just died is still serving -- so the corpse was moved to STALE, which is deliberately not an absent state. Even with the rebuild forced it would have been put straight back into the collective. Fixed by establishing the dead set before marking anything partial; mark_weights_partial already skips absent shards. 3. Nothing called report_refit, so mark_weights_partial was a one-way trip out of the serving set. A successful recovery would have left the fleet empty and the run would have died of an exhausted fleet -- a worse failure than the one being repaired. _promote_refit_shards now runs after every successful sync. 4. The other variant died of a RayActorError from ray.get(futures_inference) while ray.get(futures_train) had already returned: the broadcast completed and the shard died in the RPC epilogue. Survivors held complete weights and the run was killed anyway. Now recovered like the abort, minus the partial-weight marking, which would take a healthy fleet out of service over a transfer that worked. When no shard can be identified as absent the refit now fails attributably instead of retrying: there is no smaller membership to rebuild over, and rebuilding the full one would hang on the same silent rank -- the wedge this path exists to remove. Also makes KILL_DURING_REFIT reproducible. A refit here takes ~0.10s, so the harness could not aim at it: job 5925668 hit the RPC epilogue instead, testing something real but not what the variant claims. NRL_REFIT_HOLD_FILE lets the harness hold one specific refit open, kill inside it, and release the survivors into a collective the victim will never join. Absent in every real run, where it costs one os.path.exists per refit. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: asolergibert <asolergibert@nvidia.com>
The abort-and-rebuild machinery reached only the collective path. factory.py constructed CollectiveWeightSynchronizer with refit_timeout_s and NcclReshardWeightSynchronizer without it, and that class had no such parameter to begin with, so the reshard transport ran with no deadline at all: a generation rank dying inside a reshard refit wedged exactly as before the watchdog existed. Nothing pointed at it. recovery-reshard kills at a step boundary, so it passes through the between-refits path and never touches the missing machinery -- a green run on a transport with no protection. Plumbed the same way as the collective path, through every layer that crosses a Ray boundary: the synchronizer, both lm_policy and vllm_generation dispatchers, and BOTH generation worker entrypoints. The async one is what a real run uses, and omitting it is precisely how the collective path failed at its first refit with a TypeError out of Ray's argument validation. RefitAborted is re-raised ahead of the blanket handlers in both workers. Folded into `return False` it would be reported as a generic refit failure, which ends the run instead of triggering the rebuild -- the wedge this replaces, wearing a different error message. The watchdog now accepts several groups. This transport moves bulk weights over per-PP-stage groups and then broadcasts the remainder over the shared model_update_group, so a hang can be in either family and nothing at that level can tell which. All of them are aborted; abort() is idempotent and safe on a group that never built a communicator, and the recovery rebuilds every family anyway. Each abort is guarded separately, because the group that raises may not be the one the caller is blocked in. The signature test is now parametrized over both transports rather than asserting on the one entrypoint that had already broken. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: asolergibert <asolergibert@nvidia.com>
chaos-idle failed in job 5931102 while the run did exactly what the test exists
to verify: it died in seconds with the cause named -- "NoSurvivingShards: no
generation shard can take part in the refit; the fleet is gone". The harness
just did not recognise the exception.
Losing the only shard can surface as either of two typed errors, and which one
wins is timing, not meaning:
GenerationFleetExhausted too few shards SERVING, raised on the watchdog tick
NoSurvivingShards no shard can take part in the REFIT, raised by the
reconcile before every weight sync
The refit runs every step and the watchdog every watchdog.interval_s, so once
probe detection went from 30s to 5s the reconcile started getting there first.
Both are bounded and attributable, so both belong in the allowlist.
Not unified by having the reconcile call raise_if_exhausted() first, which looks
tidier and is a trap: an aborted refit marks every surviving shard STALE via
mark_weights_partial, so the serving set is legitimately empty between the abort
and the retry, and raising there would kill the run in the middle of the
recovery that was about to succeed.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: asolergibert <asolergibert@nvidia.com>
The abort-and-rebuild path had end-to-end coverage on the collective transport only. nccl_reshard had the deadline plumbed, the watchdog wired to both of its communicator families, and the fault-injection hook already in its receive -- but nothing ever aimed at it. recovery-reshard kills at a step boundary, so it recovers BETWEEN refits and never aborts; everything specific to aborting a reshard was resting on signature assertions. That gap matters because the two aborts are not the same operation. The collective path aborts one communicator. Reshard holds two families -- the per-PP-stage bulk groups and the shared model_update_group -- a hang can be in either, and the rebuild has to regenerate the refit plan as well as the communicators. A regression in that is invisible to every other test. Registers recovery-reshard-refit (REFIT_TRANSPORT=nccl_reshard + KILL_DURING_REFIT=true) in the L1 lane and in submit-ci.sh, which is the one combination that makes it happen: the hold file parks the receive, the victim is killed inside it, and the survivors are released into a collective it will never join. Also guards the hook itself. Two source-level assertions, because vllm_backend does `import vllm` at module scope and the default unit lane cannot import it: the hold must exist in BOTH receives, and it must sit INSIDE the RefitAbortWatchdog block. Held outside the guard the deadline clock never starts, so the victim dies during an unguarded pause and the run hangs exactly as it did before the watchdog existed -- while the test still reports a pass. Both assertions were verified to fail against those two mutations. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: asolergibert <asolergibert@nvidia.com>
…-resiliency-03-elastic-recovery
Same defect as the chaos test, in the script that kills a shard on purpose: cleanup() signals and returns, and SIGKILL does not free device memory synchronously. This one is worse placed -- the lane runs five recovery variants back to back and then chaos-serving, so there are six handoffs, and the recovery script had no pre-flight GPU check at all. A leaked engine would surface downstream as a placement-group timeout, which reads as an unrelated bug. cleanup() now waits for the memory to come back, and a pre-flight waits for $USED_GPUS before starting. It sits after the >= 3 GPU SKIP, so a host that can never run this test still skips immediately rather than waiting first. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: asolergibert <asolergibert@nvidia.com>
|
Auto-sync is disabled for ready for review pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
This was referenced Aug 12, 2026
Collaborator
Author
PR stack for #3454 (SingleController resiliency)Mirror of #3472. Part 3/4 of the stack (#3593) — see #3589 for the full stack map. Based on #3590, so the diff here is only this layer (+3437 −59, 32 files) rather than the cumulative +9051 the original showed.
Branch pushed unchanged from @asolergi-nv's fork ( |
This was referenced Aug 12, 2026
…-resiliency-03-elastic-recovery
upstream #3263 added a _should_use_nemo_gym check and a stale-in-flight abort to _sync_weights, so the minimal controller this fixture builds no longer reaches the refit -- it dies on AttributeError: no attribute '_master_config'. An empty env dict selects the native path and an empty registry makes _abort_stale_inflight a no-op. Neither is what these tests are about; both just have to exist for the refit to be reached at all. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: asolergibert <asolergibert@nvidia.com>
…-resiliency-03-elastic-recovery Signed-off-by: asolergibert <asolergibert@nvidia.com> # Conflicts: # nemo_rl/algorithms/single_controller.py
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Note
Mirror of #3472, originally opened by @asolergi-nv from the
asolergi-nv/RLfork.The branch was re-pushed to
NVIDIA-NeMo/RLunchanged (same head SHA) so the four partscan be linked as a proper PR stack —
gh stackcannot operate on fork branches.Please review and comment here. The original #3472 is left open for history.
Part 3/4 of #3454
What this fixes
With #3470 and #3471 in place a dead engine no longer wedges rollouts and no longer receives traffic — but the run still dies, and it dies in the worst possible way.
The refit communicator (
model_update_group) is built once at setup over every training and inference rank. A NCCL broadcast requires every rank in the communicator to take part, so when a generation rank is gone the next weight sync blocks forever inside NCCL: no exception, no progress, no CPU burn, and Ray still reporting every actor healthy. This is the failure the whole effort exists to remove.It is also ordered so that the hang comes first:
sync_weightsdoesray.get(futures_train)before theray.get(futures_inference)that would surface the dead actor, so the trainer blocks before anything can report the death.What it does
Reconcile before every refit, not on a death event.
reconcile_communicator()is called at the top of_sync_weights— the one point where 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 is idempotent and converges after a missed or reordered health update, instead of needing replay.Rebuild over the survivors, so training continues on what is left. Trainers are never excluded, so rank 0 stays a trainer and the broadcast root is stable.
Both NCCL transports recover, by different routes.
nccl_reshardneeds more than the plain broadcast: both communicator families rebuilt and the refit plan regenerated.Four decisions:
"Absent" is deliberately not the complement of "serving". The obvious implementation reads
serving_shards(), and it is a bug.SUSPECT(failing probes, not yet condemned) andSTALE(reloaded, holding old weights) are both withheld from traffic while their processes are alive and 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. HenceGenerationFleetMonitor.absent_shards()over{DEAD, RESTARTING, RETIRED}.Rebuilding the communicator is only half a recovery. Every refit dispatch goes through
run_all_workers_*, which walks the whole worker group — so after a loss it kept calling the dead shard's Ray actor and the next refit failed withRayActorError. Both dispatches now address only surviving DP leaders. (This shipped broken in an intermediate commit here and is fixed within this PR; kept as separate commits because the fix is instructive.)nccl_reshardcannot simply be resized. Its bulk path is a mesh-to-mesh redistribute, not a broadcast:prepare_nccl_reshard_refit_infoderives each parameter's destination placements fromgen_world_size. Reusing a plan built for the old fleet does not error — a stale mesh is still a valid mesh — it just has survivors writing the slices the dead shard owned and leaving their own unwritten. So the plan is regenerated, andinit_communicatorand the rebuild share one_build(membership)so that arithmetic exists once and every normal run exercises the rebuild path.Survivors are compacted to contiguous prefixes, not left with a hole. The reshard destination mesh is
torch.arange(offset, offset + num_gpus), so a gap silently misaligns every parameter rather than erroring. The rank arithmetic lives in a pureweight_sync/membership.py, separate from Ray dispatch, because it is the part that must be exactly right and the part that cannot be exercised without ≥3 GPUs.Tests
StatelessProcessGroup.abort()verified on 2×A6000: with a peer SIGKILLed mid-broadcast, a survivor blocked in the collective was released 0.15 s after another thread calledabort().abort()and notdestroy()— NCCL documentsdestroyas an intra-node collective every rank must call or it hangs, which is exactly what a dead rank cannot do.tests/functional/grpo_sc_generation_shard_recovery.shis the gate, runs on both refit transports in the SingleController lane, and self-skips below 3 GPUs rather than passing vacuously.Issues
List issues that this PR closes (syntax):
Usage
# Add a code snippet demonstrating how to use thisBefore your PR is "Ready for review"
Pre checks:
Additional Information