diff --git a/docs/guides/async-grpo.md b/docs/guides/async-grpo.md index ef6ec5db76..64838fab62 100644 --- a/docs/guides/async-grpo.md +++ b/docs/guides/async-grpo.md @@ -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 ``` @@ -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 diff --git a/examples/configs/grpo_math_1B.yaml b/examples/configs/grpo_math_1B.yaml index 062abddaa4..4416a0cc01 100644 --- a/examples/configs/grpo_math_1B.yaml +++ b/examples/configs/grpo_math_1B.yaml @@ -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 diff --git a/examples/configs/grpo_math_1B_megatron.yaml b/examples/configs/grpo_math_1B_megatron.yaml index f56534cabe..c6b7243ee9 100644 --- a/examples/configs/grpo_math_1B_megatron.yaml +++ b/examples/configs/grpo_math_1B_megatron.yaml @@ -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 diff --git a/examples/configs/grpo_math_8B.yaml b/examples/configs/grpo_math_8B.yaml index a483a6360e..c1ce8a8700 100644 --- a/examples/configs/grpo_math_8B.yaml +++ b/examples/configs/grpo_math_8B.yaml @@ -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" diff --git a/examples/nemo_gym/grpo_nanov3.yaml b/examples/nemo_gym/grpo_nanov3.yaml index 88d38dfb7b..2d99df87b2 100644 --- a/examples/nemo_gym/grpo_nanov3.yaml +++ b/examples/nemo_gym/grpo_nanov3.yaml @@ -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 diff --git a/examples/nemo_gym/grpo_qwen3_30ba3b_thinking_swe1.yaml b/examples/nemo_gym/grpo_qwen3_30ba3b_thinking_swe1.yaml index c19fb351b6..87ea2e9ea2 100644 --- a/examples/nemo_gym/grpo_qwen3_30ba3b_thinking_swe1.yaml +++ b/examples/nemo_gym/grpo_qwen3_30ba3b_thinking_swe1.yaml @@ -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 diff --git a/examples/nemo_gym/grpo_qwen3_30ba3b_thinking_swe2.yaml b/examples/nemo_gym/grpo_qwen3_30ba3b_thinking_swe2.yaml index 449b873e5f..3851a28e0b 100644 --- a/examples/nemo_gym/grpo_qwen3_30ba3b_thinking_swe2.yaml +++ b/examples/nemo_gym/grpo_qwen3_30ba3b_thinking_swe2.yaml @@ -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 diff --git a/examples/nemo_gym/grpo_workplace_assistant_nemotron_nano_v2_9b.yaml b/examples/nemo_gym/grpo_workplace_assistant_nemotron_nano_v2_9b.yaml index b3641dd543..afb8b7a802 100644 --- a/examples/nemo_gym/grpo_workplace_assistant_nemotron_nano_v2_9b.yaml +++ b/examples/nemo_gym/grpo_workplace_assistant_nemotron_nano_v2_9b.yaml @@ -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 diff --git a/examples/nemo_gym/nemotron-3-super/stage1_rlvr.yaml b/examples/nemo_gym/nemotron-3-super/stage1_rlvr.yaml index 03317acd30..cec953d69d 100644 --- a/examples/nemo_gym/nemotron-3-super/stage1_rlvr.yaml +++ b/examples/nemo_gym/nemotron-3-super/stage1_rlvr.yaml @@ -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 diff --git a/examples/nemo_gym/nemotron-3-super/stage2_swe1.yaml b/examples/nemo_gym/nemotron-3-super/stage2_swe1.yaml index 8a8d52ca04..1c458c961a 100644 --- a/examples/nemo_gym/nemotron-3-super/stage2_swe1.yaml +++ b/examples/nemo_gym/nemotron-3-super/stage2_swe1.yaml @@ -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 diff --git a/examples/nemo_gym/nemotron-3-super/stage2_swe2.yaml b/examples/nemo_gym/nemotron-3-super/stage2_swe2.yaml index 4c5284f5aa..e05ca0256c 100644 --- a/examples/nemo_gym/nemotron-3-super/stage2_swe2.yaml +++ b/examples/nemo_gym/nemotron-3-super/stage2_swe2.yaml @@ -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 diff --git a/examples/nemo_gym/nemotron-3-super/stage3_rlhf.yaml b/examples/nemo_gym/nemotron-3-super/stage3_rlhf.yaml index f2d2b38930..b549e12ea0 100644 --- a/examples/nemo_gym/nemotron-3-super/stage3_rlhf.yaml +++ b/examples/nemo_gym/nemotron-3-super/stage3_rlhf.yaml @@ -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 diff --git a/examples/nemo_gym/nemotron-3-ultra/ifbench_teacher.yaml b/examples/nemo_gym/nemotron-3-ultra/ifbench_teacher.yaml index 18dc57404a..0c11f1ef9c 100644 --- a/examples/nemo_gym/nemotron-3-ultra/ifbench_teacher.yaml +++ b/examples/nemo_gym/nemotron-3-ultra/ifbench_teacher.yaml @@ -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 diff --git a/examples/nemo_gym/nemotron-3-ultra/mopd.yaml b/examples/nemo_gym/nemotron-3-ultra/mopd.yaml index 203bf03e83..2e6d36b4ee 100644 --- a/examples/nemo_gym/nemotron-3-ultra/mopd.yaml +++ b/examples/nemo_gym/nemotron-3-ultra/mopd.yaml @@ -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 diff --git a/examples/nemo_gym/nemotron-3-ultra/reasoning_teacher.yaml b/examples/nemo_gym/nemotron-3-ultra/reasoning_teacher.yaml index e62f261e86..76e8464a00 100644 --- a/examples/nemo_gym/nemotron-3-ultra/reasoning_teacher.yaml +++ b/examples/nemo_gym/nemotron-3-ultra/reasoning_teacher.yaml @@ -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 diff --git a/examples/nemo_gym/nemotron-3-ultra/rlhf_teacher.yaml b/examples/nemo_gym/nemotron-3-ultra/rlhf_teacher.yaml index d757cde7a0..4e5d1b84ab 100644 --- a/examples/nemo_gym/nemotron-3-ultra/rlhf_teacher.yaml +++ b/examples/nemo_gym/nemotron-3-ultra/rlhf_teacher.yaml @@ -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 diff --git a/examples/nemo_gym/nemotron-3-ultra/student_rlvr1.yaml b/examples/nemo_gym/nemotron-3-ultra/student_rlvr1.yaml index f6c3b308ad..fc22204236 100644 --- a/examples/nemo_gym/nemotron-3-ultra/student_rlvr1.yaml +++ b/examples/nemo_gym/nemotron-3-ultra/student_rlvr1.yaml @@ -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 diff --git a/examples/nemo_gym/nemotron-3-ultra/student_rlvr2.yaml b/examples/nemo_gym/nemotron-3-ultra/student_rlvr2.yaml index 262d1679cd..4cb3c70d78 100644 --- a/examples/nemo_gym/nemotron-3-ultra/student_rlvr2.yaml +++ b/examples/nemo_gym/nemotron-3-ultra/student_rlvr2.yaml @@ -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 diff --git a/examples/nemo_gym/nemotron-3-ultra/swe_teacher.yaml b/examples/nemo_gym/nemotron-3-ultra/swe_teacher.yaml index 75224b0e55..1d2b8874c6 100644 --- a/examples/nemo_gym/nemotron-3-ultra/swe_teacher.yaml +++ b/examples/nemo_gym/nemotron-3-ultra/swe_teacher.yaml @@ -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 diff --git a/nemo_rl/algorithms/async_utils/trajectory_collector.py b/nemo_rl/algorithms/async_utils/trajectory_collector.py index ca57644afa..6bdf762efa 100644 --- a/nemo_rl/algorithms/async_utils/trajectory_collector.py +++ b/nemo_rl/algorithms/async_utils/trajectory_collector.py @@ -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. @@ -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 = ( @@ -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() @@ -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 @@ -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, @@ -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 + 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( diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index 2cc604ba2c..7df95255e6 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -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 @@ -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 {} @@ -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) @@ -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 @@ -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( @@ -5161,6 +5170,7 @@ def async_grpo_train( import traceback traceback.print_exc() + raise finally: # Finalize any pending async checkpoint before tearing down workers. diff --git a/research/template_project/configs/grpo_math_1B.yaml b/research/template_project/configs/grpo_math_1B.yaml index 9d7b2b8629..e77057620e 100644 --- a/research/template_project/configs/grpo_math_1B.yaml +++ b/research/template_project/configs/grpo_math_1B.yaml @@ -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 diff --git a/tests/unit/algorithms/test_async_utils.py b/tests/unit/algorithms/test_async_utils.py index 1e274504fb..d676215b24 100644 --- a/tests/unit/algorithms/test_async_utils.py +++ b/tests/unit/algorithms/test_async_utils.py @@ -16,6 +16,7 @@ import os import tempfile import threading +import time import unittest.mock as mock from types import SimpleNamespace @@ -1107,7 +1108,10 @@ class TestAsyncTrajectoryCollector: """Test cases for AsyncTrajectoryCollector.""" def create_local_collector( - self, replay_buffer=None, next_nemo_gym_task_index: int = 0 + self, + replay_buffer=None, + next_nemo_gym_task_index: int = 0, + max_generation_failures: int = 0, ): """Create a non-Ray collector instance for unit-testing local state.""" collector_cls = AsyncTrajectoryCollector.__ray_metadata__.modified_class @@ -1115,6 +1119,7 @@ def create_local_collector( mock_tokenizer = mock.MagicMock() task_to_env = {} master_config = self.create_mock_config() + master_config.grpo.async_grpo.max_generation_failures = max_generation_failures if replay_buffer is None: replay_buffer = mock.MagicMock() @@ -1205,7 +1210,10 @@ def create_mock_config(self) -> MasterConfig: num_prompts_per_step=2, num_generations_per_prompt=3, max_rollout_turns=1, - async_grpo=AsyncGRPOConfig.model_construct(max_trajectory_age_steps=2), + async_grpo=AsyncGRPOConfig.model_construct( + max_trajectory_age_steps=2, + max_generation_failures=0, + ), ), policy={ "max_total_sequence_length": 512, @@ -1898,6 +1906,186 @@ def test_dataloader_state_retrieval(self): ray.kill(buffer) ray.kill(mock_env) + @pytest.mark.parametrize("max_generation_failures", [0, 2]) + def test_batch_worker_failure_surfaces_after_threshold( + self, monkeypatch, max_generation_failures + ): + """Consecutive batch-worker failures become sticky past the limit.""" + collector = self.create_local_collector( + max_generation_failures=max_generation_failures + ) + collector.running = True + target_weight = 7 + + outcomes = [] + if max_generation_failures > 0: + outcomes.append(ValueError("pre-reset failure")) + outcomes.append(None) + outcomes.extend( + ValueError(f"backend failed {failure_index}") + for failure_index in range(max_generation_failures + 2) + ) + + async def collect_rollout_batch(**kwargs): + outcome = outcomes.pop(0) + if outcome is not None: + raise outcome + + monkeypatch.setattr(collector, "_collect_rollout_batch", collect_rollout_batch) + + def run_worker(*, expect_generation_wake): + collector._generation_limit_cleared.clear() + collector._generating_targets.add(target_weight) + asyncio.run( + collector._run_rollout_batch_worker( + repeated_batch=None, + generation_weight_version=4, + target_weight_version=target_weight, + num_generations=1, + use_nemo_gym=False, + ) + ) + assert target_weight not in collector._generating_targets + assert ( + collector._generation_limit_cleared.is_set() is expect_generation_wake + ) + + collector.check_health() + + if max_generation_failures > 0: + run_worker(expect_generation_wake=True) + assert collector._failure_count == 1 + collector.check_health() + + run_worker(expect_generation_wake=False) + assert collector._failure_count == 0 + collector.check_health() + + for failure_index in range(max_generation_failures + 1): + run_worker(expect_generation_wake=True) + if failure_index < max_generation_failures: + collector.check_health() + + expected_count = max_generation_failures + 1 + with pytest.raises(RuntimeError) as exc_info: + collector.check_health() + + error_message = str(exc_info.value) + assert f"{expected_count} batch-worker failure(s)" in error_message + assert f"max_generation_failures={max_generation_failures}" in error_message + assert "native batch worker" in error_message + assert "generation_weight=4" in error_message + assert "target_weight=7" in error_message + assert ( + f"ValueError('backend failed {max_generation_failures}')" in error_message + ) + assert "Worker traceback:" in error_message + assert "Traceback (most recent call last):" in error_message + + first_fatal_error = exc_info.value + run_worker(expect_generation_wake=True) + assert collector._failure_count == expected_count + 1 + + with pytest.raises(RuntimeError, match="target_weight=7") as repeated_exc_info: + collector.check_health() + + assert repeated_exc_info.value is not first_fatal_error + assert str(repeated_exc_info.value) == error_message + + def test_tolerated_worker_failure_wakes_gap_fill_pause(self, monkeypatch): + """A failed worker releases and wakes a max-age-one target for gap fill.""" + collector = self.create_local_collector(max_generation_failures=3) + collector.master_config.grpo.async_grpo.max_trajectory_age_steps = 1 + collector.current_weight_version = 4 + collector.running = True + collector.dataloader = [{"batch": 0}] + target_weight = 5 + collector._generating_targets.add(target_weight) + collector._last_limit_warning_version = collector.current_weight_version + + release_target = collector._release_target + + def release_before_wake(target_weight_version): + assert not collector._generation_limit_cleared.is_set() + release_target(target_weight_version) + + monkeypatch.setattr(collector, "_release_target", release_before_wake) + + monkeypatch.setattr( + collector, + "_should_pause_for_generation_limits", + lambda: target_weight in collector._generating_targets, + ) + + gap_fill_started = threading.Event() + + def process_gap_fill(batch): + assert target_weight not in collector._generating_targets + gap_fill_started.set() + + monkeypatch.setattr(collector, "_process_batch", process_gap_fill) + + async def fail_rollout_batch(**kwargs): + raise ValueError("worker exhausted retries") + + monkeypatch.setattr(collector, "_collect_rollout_batch", fail_rollout_batch) + + collection_thread = threading.Thread(target=collector._collection_loop) + collection_thread.start() + deadline = time.monotonic() + 1 + while collector._generation_limit_cleared.is_set(): + assert time.monotonic() < deadline, "collection loop did not enter pause" + time.sleep(0.01) + assert not gap_fill_started.is_set() + + asyncio.run( + collector._run_rollout_batch_worker( + repeated_batch=None, + generation_weight_version=4, + target_weight_version=target_weight, + num_generations=1, + use_nemo_gym=False, + ) + ) + + assert gap_fill_started.wait(timeout=1) + collection_thread.join(timeout=1) + assert not collection_thread.is_alive() + assert collector._failure_count == 1 + collector.check_health() + + def test_worker_shutdown_error_is_not_counted(self, monkeypatch): + """An in-flight worker stopping after exhaustion is not a generation failure.""" + collector = self.create_local_collector(max_generation_failures=0) + collector.running = False + collector.data_exhausted = True + target_weight = 7 + collector._generating_targets.add(target_weight) + collector._inflight_threads.add(threading.current_thread()) + collector._generation_limit_cleared.clear() + + async def fail_during_shutdown(**kwargs): + raise RuntimeError("Trajectory collection stopped before enqueue completed") + + monkeypatch.setattr(collector, "_collect_rollout_batch", fail_during_shutdown) + + asyncio.run( + collector._run_rollout_batch_worker( + repeated_batch=None, + generation_weight_version=4, + target_weight_version=target_weight, + num_generations=1, + use_nemo_gym=False, + ) + ) + + assert collector._failure_count == 0 + assert collector._fatal_error_message is None + assert not collector._generation_limit_cleared.is_set() + assert target_weight not in collector._generating_targets + assert threading.current_thread() not in collector._inflight_threads + collector.check_health() + class TestAsyncUtilsIntegration: """Integration tests for async utilities working together.""" @@ -1909,7 +2097,10 @@ def create_mock_config(self) -> MasterConfig: num_prompts_per_step=2, num_generations_per_prompt=2, max_rollout_turns=1, - async_grpo=AsyncGRPOConfig.model_construct(max_trajectory_age_steps=1), + async_grpo=AsyncGRPOConfig.model_construct( + max_trajectory_age_steps=1, + max_generation_failures=0, + ), ), policy={ "max_total_sequence_length": 512, diff --git a/tests/unit/algorithms/test_grpo.py b/tests/unit/algorithms/test_grpo.py index 479e547991..1516f6d938 100644 --- a/tests/unit/algorithms/test_grpo.py +++ b/tests/unit/algorithms/test_grpo.py @@ -290,6 +290,7 @@ def val_iter(self): async_grpo=AsyncGRPOConfig.model_construct( enabled=False, max_trajectory_age_steps=1, + max_generation_failures=0, ), seq_logprob_error_threshold=None, adv_estimator=AdvEstimatorConfig.model_construct( @@ -406,6 +407,7 @@ def test_grpo_config_nested_defaults_are_populated(): assert isinstance(first.reward_shaping, RewardShapingConfig) assert isinstance(first.reward_scaling, RewardScalingConfig) assert first.async_grpo.enabled is False + assert first.async_grpo.max_generation_failures == 0 assert first.adv_estimator.use_leave_one_out_baseline is True assert first.adv_estimator.normalize_rewards is True assert first.adv_estimator.minus_baseline is True @@ -881,11 +883,17 @@ def has_complete_batch(self): class StubAsyncTrajectoryCollector: - """Non-Ray stub of AsyncTrajectoryCollector for unit testing + """Non-Ray stub of AsyncTrajectoryCollector for unit testing. - Each method is a property that returns a MagicMock with a 'remote' attribute. + Actor methods expose MagicMocks with a ``remote`` attribute. """ + def __init__(self, health_side_effect=None): + self.check_health = MagicMock() + self.check_health.remote = MagicMock( + return_value=None, side_effect=health_side_effect + ) + @property def start_collection(self): """Start collection - returns a remote-callable mock""" @@ -973,7 +981,10 @@ def get_rollouts_state(self): def mock_async_grpo_infrastructure( - mock_batch, mock_rollout_metrics, seq_logprob_error_result=None + mock_batch, + mock_rollout_metrics, + seq_logprob_error_result=None, + collector_health_side_effect=None, ): """ Context manager that mocks all async GRPO infrastructure (Ray actors, venv, etc). @@ -990,7 +1001,9 @@ def mock_async_grpo_infrastructure( mock_batch=mock_batch, mock_rollout_metrics=mock_rollout_metrics, ) - stub_collector = StubAsyncTrajectoryCollector() + stub_collector = StubAsyncTrajectoryCollector( + health_side_effect=collector_health_side_effect + ) # Patch venv creation stack.enter_context( @@ -1165,6 +1178,49 @@ def mock_sync_grpo_infrastructure(policy): return stack +def test_async_grpo_propagates_main_loop_collector_failure(mock_grpo_components): + """A fatal collector health result aborts the trainer and still cleans up.""" + master_config = mock_grpo_components["master_config"] + master_config.grpo.max_num_steps = 1 + master_config.grpo.val_period = 0 + master_config.grpo.val_at_start = False + master_config.grpo.val_at_end = False + master_config.grpo.use_dynamic_sampling = False + master_config.policy["generation"]["colocated"]["enabled"] = False + + mock_batch = next(iter(mock_grpo_components["train_dataloader"])) + mock_rollout_metrics = {"mean_gen_tokens_per_sample": 2.0} + + with ( + mock_async_grpo_infrastructure( + mock_batch, + mock_rollout_metrics, + collector_health_side_effect=[ + None, + RuntimeError("collector health failed"), + ], + ), + pytest.raises(RuntimeError, match="collector health failed"), + ): + async_grpo_train( + mock_grpo_components["policy"], + _mock_policy_generation(), + mock_grpo_components["train_dataloader"], + mock_grpo_components["val_dataloader"], + mock_grpo_components["tokenizer"], + mock_grpo_components["loss_fn"], + mock_grpo_components["task_to_env"], + mock_grpo_components["val_task_to_env"], + mock_grpo_components["logger"], + mock_grpo_components["checkpointer"], + _initial_grpo_save_state(), + master_config, + ) + + mock_grpo_components["checkpointer"].shutdown.assert_called_once() + mock_grpo_components["policy"].shutdown.assert_called_once() + + @pytest.mark.parametrize( ("generation_config", "expected"), [ diff --git a/tests/unit/algorithms/test_grpo_router_replay_async.py b/tests/unit/algorithms/test_grpo_router_replay_async.py index 52aa8369ac..d753c90ca3 100644 --- a/tests/unit/algorithms/test_grpo_router_replay_async.py +++ b/tests/unit/algorithms/test_grpo_router_replay_async.py @@ -19,6 +19,8 @@ import torch from nemo_rl.algorithms.grpo import ( + AsyncGRPOConfig, + GRPOConfig, MasterConfig, _build_async_grpo_train_data, _initial_grpo_save_state, @@ -60,6 +62,11 @@ def _make_async_master_config(data_plane=None) -> MasterConfig: }, }, "loss_fn": SimpleNamespace(use_importance_sampling_correction=True), + "grpo": GRPOConfig.model_construct( + async_grpo=AsyncGRPOConfig.model_construct( + max_generation_failures=0, + ) + ), "data_plane": data_plane, } ) diff --git a/tests/unit/reference_configs/grpo_math_1B.yaml b/tests/unit/reference_configs/grpo_math_1B.yaml index 7f026fca47..2c1017c692 100644 --- a/tests/unit/reference_configs/grpo_math_1B.yaml +++ b/tests/unit/reference_configs/grpo_math_1B.yaml @@ -58,6 +58,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