Skip to content

fix(async-grpo): fail fast on generation worker errors #2368

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
jthomson04 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix(async-grpo): fail fast on generation worker errors #2368

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions docs/guides/async-grpo.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,7 @@ grpo:
async_grpo:
enabled: true
max_trajectory_age_steps: 1 # Maximum age, in training steps, for trajectories
max_generation_failures: 0 # Consecutive worker failures to tolerate
in_flight_weight_updates: false # Enable for faster weight synchronization
recompute_kv_cache_after_weight_updates: false # Invalidates kv cache after weight-updates
```
Expand All @@ -67,6 +68,7 @@ grpo:
async_grpo:
enabled: true
max_trajectory_age_steps: 1
max_generation_failures: 0 # Consecutive worker failures to tolerate
in_flight_weight_updates: false # Enable for faster weight synchronization
recompute_kv_cache_after_weight_updates: false # Invalidates kv cache after weight-updates

Expand Down
4 changes: 4 additions & 0 deletions examples/configs/grpo_math_1B.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,10 @@ grpo:
enabled: false # Set to true to enable async training mode
# Max age (in training steps) for trajectories used in training
max_trajectory_age_steps: 1
# Number of generation-worker failures tolerated before aborting.
# 0 (default) = fail on the first worker exception. Increase only
# when transient generation errors are expected and acceptable to drop.
max_generation_failures: 0
in_flight_weight_updates: false # Set to true to enable in-flight weight updates
recompute_kv_cache_after_weight_updates: false # Set to true to recompute kv cache after weight updates

Expand Down
1 change: 1 addition & 0 deletions examples/configs/grpo_math_1B_megatron.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ grpo:
async_grpo:
enabled: false
max_trajectory_age_steps: 1
max_generation_failures: 0

loss_fn:
reference_policy_kl_penalty: 0.01
Expand Down
1 change: 1 addition & 0 deletions examples/configs/grpo_math_8B.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ grpo:
async_grpo:
enabled: false
max_trajectory_age_steps: 1
max_generation_failures: 0

policy:
model_name: "meta-llama/Llama-3.1-8B-Instruct"
Expand Down
1 change: 1 addition & 0 deletions examples/nemo_gym/grpo_nanov3.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@ grpo:
enabled: false # Set to true to enable async training mode
# Max age (in training steps) for trajectories used in training
max_trajectory_age_steps: 1
max_generation_failures: 0

batch_multiplier: 1
use_dynamic_sampling: False
Expand Down
1 change: 1 addition & 0 deletions examples/nemo_gym/grpo_qwen3_30ba3b_thinking_swe1.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@ grpo:
stop_properly_penalty_coef: null
async_grpo:
enabled: true
max_generation_failures: 3
in_flight_weight_updates: true
loss_fn:
kl_input_clamp_value: null
Expand Down
1 change: 1 addition & 0 deletions examples/nemo_gym/grpo_qwen3_30ba3b_thinking_swe2.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@ grpo:
stop_properly_penalty_coef: null
async_grpo:
enabled: true
max_generation_failures: 3
in_flight_weight_updates: true
loss_fn:
kl_input_clamp_value: null
Expand Down
1 change: 1 addition & 0 deletions examples/nemo_gym/grpo_workplace_assistant_nemotron_nano_v2_9b.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,7 @@ grpo:
enabled: false # Set to true to enable async training mode
# Max age (in training steps) for trajectories used in training
max_trajectory_age_steps: 1
max_generation_failures: 0
in_flight_weight_updates: false # Set to true to enable in-flight weight updates
recompute_kv_cache_after_weight_updates: false # Set to true to recompute kv cache after weight updates

Expand Down
1 change: 1 addition & 0 deletions examples/nemo_gym/nemotron-3-super/stage1_rlvr.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,7 @@ grpo:
async_grpo:
enabled: true
max_trajectory_age_steps: 1
max_generation_failures: 3

@terrykong terrykong Aug 5, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

1 action item, 1 follow-up. PR-introduced.

AI-1 — after a worker gives up, the gap-fill that would top the target back up never runs, so the step never completes

Simplified to one thread, this is the collector:

for batch in dataloader:
    if no_target_slot_free():                 # slots = next max_trajectory_age_steps steps
        wait_for_event()                      # <-- PARKS. only set_weight_version() wakes it
    target = reserve_slot()
    needed = 16 - already_buffered(target)    # <-- the gap-fill, every iteration
    spawn_thread(worker, batch[:needed], target)   # returns in ms; worker runs for minutes

def worker(groups, target):
    for attempt in range(4 if gym else 1):    # retries the SAME prompts
        ... buffer each group as it completes ...
        if not missing:
            fail_count = 0; release_slot(target); return
    fail_count += 1                           # ONE increment for the whole batch
    if fail_count > max_generation_failures:
        fatal = "abort"                       # -> trainer's check_health() raises
    release_slot(target)                      # <-- frees the slot, never wakes the loop

Normally the gap-fill handles this. Say 15 of 16 groups buffered and the 16th failed all 4 attempts (max_attempts = 1 + _MAX_NEMO_GYM_STREAM_RETRIES, gym only — native gets 1). Groups are buffered individually as they stream (_enqueue_rollout_group), so the target is simply short by one. On the next loop iteration:

needed = get_trajectories_needed(target=7)   -> 1
groups = next_batch[:1]                      -> fresh prompts, not the failed one
"🎯 Reserved target weight 7 for gap-filling (need 1/16 more trajectories)"

That's the #2651 machinery and it is correct.

It never gets to run. _process_batch returns in ms; the loop's next iteration finds the slot still reserved — at max_trajectory_age_steps: 1 (L61) it is the only slot (_calculate_target_weights) — so _should_pause_for_generation_limits → True → .clear() → parks on .wait(). Minutes later the worker gives up and _release_target frees the slot but does not set _generation_limit_cleared — set only at :114 and in set_weight_version, whose callers are grpo.py:4110 (once, pre-fill) and grpo.py:4687 (post-refit). So needed is never recomputed, the target stays at 15/16, has_complete_batch stays False, and the trainer blocks. The escape at grpo.py:4272-4285 can't fire either: it needs not running, and a parked loop never reaches its finally.

Measured (repro below; verbatim methods, no GPU, deterministic 3/3):

A) generation dead                            batches  count  need  outcome
   current code, mgf=3  <- the 13 recipes           2      2     4  HUNG
   current code, mgf=0  [control]                   2      2     1  ABORTED
   fix A only: .set()                             200      3     4  HUNG
   fix A+B: .set() + .clear() un-gated              6      4     4  ABORTED

B) ONE worker exhausts its retries          steps done
   baseline: no failure at all                    8 / 8   ok
   one exhausted worker                           2       STRANDED
   one exhausted worker + fix A+B                 8 / 8   ok

Net for this recipe: 0 latches fatal before the park matters and aborts cleanly; 3 latches nothing and hangs. The tolerance setting is strictly worse than the default here.

Action — two changes; the first alone is not enough:

A. In _run_rollout_batch_worker's except, after the increment: self._generation_limit_cleared.set().

B. Move .clear() at :324 out of the if self._last_limit_warning_version != self.current_weight_version: guard at :309. That guard logs once per weight version, but .clear() is nested in it while .wait() is not — so after the first warning the event can never be re-armed and the pause degenerates into a spin (row "fix A only": 200 dataloader batches burned, zero GPU work).

Please don't put the set() in _release_target — it also runs on the success path, where parking is correct.

A startup validator on max_generation_failures < max_trajectory_age_steps + 1 isn't a substitute: at max_age: 1 it would reject every non-zero value, deleting the feature for the recipes that want it.

Caveats. (a) The park needs the worker to fail slower than one loop iteration (~tens of ms; several blocking Ray round-trips) — anything with a network or subprocess round trip qualifies. (b) Event check-then-clear still has a lost-wakeup window; it never fired here, but a Condition with a predicate loop would be airtight if you prefer.

Follow-up

A regression test for a worker exhausting its retries at max_trajectory_age_steps: 1. No current test could catch this — they all drive _run_rollout_batch_worker directly and never exercise the loop's pause path.

repro (save as repro.py, run from repo root, ~60 s) 
#!/usr/bin/env -S uv run --script
# /// script
# requires-python = ">=3.11"
# dependencies = []
# ///
"""PR #2368 repro. No GPU/ray/torch. Run from the repo root:  ./repro.py

The six load-bearing methods are extracted VERBATIM from trajectory_collector.py via
ast and exec'd, so the control flow under test is the real code. Modeled only:
  ray.get/.remote -> pass-through | replay buffer -> per-target group counts
  _process_batch  -> reserve + spawn                  (mirrors :390-462)
  worker epilogue -> increment/compare/latch/release  (mirrors :854-895)

Scenario A: generation is dead      -> can the counter ever reach the threshold?
Scenario B: ONE worker exhausts its retries, all others succeed
                                    -> does the gap-fill top the short target back up?
"""
import ast, sys, threading, time, io, contextlib
from contextlib import contextmanager
from pathlib import Path

SRC = Path(sys.argv[1] if len(sys.argv) > 1
           else "nemo_rl/algorithms/async_utils/trajectory_collector.py")
WANT = ["_calculate_target_weights", "_get_next_target_for_generation", "set_weight_version",
        "_should_pause_for_generation_limits", "_collection_loop", "_release_target"]


def extract(path):
    txt = path.read_text(); lines = txt.splitlines()
    cls = next(n for n in ast.walk(ast.parse(txt))
               if isinstance(n, ast.ClassDef) and n.name == "AsyncTrajectoryCollector")
    out = {}
    for n in cls.body:
        if isinstance(n, ast.FunctionDef) and n.name in WANT:
            b = "\n".join(lines[n.lineno - 1:n.end_lineno])
            out[n.name] = "\n".join(l[4:] if l.startswith("    ") else l for l in b.splitlines())
    assert not set(WANT) - set(out), sorted(set(WANT) - set(out))
    return out


class _Ref:
    def __init__(s, f): s._f = f
    def remote(s, *a, **k): return s._f(*a, **k)


class _Ray:
    @staticmethod
    def get(v): return v


class Buf:
    def __init__(s, n): s.n, s.counts, s.last = n, {}, -1
    @property
    def get_last_target_weight_already_generated(s): return _Ref(lambda: s.last)
    @property
    def get_trajectories_needed(s):
        return _Ref(lambda t, num, m=None: max(0, s.n - s.counts.get(t, 0)))
    def complete(s, t): return s.counts.get(t, 0) >= s.n
    def fill(s, t): s.counts[t] = s.n


@contextmanager
def _null(*a, **k): yield


class _T:
    def time(s, *a, **k): return _null()
    def record(s, *a, **k): pass


class Cfg:
    class grpo:
        num_prompts_per_step = 4
        class async_grpo: max_trajectory_age_steps = 1


class C:
    def __init__(s, limit, mode, fix, fail_target=2, delay=0.15):
        s.master_config, s.replay_buffer = Cfg, Buf(Cfg.grpo.num_prompts_per_step)
        s.running, s.dataloader = True, None
        s.data_exhausted = s.collection_failed = False
        s.initial_weight_version = s.current_weight_version = 0
        s._last_limit_warning_version = -1
        s._generating_targets, s._inflight_threads = set(), set()
        s._generation_check_lock = threading.Lock(); s._threads_lock = threading.Lock()
        s._generation_limit_cleared = threading.Event(); s._generation_limit_cleared.set()
        s._manual_pause_cleared = threading.Event(); s._manual_pause_cleared.set()
        s._refit_pause_cleared = threading.Event(); s._refit_pause_cleared.set()
        s._efficiency_timer = _T()
        s._failure_lock = threading.Lock()
        s._failure_count, s._fatal_error_message = 0, None
        s._max_generation_failures = limit
        s.mode, s.fix, s.fail_target, s.delay = mode, fix, fail_target, delay
        s.spawned = s.batches = 0
        s.failed_once = False

    def _process_batch(s, batch):                       # mirrors :390-462
        s.batches += 1
        t = s._get_next_target_for_generation(s.current_weight_version)
        if t is None: return
        s.spawned += 1
        threading.Thread(target=s._worker, args=(t,), daemon=True).start()

    def _worker(s, target):                             # mirrors :854-895
        time.sleep(s.delay)
        fails = s.mode == "all_fail" or (
            s.mode == "one" and target == s.fail_target and not s.failed_once)
        if fails:
            if s.mode != "all_fail": s.failed_once = True
            with s._failure_lock:
                s._failure_count += 1
                if (s._failure_count > s._max_generation_failures
                        and not s._fatal_error_message):
                    s._fatal_error_message = (
                        f"aborting: {s._failure_count} failure(s) exceeded "
                        f"max_generation_failures={s._max_generation_failures}")
        else:
            with s._failure_lock: s._failure_count = 0   # :855 success resets
            s.replay_buffer.fill(target)
        s._release_target(target)                        # verbatim :647
        if s.fix in ("set-only", "set+ungate"):
            s._generation_limit_cleared.set()

    def check_health(s):
        with s._failure_lock: m = s._fatal_error_message
        if m: raise RuntimeError(m)


def install(methods, ungate):
    import typing
    ns = {"ray": _Ray, "_threading": threading, "Optional": typing.Optional,
          "Any": typing.Any, "StatefulDataLoader": object, "BatchedDataDict": dict,
          "DatumSpec": dict}
    for name, src in methods.items():
        if name == "_collection_loop" and ungate:
            # move .clear() OUT of the `log once per weight version` guard
            src = src.replace("                    self._generation_limit_cleared.clear()", "", 1)
            src = src.replace("                # Efficiently wait for generation limits",
                              "                self._generation_limit_cleared.clear()\n"
                              "                # Efficiently wait for generation limits", 1)
        exec(compile("from __future__ import annotations\n" + src, f"<{name}>", "exec"), ns)
        setattr(C, name, ns[name])


def run(limit, mode, fix, steps=8, settle=2.5, train_time=0.35, n_batches=200):
    install(METHODS, ungate=(fix == "set+ungate"))
    c = C(limit, mode, fix); c.dataloader = iter(range(n_batches))
    with contextlib.redirect_stdout(io.StringIO()):
        threading.Thread(target=c._collection_loop, daemon=True).start()
        fatal, done = False, 0
        for _ in range(steps):
            end, advanced = time.time() + settle, False
            while time.time() < end:
                try: c.check_health()
                except RuntimeError: fatal = True; break
                if c.replay_buffer.complete(c.current_weight_version):
                    time.sleep(train_time)
                    c.current_weight_version += 1
                    c.set_weight_version(c.current_weight_version)
                    done += 1; advanced = True; break
                time.sleep(0.02)
            if fatal or not advanced: break
    return dict(count=c._failure_count, need=limit + 1, batches=c.batches,
                steps=done, fatal=fatal)


if __name__ == "__main__":
    METHODS = extract(SRC)
    print(f"\nverbatim from {SRC.name}: {', '.join(sorted(METHODS))}")

    print("\nA) generation dead  (max_age=1) — can the counter reach the threshold?")
    print(f"   {'variant':46} {'batches':>8} {'count':>6} {'need':>5}  outcome")
    for lbl, lim, fix in (("current code, mgf=3  <- the 13 recipes", 3, "none"),
                          ("current code, mgf=0  [control]        ", 0, "none"),
                          ("fix A: .set() only                    ", 3, "set-only"),
                          ("fix A+B: .set() + .clear() un-gated    ", 3, "set+ungate")):
        r = run(lim, "all_fail", fix)
        print(f"   {lbl:46} {r['batches']:>8} {r['count']:>6} {r['need']:>5}  "
              f"{'ABORTED' if r['fatal'] else 'HUNG'}")

    print("\nB) ONE worker exhausts its retries, all others succeed (max_age=1, mgf=3)")
    print(f"   {'variant':46} {'steps done':>11}  outcome")
    res = {}
    for lbl, mode, fix in (("baseline: no failure at all           ", "none", "none"),
                           ("one exhausted worker                  ", "one", "none"),
                           ("one exhausted worker + fix A+B        ", "one", "set+ungate")):
        r = run(3, mode, fix); res[lbl.strip()] = r
        print(f"   {lbl:46} {r['steps']:>11}  "
              f"{'ok' if r['steps'] >= 8 else 'STRANDED at step %d' % r['steps']}")

    ok = (res["one exhausted worker"]["steps"] < 8
          and res["one exhausted worker + fix A+B"]["steps"] >= 8)
    print("\n" + ("REPRODUCED" if ok else "did not reproduce"))
    sys.exit(0 if ok else 1)

@jthomson04 jthomson04 Aug 6, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Addressed in 87f55d4. A failed batch worker now sets _generation_limit_cleared after releasing its failure state so a paused collector can re-evaluate and gap-fill the released target, and the collection loop clears the event on every limit pause rather than only in the log-once branch. I added a regression using the real collection loop and worker path with max_trajectory_age_steps=1; it verifies the tolerated failure wakes the loop, releases the reservation, and starts gap filling without latching a fatal error.

@terrykong terrykong Aug 6, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

fixed in 87f55d4

in_flight_weight_updates: true
recompute_kv_cache_after_weight_updates: false

Expand Down
1 change: 1 addition & 0 deletions examples/nemo_gym/nemotron-3-super/stage2_swe1.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,7 @@ grpo:
async_grpo:
enabled: true
max_trajectory_age_steps: 1
max_generation_failures: 3
in_flight_weight_updates: true
recompute_kv_cache_after_weight_updates: false

Expand Down
1 change: 1 addition & 0 deletions examples/nemo_gym/nemotron-3-super/stage2_swe2.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,7 @@ grpo:
async_grpo:
enabled: true
max_trajectory_age_steps: 1
max_generation_failures: 3
in_flight_weight_updates: true
recompute_kv_cache_after_weight_updates: false

Expand Down
1 change: 1 addition & 0 deletions examples/nemo_gym/nemotron-3-super/stage3_rlhf.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,7 @@ grpo:
async_grpo:
enabled: true
max_trajectory_age_steps: 1
max_generation_failures: 3
in_flight_weight_updates: true
recompute_kv_cache_after_weight_updates: false

Expand Down
1 change: 1 addition & 0 deletions examples/nemo_gym/nemotron-3-ultra/ifbench_teacher.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -91,6 +91,7 @@ grpo:
async_grpo:
enabled: true
max_trajectory_age_steps: 1
max_generation_failures: 3
in_flight_weight_updates: true
recompute_kv_cache_after_weight_updates: false

Expand Down
1 change: 1 addition & 0 deletions examples/nemo_gym/nemotron-3-ultra/mopd.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -106,6 +106,7 @@ grpo:
async_grpo:
enabled: true
max_trajectory_age_steps: 1
max_generation_failures: 3
in_flight_weight_updates: true
recompute_kv_cache_after_weight_updates: false

Expand Down
1 change: 1 addition & 0 deletions examples/nemo_gym/nemotron-3-ultra/reasoning_teacher.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -95,6 +95,7 @@ grpo:
async_grpo:
enabled: true
max_trajectory_age_steps: 1
max_generation_failures: 3
in_flight_weight_updates: true
recompute_kv_cache_after_weight_updates: false

Expand Down
1 change: 1 addition & 0 deletions examples/nemo_gym/nemotron-3-ultra/rlhf_teacher.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -92,6 +92,7 @@ grpo:
async_grpo:
enabled: true
max_trajectory_age_steps: 1
max_generation_failures: 3
in_flight_weight_updates: true
recompute_kv_cache_after_weight_updates: false

Expand Down
1 change: 1 addition & 0 deletions examples/nemo_gym/nemotron-3-ultra/student_rlvr1.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -88,6 +88,7 @@ grpo:
async_grpo:
enabled: true
max_trajectory_age_steps: 1
max_generation_failures: 3
in_flight_weight_updates: true
recompute_kv_cache_after_weight_updates: false

Expand Down
1 change: 1 addition & 0 deletions examples/nemo_gym/nemotron-3-ultra/student_rlvr2.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -89,6 +89,7 @@ grpo:
async_grpo:
enabled: true
max_trajectory_age_steps: 1
max_generation_failures: 3
in_flight_weight_updates: true
recompute_kv_cache_after_weight_updates: false

Expand Down
1 change: 1 addition & 0 deletions examples/nemo_gym/nemotron-3-ultra/swe_teacher.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -113,6 +113,7 @@ grpo:
async_grpo:
enabled: true
max_trajectory_age_steps: 1
max_generation_failures: 3
in_flight_weight_updates: true
recompute_kv_cache_after_weight_updates: false

Expand Down
73 changes: 66 additions & 7 deletions nemo_rl/algorithms/async_utils/trajectory_collector.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -126,6 +126,15 @@ def __init__(
# Timer for efficiency metrics
self._efficiency_timer = ThreadSafeTimer(context={"worker": "collector"})

# Failure tracking for rollout batch workers. _failure_lock guards both
# _failure_count and _fatal_error_message.
self._failure_lock: _threading.Lock = _threading.Lock()
self._failure_count: int = 0
self._fatal_error_message: str | None = None
self._max_generation_failures = (
self.master_config.grpo.async_grpo.max_generation_failures
)

def _calculate_target_weights(self, generation_weight_version: int) -> list[int]:
"""Calculate target weight versions for given generation weight version.

Expand Down Expand Up @@ -291,6 +300,8 @@ def _collection_loop(self):

# Check if generation limits require pausing collection
if self._should_pause_for_generation_limits() and self.running:
self._generation_limit_cleared.clear()

# Only log warning once per weight version
if self._last_limit_warning_version != self.current_weight_version:
max_trajectory_age = (
Expand All @@ -307,8 +318,6 @@ def _collection_loop(self):
)
self._last_limit_warning_version = self.current_weight_version

self._generation_limit_cleared.clear() # Clear the event to pause

# Efficiently wait for generation limits to be cleared (no polling!)
with self._efficiency_timer.time("idle/generation_limit_pause"):
self._generation_limit_cleared.wait()
Expand Down Expand Up @@ -472,6 +481,21 @@ def _run_rollout_batch() -> None:
def get_weight_version(self) -> int:
return self.current_weight_version

def check_health(self) -> None:
"""Raise the stored fatal worker error, if any.

Called by the trainer between sampling iterations. When a generation
worker has recorded a fatal failure (consecutive count exceeded
max_generation_failures), this raises it so the training job dies
instead of stalling on an empty replay buffer. Safe to call
repeatedly: returns silently when no fatal error is set, and raises
every time once one is.
"""
with self._failure_lock:
error_message = self._fatal_error_message
if error_message is not None:
raise RuntimeError(error_message)

def pause(self) -> None:
"""Pause trajectory collection."""
self._manual_pause_cleared.clear() # Signal collection to pause
Expand Down Expand Up @@ -814,6 +838,7 @@ async def _run_rollout_batch_worker(
) -> None:
"""Own one target reservation while collecting its rollout batch."""
worker_start = time.perf_counter()
wake_generation_limits_after_cleanup = False
try:
await self._collect_rollout_batch(
repeated_batch=repeated_batch,
Expand All @@ -822,22 +847,56 @@ async def _run_rollout_batch_worker(
num_generations=num_generations,
use_nemo_gym=use_nemo_gym,
)
with self._failure_lock:
if self._fatal_error_message is None:
self._failure_count = 0
except Exception as error:
if not self.running:
return

self._efficiency_timer.record(
"wasted/failed_trajectory", time.perf_counter() - worker_start
)
backend = "NeMo-Gym" if use_nemo_gym else "native"
print(
f"❌ Error in {backend} batch worker "
f"(target_weight={target_weight_version}): {error}"
)
import traceback

traceback.print_exc()
failure_traceback = traceback.format_exc()
with self._failure_lock:
self._failure_count += 1

@jthomson04 jthomson04 Jul 29, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The counter is process-lifetime, so max_generation_failures: 5 means "5 failures over the entire run", not "5 in a row". On a 10k-step run any nonzero budget is exhausted early and the setting collapses back to fail-fast — which makes it hard to use for its stated purpose ("transient generation errors ... acceptable to drop").

Suggest resetting on success so the threshold counts consecutive failures:

try:
    await self._collect_rollout_batch(...)
    with self._failure_lock:
        self._failure_count = 0
except Exception as error:
    ...

and updating the AsyncGRPOConfig comment accordingly. That makes the knob discriminate transient flakes from a genuinely broken backend, which is the distinction this PR is drawing.

@terrykong terrykong Aug 5, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

1 action item. PR-introduced, independent of the other finding. Confirmed with a repro (below).

AI-1 — normal dataloader exhaustion is counted as a generation failure

This isn't a rare race — it's the default end-of-data path. The loop spawns a worker for the last batch, immediately exhausts the dataloader, and _collection_loop's finally sets self.running = False while that worker is still generating. The worker then hits raise RuntimeError("Trajectory collection stopped before enqueue completed") — note it drops the group without even attempting the add — which surfaces as batch_error from _collect_rollout_batch and lands in this counting block as an ordinary generation failure.

With max_generation_failures: 0 — what every non-agentic config resolves to — that latches a fatal on the first such worker.

Repro output (_collection_loop and _enqueue_rollout_group extracted verbatim; no GPU/ray/torch):

control  (running=True)          : buffered OK
after exhaustion                 : running=False, data_exhausted=True
in-flight worker enqueue         : RuntimeError('Trajectory collection stopped before enqueue completed')
_failure_count                   : 1 (max_generation_failures=0)

What the user sees at the next check_health():
   AsyncTrajectoryCollector aborting: 1 batch-worker failure(s) exceeded
   max_generation_failures=0. Last failure in native batch worker ...
   RuntimeError('Trajectory collection stopped before enqueue completed')

What they should have seen (grpo.py:4272-4285):
   Trajectory collector stopped: dataloader exhausted ...
   Increase data.train.max_num_epochs or use a larger dataset.

Cost: the run was ending anyway, so this isn't a corrupted-training issue — but the user gets an error blaming generation for what is actually "your dataset ran out", and the job dies at the next check_health() instead of training the steps still sitting in the buffer.

Action: don't count shutdown-induced enqueue aborts. Either guard the counting block with if self.running:, or raise a distinct sentinel from _enqueue_rollout_group and re-raise it past the counter.

repro (save as repro_exhaust.py, run from repo root) 
#!/usr/bin/env -S uv run --script
# /// script
# requires-python = ">=3.11"
# dependencies = []
# ///
"""PR #2368: does NORMAL dataloader exhaustion get counted as a generation failure?

Run from the repo root:  ./repro_exhaust.py

_collection_loop and _enqueue_rollout_group are extracted VERBATIM from
trajectory_collector.py via ast and exec'd, so the running=False -> raise chain is
the real code. Modeled: ray.get/.remote pass-through, replay buffer, and the
worker's except block (mirrors :857-879).

Sequence under test:
  1. the dataloader runs out -> _collection_loop's `finally` sets self.running = False   (:349)
  2. a still-in-flight worker calls _enqueue_rollout_group -> `while self.running:` is
     False -> raise RuntimeError("Trajectory collection stopped before enqueue completed")
  3. that lands in the batch-worker except block -> _failure_count += 1                   (:866)
  4. with max_generation_failures=0 that latches a fatal, and check_health() raises it
     instead of the intended "dataloader exhausted / increase max_num_epochs" message
"""
import ast, sys, asyncio, threading, time, io, contextlib
from contextlib import contextmanager
from pathlib import Path

SRC = Path(sys.argv[1] if len(sys.argv) > 1
           else "nemo_rl/algorithms/async_utils/trajectory_collector.py")
WANT = ["_collection_loop", "_enqueue_rollout_group", "_should_pause_for_generation_limits",
        "_calculate_target_weights", "_get_next_target_for_generation", "_release_target"]


def extract(path):
    txt = path.read_text(); lines = txt.splitlines()
    cls = next(n for n in ast.walk(ast.parse(txt))
               if isinstance(n, ast.ClassDef) and n.name == "AsyncTrajectoryCollector")
    out = {}
    for n in cls.body:
        if isinstance(n, (ast.FunctionDef, ast.AsyncFunctionDef)) and n.name in WANT:
            b = "\n".join(lines[n.lineno - 1:n.end_lineno])
            out[n.name] = "\n".join(l[4:] if l.startswith("    ") else l for l in b.splitlines())
    assert not set(WANT) - set(out), sorted(set(WANT) - set(out))
    return out


class _Ref:
    def __init__(s, f): s._f = f
    def remote(s, *a, **k): return s._f(*a, **k)


class _AsyncRef:
    """replay_buffer.add.remote(...) is awaited in _enqueue_rollout_group."""
    def __init__(s, f): s._f = f
    def remote(s, *a, **k):
        async def go(): return s._f(*a, **k)
        return go()


class _Ray:
    @staticmethod
    def get(v): return v


class Buf:
    def __init__(s, n): s.n, s.counts, s.last = n, {}, -1
    @property
    def get_last_target_weight_already_generated(s): return _Ref(lambda: s.last)
    @property
    def get_trajectories_needed(s):
        return _Ref(lambda t, num, m=None: max(0, s.n - s.counts.get(t, 0)))
    @property
    def add(s):
        def _add(group, gwv, twv):
            s.counts[twv] = s.counts.get(twv, 0) + 1
            return "success"
        return _AsyncRef(_add)


class _FakeBatch(dict):
    def to(s, dev): return s


class _Result:
    def __init__(s):
        s.final_batch = _FakeBatch()
        s.rollout_metrics = {}
        s.group_index = 0
        s.task_index = None


@contextmanager
def _null(*a, **k): yield


class _T:
    def time(s, *a, **k): return _null()
    def record(s, *a, **k): pass


class Cfg:
    class grpo:
        num_prompts_per_step = 4
        class async_grpo: max_trajectory_age_steps = 1


class C:
    def __init__(s, limit):
        s.master_config, s.replay_buffer = Cfg, Buf(4)
        s.running, s.dataloader = True, None
        s.data_exhausted = s.collection_failed = False
        s.initial_weight_version = s.current_weight_version = 0
        s._last_limit_warning_version = -1
        s._generating_targets, s._inflight_threads = set(), set()
        s._generation_check_lock = threading.Lock(); s._threads_lock = threading.Lock()
        s._generation_limit_cleared = threading.Event(); s._generation_limit_cleared.set()
        s._manual_pause_cleared = threading.Event(); s._manual_pause_cleared.set()
        s._refit_pause_cleared = threading.Event(); s._refit_pause_cleared.set()
        s._efficiency_timer = _T()
        s._has_distillation_teachers = False
        s.tokenizer = None
        s._failure_lock = threading.Lock()
        s._failure_count, s._fatal_error_message = 0, None
        s._max_generation_failures = limit

    def _process_batch(s, batch):
        pass  # not under test here

    # mirrors _run_rollout_batch_worker's except block, :857-879
    def count_worker_failure(s, error, backend="native", gwv=0, twv=0):
        with s._failure_lock:
            s._failure_count += 1
            n = s._failure_count
            fatal = n > s._max_generation_failures
            if fatal and s._fatal_error_message is None:
                s._fatal_error_message = (
                    f"AsyncTrajectoryCollector aborting: {n} batch-worker failure(s) "
                    f"exceeded max_generation_failures={s._max_generation_failures}. "
                    f"Last failure in {backend} batch worker for generation_weight={gwv}, "
                    f"target_weight={twv}: {error!r}")

    def check_health(s):
        with s._failure_lock: m = s._fatal_error_message
        if m: raise RuntimeError(m)


def install(methods):
    import typing
    ns = {"ray": _Ray, "_threading": threading, "asyncio": asyncio, "time": time,
          "Optional": typing.Optional, "Any": typing.Any, "StatefulDataLoader": object,
          "BatchedDataDict": dict, "DatumSpec": dict,
          "NEMO_GYM_TASK_INDEX_KEY": "_ng_task_index",
          "_REPLAY_BUFFER_MAX_BACKOFF_SECONDS": 0.5}
    for name, src in methods.items():
        exec(compile("from __future__ import annotations\n" + src, f"<{name}>", "exec"), ns)
        setattr(C, name, ns[name])


def enqueue(c):
    """Run the verbatim _enqueue_rollout_group exactly as a worker would."""
    return asyncio.run(c._enqueue_rollout_group(
        rollout_result=_Result(), generation_weight_version=0, target_weight_version=1,
        expected_prompt_groups=4, buffered_group_indices=set(),
        collection_started_at=time.perf_counter()))


if __name__ == "__main__":
    install(extract(SRC))
    print(f"\nverbatim: _collection_loop, _enqueue_rollout_group (+4 helpers) "
          f"from {SRC.name}\n")

    # --- control: collector still running -> the group buffers fine ---
    c = C(limit=0); c.dataloader = iter([])
    with contextlib.redirect_stdout(io.StringIO()):
        try:
            enqueue(c); control = "buffered OK"
        except Exception as e:
            control = f"raised {e!r}"
    print(f"control  (running=True)          : {control}")

    # --- the real sequence: let the dataloader exhaust, THEN enqueue ---
    c = C(limit=0); c.dataloader = iter(range(2))
    with contextlib.redirect_stdout(io.StringIO()):
        c._collection_loop()                    # VERBATIM: exhausts, finally sets running=False
    print(f"after exhaustion                 : running={c.running}, "
          f"data_exhausted={c.data_exhausted}")

    with contextlib.redirect_stdout(io.StringIO()):
        try:
            enqueue(c); outcome = None
        except Exception as e:
            outcome = e
    print(f"in-flight worker enqueue         : {outcome!r}")

    if outcome is None:
        print("\nNOT reproduced: enqueue did not raise"); sys.exit(1)

    c.count_worker_failure(outcome)             # mirrors :857-879
    print(f"_failure_count                   : {c._failure_count} "
          f"(max_generation_failures={c._max_generation_failures})")

    try:
        c.check_health(); print("\nNOT reproduced: check_health stayed silent"); sys.exit(1)
    except RuntimeError as e:
        print("\nWhat the user sees at the next check_health():\n")
        print("   " + str(e)[:300])
        print("\nWhat they should have seen (grpo.py:4272-4285):\n")
        print("   RuntimeError: Trajectory collector stopped: dataloader exhausted ...")
        print("   The dataset ran out of data before training could start.")
        print("   Increase data.train.max_num_epochs or use a larger dataset.")
        print("\nREPRODUCED: normal exhaustion is reported as a generation failure")
        sys.exit(0)

@jthomson04 jthomson04 Aug 6, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Addressed in 87f55d4. _run_rollout_batch_worker now skips failure accounting when the collector has already stopped, while the existing finally block still releases the target and removes the worker thread. The new shutdown regression simulates the enqueue-side error after normal exhaustion and verifies the failure count/fatal diagnostic remain untouched and cleanup completes.

@terrykong terrykong Aug 6, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

fixed in 87f55d4

failure_count = self._failure_count
failure_limit = self._max_generation_failures
is_fatal = failure_count > failure_limit
if is_fatal and self._fatal_error_message is None:
self._fatal_error_message = (
"AsyncTrajectoryCollector aborting: "
f"{failure_count} batch-worker failure(s) exceeded "
f"max_generation_failures={failure_limit}. "
f"Last failure in {backend} batch worker for "
f"generation_weight={generation_weight_version}, "
f"target_weight={target_weight_version}: {error!r}\n"
f"Worker traceback:\n{failure_traceback}"
)
wake_generation_limits_after_cleanup = True
print(
f"[AsyncTrajectoryCollector] {backend} batch worker FAILED "
f"(failure {failure_count}, tolerating {failure_limit}) "
f"generation_weight={generation_weight_version} "
f"target_weight={target_weight_version}\n{failure_traceback}",
flush=True,
)
if is_fatal:
print(
f"[AsyncTrajectoryCollector] FATAL: failure count "
f"{failure_count} exceeds threshold {failure_limit}; trainer "
"will be notified on the next check_health() call.",
flush=True,
)
finally:
self._release_target(target_weight_version)
with self._threads_lock:
self._inflight_threads.discard(_threading.current_thread())
if wake_generation_limits_after_cleanup:
self._generation_limit_cleared.set()

@staticmethod
def _build_task_index_map(
Expand Down
12 changes: 11 additions & 1 deletion nemo_rl/algorithms/grpo.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -183,6 +183,10 @@ class AsyncGRPOConfig(BaseModel, extra="allow"):
# async replay buffer. Trajectories older than this are excluded during
# sampling; buffer sizing also scales with this value.
max_trajectory_age_steps: int = 1
# Generation-worker failures tolerated before the AsyncTrajectoryCollector
# aborts the run. A successful batch worker resets the count.
# 0 makes the very first worker exception fatal.
max_generation_failures: int = 0
# Does the weight synchronization as soon as the training is done
# without waiting for the pending generations to finish.
in_flight_weight_updates: bool = False
Expand Down Expand Up @@ -3986,6 +3990,7 @@ def async_grpo_train(
assert master_config.loss_fn.use_importance_sampling_correction, (
"Importance sampling correction must be enabled for async GRPO for good convergence due to off-policy samples!"
)
max_generation_failures = master_config.grpo.async_grpo.max_generation_failures

if router_replay_enabled(master_config.policy) and (
master_config.data_plane or {}
Expand Down Expand Up @@ -4176,7 +4181,9 @@ def async_grpo_train(
print("📦 Started continuous background trajectory collection")

print(
f"🚀 Starting async GRPO training with buffer_size={optimal_buffer_size}, max_age={max_trajectory_age_steps} steps"
f"🚀 Starting async GRPO training with buffer_size={optimal_buffer_size}, "
f"max_age={max_trajectory_age_steps} steps, "
f"max_generation_failures={max_generation_failures}"
)

print("⏳ Preparing policy generation for training...", flush=True)
Expand Down Expand Up @@ -4281,6 +4288,7 @@ def async_grpo_train(
wait_iterations = 0
while True:
buffer_size_current = ray.get(replay_buffer.size.remote())
ray.get(trajectory_collector.check_health.remote())
current_step_ready = ray.get(
replay_buffer.has_complete_batch.remote(
step, num_prompts_per_step, max_trajectory_age_steps
Expand Down Expand Up @@ -4357,6 +4365,7 @@ def async_grpo_train(
# Main training loop
try:
while step < master_config.grpo.max_num_steps:
ray.get(trajectory_collector.check_health.remote())
refit_metrics: dict[str, float] = {}
early_stop_message: Optional[str] = None
print(
Expand Down Expand Up @@ -5161,6 +5170,7 @@ def async_grpo_train(
import traceback

traceback.print_exc()
raise

finally:
# Finalize any pending async checkpoint before tearing down workers.
Expand Down
1 change: 1 addition & 0 deletions research/template_project/configs/grpo_math_1B.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,7 @@ grpo:
enabled: false # Set to true to enable async training mode
# Max age (in training steps) for trajectories used in training
max_trajectory_age_steps: 1
max_generation_failures: 0
in_flight_weight_updates: false # Set to true to enable in-flight weight updates
recompute_kv_cache_after_weight_updates: false # Set to true to recompute kv cache after weight updates

Expand Down
Loading
Loading