diff --git a/docs/source/async_distillation_trainer.md b/docs/source/async_distillation_trainer.md index 4e32137e2e9..a8fa0a13ea0 100644 --- a/docs/source/async_distillation_trainer.md +++ b/docs/source/async_distillation_trainer.md @@ -74,7 +74,8 @@ logprob for at each `beta` regime. After every `weight_sync_steps` training steps, the updated student weights are transferred to its vLLM server via NCCL. As with [`experimental.async_grpo.AsyncGRPOTrainer`], generation runs ahead of training, so samples may reflect a slightly stale -policy; `max_staleness` controls how many weight updates a sample can lag behind before being discarded. +policy; `max_staleness` controls how many weight updates a sample can lag behind before in-flight work is cancelled or +a queued sample is discarded. ## Quick start @@ -229,6 +230,7 @@ A **rollout** is one prompt taken all the way through: generated by the student, | `rollout/inflight` | rollouts in flight, i.e. generating or being scored | | `rollout/vllm_retry_total` | retried vLLM requests, to either the student's server or a teacher's. A degraded server otherwise looks like unexplained slowness. It sits here rather than in `completions/` because it counts requests to a server, not generated text | | `rollout/backpressure_s` | how long generation was blocked because the rollout queue was full. See [the rollout queue](#the-rollout-queue) | +| `rollout/stale_samples_total` | in-flight samples cancelled because the policy moved more than `max_staleness` versions past the one they started at | ### Samples arriving from the queue diff --git a/docs/source/async_grpo_trainer.md b/docs/source/async_grpo_trainer.md index be68a531998..b924cb49b3b 100644 --- a/docs/source/async_grpo_trainer.md +++ b/docs/source/async_grpo_trainer.md @@ -36,7 +36,7 @@ The rollout worker runs in a separate process spawned from the trainer, so rewar After every `weight_sync_steps` training steps, the updated weights are transferred to the vLLM server via NCCL so that subsequent generations reflect the latest policy. -Because generation and training run concurrently, the training samples may have been generated by a slightly older version of the model. The `max_staleness` parameter controls how many weight updates a sample can lag behind before being discarded. +Because generation and training run concurrently, the training samples may have been generated by a slightly older version of the model. The `max_staleness` parameter controls how many weight updates a sample can lag behind before being discarded. The worker applies the same limit to its in-flight rollouts: when the policy advances, it cancels the generations of any group that started too many versions ago, so vLLM does not finish work the trainer would drop. The number of concurrent requests sent to the vLLM server is controlled by `max_inflight_tasks`. By default it is set automatically to `max_staleness × per_device_train_batch_size × gradient_accumulation_steps × num_processes` — the maximum number of samples the trainer can consume before they become stale. Generating more than this is wasteful since the excess samples will be discarded. @@ -259,6 +259,9 @@ A **rollout** is **one full** conversation: a prompt generated to completion, in | `rollout/score_s`, `rollout/score_wait_s`, `rollout/score_block_s` | scoring: time to score a group, group wait time to be scored, and how long generation was blocked because the scoring queue was full | | `rollout/vllm_retry_total` | retried vLLM requests. A degraded server otherwise looks like unexplained slowness. It sits here rather than in `completions/` because it counts requests to the server, not generated text: a retried request produced no completion at all | | `rollout/backpressure_s` | how long generation was blocked because the rollout queue was full. See [the rollout queue](#the-rollout-queue) | +| `rollout/failed_total` | rollouts that raised (a request that failed every retry, a completion that could not be parsed). The rollout is dropped and its group scored with the rest; only a group where every rollout failed takes the worker down | +| `rollout/dropped_groups_total` | groups dropped because a single rollout survived: a group-relative advantage needs at least two | +| `rollout/stale_groups_total` | in-flight groups cancelled because the policy moved more than `max_staleness` versions past the one they started at. The trainer would drop their samples anyway (see `sample/dropped_stale_total`), so finishing them only burns vLLM compute | ### Tools diff --git a/tests/experimental/test_async_distillation_trainer.py b/tests/experimental/test_async_distillation_trainer.py index 9a4ee90af24..dce8192661c 100644 --- a/tests/experimental/test_async_distillation_trainer.py +++ b/tests/experimental/test_async_distillation_trainer.py @@ -25,6 +25,7 @@ import pytest import torch +from accelerate import PartialState from datasets import Dataset, load_dataset from transformers import AutoTokenizer from transformers.testing_utils import torch_device @@ -356,6 +357,24 @@ def _bare_loop(tokenizer, teacher_server_urls): TWO_TEACHERS = {"math": "http://math:8002", "code": "http://code:8003"} +def _rollout_loop(dataset, **kwargs): + PartialState() + ctx = mp.get_context("spawn") + loop_kwargs = dict( + model_name="test", + dataset=dataset, + processing_class=MagicMock(), + rollout_buffer=ctx.Queue(), + metrics_queue=ctx.Queue(), + model_version_value=ctx.Value("i", 0), + heartbeat_value=ctx.Value("d", 0.0), + failed_event=ctx.Event(), + exception_info_queue=ctx.Queue(), + ) + loop_kwargs.update(kwargs) + return _AsyncRolloutLoop(**loop_kwargs) + + class TestWorkerMetrics: """The worker's payload has the same shape as the trainer's sink, so draining it in `log()` is an append.""" @@ -922,6 +941,39 @@ def test_stops_once_the_prompt_target_is_reached(self, trained, before_resume, s assert control.should_training_stop is should_stop +class TestGenerateLoop(TrlTestCase): + def test_stale_in_flight_samples_are_cancelled_when_the_policy_advances(self): + loop = _rollout_loop( + Dataset.from_dict({"prompt": [f"q{i}" for i in range(8)]}), max_inflight_tasks=2, max_staleness=0 + ) + cancelled = [] + + async def run(): + stop = asyncio.Event() + + async def generate_and_score_one(prompt_id, row): + version = loop.model_version + if version == 0: + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + cancelled.append(prompt_id) + raise + stop.set() + return types.SimpleNamespace(completion_mask=[1], model_version=version, enqueued_at=None) + + loop._generate_and_score_one = generate_and_score_one + task = asyncio.create_task(loop._generate_loop(stop)) + await asyncio.sleep(0.1) + loop._model_version_value.value = 1 + await asyncio.wait_for(task, 5) + + asyncio.run(run()) + assert len(cancelled) == 2 + assert loop.rollout_buffer.get(timeout=5).model_version == 1 + assert loop._metrics_queue.get(timeout=5)["rollout/stale_samples_total"] == 2 + + class TestRolloutStateCheckpoint(TrlTestCase): """Prompt-index checkpoint/resume logic — no GPU or vLLM required.""" @@ -974,17 +1026,8 @@ def record(*_args, **_kwargs): assert written_before_super == [True] def test_rollout_loop_skips_to_start_index(self): - ctx = mp.get_context("spawn") - loop = _AsyncRolloutLoop( - model_name="test", - dataset=Dataset.from_dict({"prompt": [f"row_{i}" for i in range(10)]}), - processing_class=MagicMock(), - rollout_buffer=ctx.Queue(), - model_version_value=ctx.Value("i", 0), - heartbeat_value=ctx.Value("d", 0.0), - failed_event=ctx.Event(), - exception_info_queue=ctx.Queue(), - metrics_queue=ctx.Queue(), + loop = _rollout_loop( + Dataset.from_dict({"prompt": [f"row_{i}" for i in range(10)]}), dataset_start_index=3, ) _prompt_id, row = next(loop._repeat_iterator()) diff --git a/tests/experimental/test_async_grpo_trainer.py b/tests/experimental/test_async_grpo_trainer.py index 6f2fb5e98a1..636756995f6 100644 --- a/tests/experimental/test_async_grpo_trainer.py +++ b/tests/experimental/test_async_grpo_trainer.py @@ -645,7 +645,7 @@ def test_rollout_loop_skips_to_start_index(self): dataset = Dataset.from_dict({"prompt": [f"row_{i}" for i in range(10)]}) loop = self._make_rollout_loop(dataset, dataset_start_index=3) it = loop._repeat_iterator() - _group_id, row = next(it) + _group_id, _index, row = next(it) assert row["prompt"] == "row_3" def test_inner_training_loop_sets_dataset_start_index_from_file(self): @@ -1373,6 +1373,143 @@ def test_epoch_stop_is_fork_independent(self): assert forked.state.global_step > no_fork.state.global_step +def _strict_reward(completions, answer, **kwargs): + return [1.0 for _ in zip(completions, answer, strict=True)] + + +_ROLLOUT = ( + [{"role": "assistant", "content": "c"}], + [1, 2], + [TrainingSequence([1, 2], [0, 1], [0.0, -0.1], "r")], + 0, + 0, + None, +) + + +class TestGenerateLoop(TrlTestCase): + def _loop(self, num_generations, max_inflight_tasks, max_staleness=4): + PartialState() + with patch("trl.experimental.async_grpo.async_rollout_worker.add_response_schema", side_effect=lambda x: x): + return _AsyncRolloutLoop( + model_name="test", + dataset=Dataset.from_dict({"prompt": [f"q{i}" for i in range(8)]}), + reward_funcs=[dummy_reward_func], + processing_class=MagicMock(), + rollout_buffer=queue.Queue(), + metrics_queue=queue.Queue(), + model_version_value=mp.Value("i", 0), + heartbeat_value=mp.Value("d", 0.0), + failed_event=mp.Event(), + exception_info_queue=queue.Queue(), + num_generations=num_generations, + max_inflight_tasks=max_inflight_tasks, + max_staleness=max_staleness, + ) + + async def _groups(self, loop, generate_one, n, bump_version_to=None): + """Run the generate loop until `n` groups reach the score queue.""" + loop._generate_one = generate_one + stop = asyncio.Event() + task = asyncio.create_task(loop._generate_loop(stop)) + if bump_version_to is not None: + await asyncio.sleep(0.1) + loop._model_version_value.value = bump_version_to + groups = [await asyncio.wait_for(loop._groups_to_score.get(), 5) for _ in range(n)] + stop.set() + await task + return groups + + def test_failed_rollout_is_dropped_and_the_group_scored_with_the_rest(self): + loop = self._loop(num_generations=4, max_inflight_tasks=4) + calls = itertools.count() + + async def generate_one(prompt, tool_dict, tools, group_id): + if next(calls) == 1: + raise RuntimeError("boom") + return _ROLLOUT + + (group,) = asyncio.run(self._groups(loop, generate_one, 1)) + assert group.group_id == 0 + assert len(group.completions) == 3 + + def test_group_with_a_single_surviving_rollout_is_dropped(self): + loop = self._loop(num_generations=3, max_inflight_tasks=3) + calls = itertools.count() + + async def generate_one(prompt, tool_dict, tools, group_id): + if next(calls) < 2: + raise RuntimeError("boom") + return _ROLLOUT + + (group,) = asyncio.run(self._groups(loop, generate_one, 1)) + assert group.group_id == 1 + + def test_group_where_every_rollout_fails_raises(self): + loop = self._loop(num_generations=2, max_inflight_tasks=2) + + async def generate_one(prompt, tool_dict, tools, group_id): + raise RuntimeError("boom") + + loop._generate_one = generate_one + with pytest.raises(RuntimeError, match="boom"): + asyncio.run(asyncio.wait_for(loop._generate_loop(asyncio.Event()), 5)) + + def test_stale_in_flight_groups_are_cancelled_when_the_policy_advances(self): + loop = self._loop(num_generations=2, max_inflight_tasks=4, max_staleness=1) + cancelled = [] + + async def generate_one(prompt, tool_dict, tools, group_id): + if loop.model_version == 0: + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + await asyncio.sleep(0) + cancelled.append(group_id) + raise + assert len(cancelled) == 4 + return _ROLLOUT + + groups = asyncio.run(self._groups(loop, generate_one, 2, bump_version_to=2)) + assert len(cancelled) == 4 + assert [g.group_id for g in groups] == [2, 3] + assert [g.model_version for g in groups] == [2, 2] + + def test_partially_dispatched_stale_group_is_regenerated_as_a_smaller_group(self): + loop = self._loop(num_generations=4, max_inflight_tasks=2, max_staleness=0) + + async def generate_one(prompt, tool_dict, tools, group_id): + if loop.model_version == 0: + await asyncio.Event().wait() + return _ROLLOUT + + (group,) = asyncio.run(self._groups(loop, generate_one, 1, bump_version_to=1)) + assert group.group_id == 0 + assert len(group.completions) == 2 + assert group.model_version == 1 + + def test_stale_group_with_one_undispatched_rollout_is_skipped(self): + loop = self._loop(num_generations=4, max_inflight_tasks=3, max_staleness=0) + + async def generate_one(prompt, tool_dict, tools, group_id): + if loop.model_version == 0: + await asyncio.Event().wait() + if group_id == 0: + raise RuntimeError("the stale tail must not be dispatched") + return _ROLLOUT + + (group,) = asyncio.run(self._groups(loop, generate_one, 1, bump_version_to=1)) + assert group.group_id == 1 + assert group.model_version == 1 + + def test_reward_kwargs_are_trimmed_to_the_surviving_rollouts(self): + PartialState() + group = _group([[_ROLLOUT[2][0]]] * 3, [[1, 2]] * 3) + group.reward_kwargs = {"answer": ["a"] * 4} + samples = asyncio.run(_bare_loop([_strict_reward])._score_group(group)) + assert len(samples) == 3 + + @require_peft class TestValidateLoraForVLLMSync(TrlTestCase): model_id = "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5" diff --git a/trl/experimental/async_distillation/async_distillation_config.py b/trl/experimental/async_distillation/async_distillation_config.py index 9a24e9a2917..00d5aea51d5 100644 --- a/trl/experimental/async_distillation/async_distillation_config.py +++ b/trl/experimental/async_distillation/async_distillation_config.py @@ -149,8 +149,8 @@ class AsyncDistillationConfig(_BaseConfig): `-1` (auto), which sets it to `max_staleness * per_device_train_batch_size * gradient_accumulation_steps * num_processes`. max_staleness (`int`, *optional*, defaults to `4`): - Maximum number of weight update steps a rollout sample can lag behind the current model version before - being discarded. + Maximum number of weight update steps a rollout sample can lag behind the current model version before an + in-flight sample is cancelled or a queued sample is discarded. queue_maxsize (`int`, *optional*, defaults to `1024`): Maximum number of rollout samples to buffer in the rollout queue. weight_sync_steps (`int`, *optional*, defaults to `1`): @@ -358,7 +358,7 @@ class AsyncDistillationConfig(_BaseConfig): default=4, metadata={ "help": "Maximum number of weight update steps a rollout sample can lag behind the current model " - "version before being discarded." + "version before an in-flight sample is cancelled or a queued sample is discarded." }, ) queue_maxsize: int = field( diff --git a/trl/experimental/async_distillation/async_distillation_trainer.py b/trl/experimental/async_distillation/async_distillation_trainer.py index 127b62799f1..0bcd5c48230 100644 --- a/trl/experimental/async_distillation/async_distillation_trainer.py +++ b/trl/experimental/async_distillation/async_distillation_trainer.py @@ -1133,6 +1133,7 @@ def __init__( teacher_top_k=self.args.teacher_top_k, teacher_temperature=self.args.teacher_temperature, max_tokens=self.args.max_completion_length, + max_staleness=self.args.max_staleness, temperature=self.args.temperature, top_p=self.args.top_p, top_k=self.args.top_k, diff --git a/trl/experimental/async_distillation/async_rollout_worker.py b/trl/experimental/async_distillation/async_rollout_worker.py index 756c22a94ce..8354a6ab35e 100644 --- a/trl/experimental/async_distillation/async_rollout_worker.py +++ b/trl/experimental/async_distillation/async_rollout_worker.py @@ -208,6 +208,7 @@ def __init__( teacher_top_k: int = 8, teacher_temperature: float = 1.0, max_tokens: int = 32, + max_staleness: int = 4, temperature: float = 1.0, top_p: float = 1.0, top_k: int = 0, @@ -243,6 +244,7 @@ def __init__( self.max_inflight_tasks = max_inflight_tasks self.queue_maxsize = queue_maxsize self.max_tokens = max_tokens + self.max_staleness = max_staleness self.temperature = temperature self.top_p = top_p self.top_k = top_k @@ -321,19 +323,41 @@ async def _resolve_teacher_model_names(self) -> None: logger.info(f"teacher {teacher_id!r} at {url} serves {self.teacher_model_names[teacher_id]}") async def _generate_loop(self, stop_event: asyncio.Event) -> None: - inflight_tasks: dict[asyncio.Task, int] = {} + # Keep the dispatch version beside the slot: a sample does not expose its model version until its task returns. + inflight_tasks: dict[asyncio.Task, tuple[int, int]] = {} free_slots = set(range(self.max_inflight_tasks)) work_iter = self._repeat_iterator() + last_version = self.model_version self._generation_start_time = time.monotonic() try: while True: self._heartbeat_value.value = time.time() + + version = self.model_version + if version != last_version: + last_version = version + stale_tasks = [ + task + for task, (_slot, task_version) in inflight_tasks.items() + if version - task_version > self.max_staleness + ] + for task in stale_tasks: + task.cancel() + if stale_tasks: + await asyncio.gather(*stale_tasks, return_exceptions=True) + for task in stale_tasks: + slot, _task_version = inflight_tasks.pop(task) + free_slots.add(slot) + if stale_tasks: + self._counters["rollout/stale_samples_total"] += len(stale_tasks) + logger.info(f"cancelled {len(stale_tasks)} stale rollout(s) at version {version}") + while free_slots and not stop_event.is_set(): prompt_id, row = next(work_iter) slot = free_slots.pop() task = asyncio.create_task(self._generate_and_score_one(prompt_id, row)) - inflight_tasks[task] = slot + inflight_tasks[task] = (slot, self.model_version) if not inflight_tasks: if stop_event.is_set(): @@ -346,7 +370,7 @@ async def _generate_loop(self, stop_event: asyncio.Event) -> None: continue for task in done: - slot = inflight_tasks.pop(task) + slot, _task_version = inflight_tasks.pop(task) free_slots.add(slot) if task.exception() is not None: raise task.exception() diff --git a/trl/experimental/async_grpo/async_grpo_trainer.py b/trl/experimental/async_grpo/async_grpo_trainer.py index 0066dd76c66..68537118e50 100644 --- a/trl/experimental/async_grpo/async_grpo_trainer.py +++ b/trl/experimental/async_grpo/async_grpo_trainer.py @@ -1321,6 +1321,7 @@ def __init__( queue_maxsize=self.args.queue_maxsize, vllm_server_url=self.args.vllm_server_base_url, max_tokens=self.args.max_completion_length, + max_staleness=self.args.max_staleness, temperature=self.args.temperature, top_p=self.args.top_p, top_k=self.args.top_k, diff --git a/trl/experimental/async_grpo/async_rollout_worker.py b/trl/experimental/async_grpo/async_rollout_worker.py index 43c8406b009..a9f5edf0c48 100644 --- a/trl/experimental/async_grpo/async_rollout_worker.py +++ b/trl/experimental/async_grpo/async_rollout_worker.py @@ -315,6 +315,7 @@ def __init__( vllm_server_url: str = "http://localhost:8000", lora_name: str | None = None, max_tokens: int = 32, + max_staleness: int = 4, temperature: float = 1.0, top_p: float = 1.0, top_k: int = 0, @@ -369,6 +370,7 @@ def __init__( self.max_inflight_tasks = max_inflight_tasks self.queue_maxsize = queue_maxsize self.max_tokens = max_tokens + self.max_staleness = max_staleness self.temperature = temperature self.top_p = top_p self.top_k = top_k @@ -515,14 +517,45 @@ async def _generate_loop(self, stop_event: asyncio.Event) -> None: inflight_tasks: dict[asyncio.Task, tuple[int, int, Any, object, Messages]] = {} free_slots = set(range(self.max_inflight_tasks)) work_iter = self._repeat_iterator() + last_version = self.model_version self._generation_start_time = time.monotonic() try: while True: # Wall-clock for cross-process comparison; parent uses time.time() in check_health. self._heartbeat_value.value = time.time() + + version = self.model_version + if version != last_version: + last_version = version + stale = { + group_id + for group_id, group in pending_groups.items() + if version - group.model_version > self.max_staleness + } + stale_tasks = [task for task, values in inflight_tasks.items() if values[0] in stale] + for task in stale_tasks: + task.cancel() + if stale_tasks: + await asyncio.gather(*stale_tasks, return_exceptions=True) + for task in stale_tasks: + _group_id, slot, name, environment, _prompt = inflight_tasks.pop(task) + free_slots.add(slot) + if environment is not None: + self._environment_pool[name].append(environment) + for group_id in stale: + del pending_groups[group_id] + del pending_completed[group_id] + if stale: + self._counters["rollout/stale_groups_total"] += len(stale) + logger.info(f"[generate] cancelled {len(stale)} stale group(s) at version {version}") + while free_slots and not stop_event.is_set(): - group_id, row = next(work_iter) + group_id, index, row = next(work_iter) + # Missing state at a nonzero index means stale cancellation discarded the group's earlier work. + # Do not dispatch a lone remainder: it cannot produce a group-relative advantage. + if group_id not in pending_groups and self.num_generations - index < 2: + continue slot = free_slots.pop() # The environment is selected per example via its `environment` field (multi-env); only its tools # are exposed in the example's prompt. When there are no environments, every example shares the @@ -598,7 +631,7 @@ async def _generate_loop(self, stop_event: asyncio.Event) -> None: env_rewards=[], rollout_rewards=[], ) - pending_completed[group_id] = 0 + pending_completed[group_id] = index task = asyncio.create_task( self._generate_one(prompt, tool_dict=tool_dict, tools=tools, group_id=group_id) @@ -619,45 +652,58 @@ async def _generate_loop(self, stop_event: asyncio.Event) -> None: for task in done: group_id, slot, name, environment, prompt = inflight_tasks.pop(task) free_slots.add(slot) - if task.exception() is not None: - raise task.exception() - - ( - completion, - completion_ids, - sequences, - tool_call_count, - tool_failure_count, - rollout_reward, - ) = task.result() group = pending_groups[group_id] - group.prompts.append(prompt) - group.completions.append(completion) - group.completions_ids.append(completion_ids) - group.completions_sequences.append(sequences) - group.tool_call_counts.append(tool_call_count) - group.tool_failure_counts.append(tool_failure_count) - group.rollout_rewards.append(rollout_reward) - # The environment owns the reward: score it now, while this rollout's environment still holds its - # final state and before returning it to the pool. `get_reward` may be async awaiting yields to - # inflight requests instead of halting them. The env is returned to the pool only after scoring, so - # a concurrent rollout can't draw and reset it during the await. Record `(env class, reward)` so - # `_score_group` can place it in the matching env's reward column; rollouts whose env owns no reward - # record `None` (turned into NaN and ignored) to stay aligned with the group's other per-rollout lists. - if self._env_reward_types: - env_type = type(environment) - if env_type in self._env_reward_types: - get_reward = environment.get_reward - reward = await get_reward() if inspect.iscoroutinefunction(get_reward) else get_reward() - group.env_rewards.append((env_type, reward)) - else: - group.env_rewards.append(None) + error = task.exception() + if error is not None: + logger.warning(f"[generate] rollout failed for group {group_id}, dropping it", exc_info=error) + self._counters["rollout/failed_total"] += 1 + else: + ( + completion, + completion_ids, + sequences, + tool_call_count, + tool_failure_count, + rollout_reward, + ) = task.result() + group.prompts.append(prompt) + group.completions.append(completion) + group.completions_ids.append(completion_ids) + group.completions_sequences.append(sequences) + group.tool_call_counts.append(tool_call_count) + group.tool_failure_counts.append(tool_failure_count) + group.rollout_rewards.append(rollout_reward) + # The environment owns the reward: score it now, while this rollout's environment still holds + # its final state and before returning it to the pool. `get_reward` may be async awaiting + # yields to inflight requests instead of halting them. The env is returned to the pool only + # after scoring, so a concurrent rollout can't draw and reset it during the await. Record + # `(env class, reward)` so `_score_group` can place it in the matching env's reward column; + # rollouts whose env owns no reward record `None` (turned into NaN and ignored) to stay aligned + # with the group's other per-rollout lists. + if self._env_reward_types: + env_type = type(environment) + if env_type in self._env_reward_types: + get_reward = environment.get_reward + reward = ( + await get_reward() if inspect.iscoroutinefunction(get_reward) else get_reward() + ) + group.env_rewards.append((env_type, reward)) + else: + group.env_rewards.append(None) + self._total_completion_tokens += len(completion_ids) if environment is not None: self._environment_pool[name].append(environment) - self._total_completion_tokens += len(completion_ids) pending_completed[group_id] += 1 if pending_completed[group_id] == self.num_generations: + del pending_groups[group_id] + del pending_completed[group_id] + if not group.completions: + raise error + if len(group.completions) < 2: + logger.warning(f"[generate] dropping group {group_id}: a single rollout succeeded") + self._counters["rollout/dropped_groups_total"] += 1 + continue group.queued_at = time.monotonic() t_blocked = None while True: @@ -673,8 +719,6 @@ async def _generate_loop(self, stop_event: asyncio.Event) -> None: if t_blocked is not None: # Generation held back by scoring self._push_metrics({"rollout/score_block_s": time.monotonic() - t_blocked}) - del pending_groups[group_id] - del pending_completed[group_id] finally: for task in inflight_tasks: task.cancel() @@ -839,7 +883,7 @@ def _push_rollout_metrics( } ) - def _repeat_iterator(self) -> Iterator[tuple[int, dict[str, Any]]]: + def _repeat_iterator(self) -> Iterator[tuple[int, int, dict[str, Any]]]: group_id = 0 while True: try: @@ -847,8 +891,8 @@ def _repeat_iterator(self) -> Iterator[tuple[int, dict[str, Any]]]: except StopIteration: self._dataset_iter = iter(self.dataset) row = next(self._dataset_iter) - for _ in range(self.num_generations): - yield group_id, row + for index in range(self.num_generations): + yield group_id, index, row group_id += 1 async def _generate_one( @@ -982,11 +1026,12 @@ async def _generate_one_turn(self, prompt_ids: list[int]) -> tuple[list[int], li return choice["token_ids"], choice["logprobs"]["token_logprobs"] async def _score_group(self, group: RolloutGroup) -> list[RolloutSample]: + n = len(group.completions) kwargs = dict( completions=group.completions, prompts=group.prompts, completion_ids=group.completions_ids, - **group.reward_kwargs, + **{key: values[:n] for key, values in group.reward_kwargs.items()}, ) all_rewards = await asyncio.gather( *[ diff --git a/trl/experimental/async_grpo/openenv_harness.py b/trl/experimental/async_grpo/openenv_harness.py index 16cd00164c7..f79258dbe24 100644 --- a/trl/experimental/async_grpo/openenv_harness.py +++ b/trl/experimental/async_grpo/openenv_harness.py @@ -17,11 +17,12 @@ import asyncio import functools import json +import threading import time import uuid from collections.abc import Callable -from concurrent.futures import ThreadPoolExecutor -from dataclasses import dataclass +from concurrent.futures import Future, ThreadPoolExecutor +from dataclasses import dataclass, field from typing import Any, Protocol, TypedDict, cast from accelerate.logging import get_logger @@ -83,6 +84,13 @@ class HarnessTurn: tool_calls: list[dict] # the tool calls the assistant emitted (empty for a pure-text turn) +@dataclass +class _CancellationToken: + event: threading.Event = field(default_factory=threading.Event) + session: object | None = None + sampling_future: Future | None = None + + def _tools_to_schema(tools: list) -> list[dict] | None: """Convert OpenEnv `Tool`s (MCP spec) into the OpenAI function schema `apply_chat_template` expects.""" if not tools: @@ -147,14 +155,30 @@ def __init__( self._session_pool = ThreadPoolExecutor( max_workers=max(1, self.max_inflight_tasks), thread_name_prefix="harness-session" ) - # In-flight sessions, so `_run_loops` can close them on stop (see there). set ops are atomic under the GIL. + # In-flight sessions, so `_run_loops` can close them on stop (see there). Set ops are atomic under the GIL. self._live_sessions: set = set() async def _generate_one(self, prompt, tool_dict, tools, group_id=0): # TODO(@openenv): provide an async version for performance # OpenEnv's harness layer is synchronous, so run the whole session on the pool. loop = asyncio.get_running_loop() - result, metrics = await loop.run_in_executor(self._session_pool, self._run_session, prompt, group_id) + token = _CancellationToken() + future = loop.run_in_executor(self._session_pool, self._run_session, prompt, group_id, token) + try: + result, metrics = await asyncio.shield(future) + except asyncio.CancelledError: + # Executor cancellation does not stop its thread, and white-box sampling runs in a separate event-loop + # future. Signal both paths, then keep the worker slot occupied until its thread has unwound. + token.event.set() + if token.sampling_future is not None: + token.sampling_future.cancel() + if token.session is not None: + try: + await asyncio.to_thread(token.session.close) + except Exception: + logger.warning("closing cancelled harness session failed", exc_info=True) + await asyncio.gather(future, return_exceptions=True) + raise # Pushed here and not in `_run_session`: the accumulators are plain dicts, and the pool runs many sessions at # once, so a push off the event loop would race the score loop's. if metrics is not None: @@ -179,7 +203,7 @@ def _close(session): finally: self._session_pool.shutdown(wait=True) - def _run_session(self, prompt, group_id=0): + def _run_session(self, prompt, group_id=0, token: _CancellationToken | None = None): """Drive one OpenEnv session to completion, on a pool thread. Returns `(the _generate_one tuple, rollout metrics or None)`. The metrics are handed back rather than pushed @@ -202,16 +226,20 @@ def _run_session(self, prompt, group_id=0): except Exception: logger.warning("harness session create failed; scoring rollout as unscorable", exc_info=True) return self._EMPTY_ROLLOUT, None - self._live_sessions.add(session) # tracked so a stop can close it (unblocks wait_for_completion below) + token = token or _CancellationToken() + token.session = session + self._live_sessions.add(session) timed_out = False trace: list[TraceEntry] = [] tool_calls_by_name: dict[str, int] = {} try: + if token.event.is_set(): + return self._EMPTY_ROLLOUT, None if self._adapter is not None: # white-box: the adapter runs the tool loop, calling `_sample_turn` each turn. turns: list[TurnRecord] = [] result = self._adapter.run_white_box( - functools.partial(self._sample_turn, turns), session, self._limits + functools.partial(self._sample_turn, turns, token), session, self._limits ) completion = result.messages tool_call_count = int(result.metrics.get("tool_calls", len(result.tool_trace))) @@ -267,7 +295,9 @@ def _run_session(self, prompt, group_id=0): except Exception: logger.warning("harness session close failed", exc_info=True) - def _sample_turn(self, turns: list[TurnRecord], messages, tools, sampling) -> ModelStepResult: + def _sample_turn( + self, turns: list[TurnRecord], token: _CancellationToken, messages, tools, sampling + ) -> ModelStepResult: """OpenEnv `ModelStep`: sample one assistant turn against vLLM and record a `TurnRecord` into `turns`.""" prompt_ids = self.tokenizer.apply_chat_template( messages, @@ -279,7 +309,14 @@ def _sample_turn(self, turns: list[TurnRecord], messages, tools, sampling) -> Mo **self.chat_template_kwargs, ) # ModelStep is sync on a pool thread; bridge the async vLLM POST onto the loop's event loop. - turn_ids, logprobs = asyncio.run_coroutine_threadsafe(self._generate_one_turn(prompt_ids), self._loop).result() + future = asyncio.run_coroutine_threadsafe(self._generate_one_turn(prompt_ids), self._loop) + token.sampling_future = future + if token.event.is_set(): + future.cancel() + try: + turn_ids, logprobs = future.result() + finally: + token.sampling_future = None turns.append(TurnRecord(prompt_ids, turn_ids, logprobs)) message = parse_response(self.tokenizer, turn_ids, prefix=prompt_ids) return ModelStepResult(