Buckets:
| """Async GRPO glue: our episodes as the rollout worker of TRL's AsyncGRPOTrainer. | |
| TRL's built-in rollout loop only continues a conversation through tool calls and re-renders the chat each | |
| turn. Our episode is one token stream in which the target's replies come back as user messages (rollout.py), | |
| so, like TRL's own HarnessRolloutWorker, we subclass its rollout loop and replace `_generate_one` (one | |
| conversation) with rollout.play_episode. TRL still does everything around it: the requests to the vLLM server | |
| (naming the current LoRA version, returning top-k/top-p candidate sets and sampler log-probs), turning the | |
| turns into training rows, group advantages, staleness, weight sync, and the queue to the trainer. | |
| The loop runs in a CUDA-free child process that TRL spawns. Each episode plays in a thread of our own pool, | |
| since the target and judge calls are blocking HTTP. Its attacker turns are handed to the loop's event loop, | |
| where TRL's `_generate_one_turn` posts them to the server. | |
| Every turn goes out under the policy version live at that moment, so a long episode can span weight syncs. | |
| Each turn's log-probs come from the policy that sampled it, which is what the trainer's importance ratio | |
| needs. Staleness is counted from when the episode's group started. | |
| Every finished episode is appended to rollouts.jsonl, and its reward parts, outcome rates, statuses and | |
| timings go to wandb under episode/, status/ and timing/. | |
| """ | |
| import asyncio | |
| import json | |
| import time | |
| import uuid | |
| from concurrent.futures import ThreadPoolExecutor | |
| from typing import Any | |
| from trl.chat_template_utils import add_response_schema | |
| from trl.experimental.async_grpo.async_rollout_worker import ( | |
| AsyncRolloutWorker, | |
| TurnRecord, | |
| _AsyncRolloutLoop, | |
| _chain_to_sequences, | |
| ) | |
| import rewards | |
| import rollout | |
| from episode import DROPPED_STATUSES, Episode, Job, to_record | |
| # Reward parts averaged over scored (not dropped) episodes, as episode/<name> in wandb. | |
| REWARD_PARTS = ( | |
| "outcome_score", | |
| "attack_quality_reward", | |
| "early_done_bonus", | |
| "false_done_penalty", | |
| "format_penalty", | |
| ) | |
| # How an episode can end (rollout.play_episode, rewards.judge_episode), as status/<name>_frac in wandb. | |
| STATUSES = ( | |
| "ok", | |
| "attacker_done", | |
| "no_message", | |
| "context_full", | |
| "verified_leak", | |
| "errored", | |
| "judge_errored", | |
| ) | |
| class AttackerRolloutLoop(_AsyncRolloutLoop): | |
| """TRL's rollout loop with one of our episodes as each conversation; its reward is the episode's reward.""" | |
| _provides_rollout_reward = True | |
| def __init__( | |
| self, | |
| *, | |
| jobs: list[Job], | |
| config: dict[str, Any], | |
| log_path: str, | |
| **loop_kwargs: Any, | |
| ) -> None: | |
| # TRL parses responses for tool calls and refuses tokenizers whose chat template it does not recognize. | |
| # We never parse responses, so an unrecognized template (e.g. a fine-tune's edited one) gets a placeholder. | |
| tokenizer = loop_kwargs["processing_class"] | |
| if ( | |
| getattr(tokenizer, "response_template", None) is None | |
| and getattr(tokenizer, "response_schema", None) is None | |
| ): | |
| try: | |
| add_response_schema(tokenizer) | |
| except ValueError: | |
| tokenizer.response_schema = {} | |
| # Our stream only ever grows by appending, so turns never drift. If one ever did, forking it into its own | |
| # row keeps its attacker tokens trained, where realigning would silently turn them into context. | |
| loop_kwargs["fork_threshold_tokens"] = 0 | |
| super().__init__(**loop_kwargs) | |
| self.config = config | |
| self.log_path = log_path | |
| self.jobs_by_seed = rollout.jobs_by_seed(jobs, config["episode"]["max_turns"]) | |
| self.chat = rollout.ChatTokens(self.tokenizer) | |
| self.reward_func_names.append( | |
| "episode_reward" | |
| ) # the column the rollout reward is logged under | |
| # One thread per episode that can be in flight, so no episode waits for a thread. | |
| self._episode_pool = ThreadPoolExecutor( | |
| max_workers=max(1, self.max_inflight_tasks), thread_name_prefix="episode" | |
| ) | |
| async def _run_loops(self, stop_event: Any) -> None: | |
| try: | |
| await super()._run_loops(stop_event) | |
| finally: | |
| self._episode_pool.shutdown(wait=False, cancel_futures=True) | |
| async def _post(self, path: str, payload: dict, timeout: float, max_retries: int = 3) -> dict: | |
| """TRL's POST to the vLLM server, with /inference/v1/generate's top log-probs put back in the shape TRL reads. | |
| vLLM 0.32 returns each top log-prob candidate as {"token_id": 123, "logprob": ...}, while TRL (pinned commit | |
| and main as of 2026-10-28) still reads {"token": "token_id:123", ...} and crashes with KeyError: 'token'. | |
| Spelling the id the old way lets TRL's own parser key the candidates by token id unchanged. | |
| """ | |
| output = await super()._post(path, payload, timeout, max_retries) | |
| if path == "/inference/v1/generate": | |
| for choice in output.get("choices", []): | |
| for token in (choice.get("logprobs") or {}).get("content") or []: | |
| for candidate in token.get("top_logprobs", []): | |
| if "token" not in candidate and "token_id" in candidate: | |
| candidate["token"] = f"token_id:{candidate['token_id']}" | |
| return output | |
| async def _generate_one( | |
| self, | |
| prompt: list[dict[str, str]], | |
| tool_dict: dict[str, Any], | |
| tools: list[Any], | |
| group_id: int = 0, | |
| ) -> tuple[list[dict[str, str]], list[int], Any, int, int, float | None]: | |
| """Play, judge and score one episode; return it the way TRL's loop expects one conversation.""" | |
| started = time.monotonic() | |
| started_version = self.model_version # the policy live when this episode's first turn goes out | |
| job = self.jobs_by_seed[prompt[-1]["content"]] | |
| turns: list[TurnRecord] = [] # one per attacker turn, filled in by the episode thread | |
| ep = await asyncio.get_running_loop().run_in_executor( | |
| self._episode_pool, self.play, job, turns | |
| ) | |
| # A dropped episode (a target or judge call failed) says nothing about the attacker: no reward, so it is | |
| # left out of its group's baseline, and no training rows. | |
| dropped = ep["status"] in DROPPED_STATUSES | |
| sequences, tally = _chain_to_sequences(turns, uuid.uuid4().hex, self._fork_threshold_tokens) | |
| if dropped: | |
| sequences = [] | |
| completion_ids = [token for turn in turns for token in turn.output_ids] | |
| self._push_rollout_metrics( | |
| turns=len(turns), | |
| sequences=len(sequences), | |
| completion_ids=completion_ids, | |
| tally=tally, | |
| loop_exhausted=False, | |
| duration_s=time.monotonic() - started, | |
| ) | |
| self.record_episode(ep, dropped, group_id, started_version) | |
| # The conversation as messages, for TRL's sample logging: the attacker is the assistant. | |
| completion = [ | |
| {"role": "assistant" if turn["role"] == "attacker" else "user", "content": turn["content"]} | |
| for turn in ep["transcript"] | |
| ] | |
| return completion, completion_ids, sequences, 0, 0, None if dropped else ep["reward"] | |
| def play(self, job: Job, turns: list[TurnRecord]) -> Episode: | |
| """Runs in an episode thread: play the episode against TRL's vLLM requests, then judge and score it.""" | |
| def generate_turn(prompt_ids: list[int]) -> tuple[list[int], list[float], str]: | |
| # TRL's request coroutine runs on the loop's event loop; this thread waits for its result. | |
| ids, logprobs, sampling_mask, sampler_logprobs = asyncio.run_coroutine_threadsafe( | |
| self._generate_one_turn(prompt_ids), self._loop | |
| ).result() | |
| turns.append( | |
| TurnRecord( | |
| list(prompt_ids), | |
| list(ids), | |
| list(logprobs), | |
| output_sampling_mask=sampling_mask, | |
| output_sampler_logprobs=sampler_logprobs, | |
| ) | |
| ) | |
| finish_reason = "stop" if ids and ids[-1] in self.eos_token_ids else "length" | |
| return list(ids), list(logprobs), finish_reason | |
| ep = rollout.new_episode(job, self.chat, self.config["episode"]["max_turns"]) | |
| rollout.play_episode(ep, self.chat, generate_turn, self.config) | |
| rewards.score_episode(ep, self.config["reward"]) | |
| return ep | |
| def record_episode(self, ep: Episode, dropped: bool, group_id: int, started_version: int) -> None: | |
| """Append the episode to rollouts.jsonl and add it to the wandb rates. Runs on the event loop's thread, | |
| so writes never interleave.""" | |
| record = to_record( | |
| ep, | |
| policy_version=self.model_version, | |
| started_policy_version=started_version, | |
| group=group_id, | |
| dropped=dropped, | |
| ) | |
| with open(self.log_path, "a", encoding="utf-8") as fout: | |
| fout.write(json.dumps(record, ensure_ascii=False) + "\n") | |
| # Each rate is (sum, count), reduced over the logging window by the trainer. | |
| self.add_rate("episode/dropped_frac", float(dropped)) | |
| # Every episode counts toward every status, so each fraction is over all episodes. | |
| for status in STATUSES: | |
| self.add_rate(f"status/{status}_frac", float(ep["status"] == status)) | |
| for name, seconds in ep["timing"].items(): | |
| self.add_rate(f"timing/{name}_s", seconds) | |
| # Weight updates that happened while this episode played. TRL drops a whole group once it is more than | |
| # max_staleness updates old (counted from the group's start, logged as sample/dropped_stale_total), so | |
| # this creeping toward max_staleness means long conversations are being thrown away. | |
| self.add_rate("episode/versions_spanned", self.model_version - started_version) | |
| if dropped: | |
| return | |
| self.add_rate("episode/reward", ep["reward"]) | |
| for name in REWARD_PARTS: | |
| self.add_rate(f"episode/{name}", ep[name]) | |
| self.add_rate("episode/asr", ep["success"]) | |
| self.add_rate("episode/partial_frac", float(ep["outcome"] == "partial")) | |
| self.add_rate("episode/verified_leak_frac", float(bool(ep["leaked"]))) | |
| self.add_rate("episode/turns", ep["turns"]) | |
| def add_rate(self, key: str, value: float) -> None: | |
| self._rates[key][0] += value | |
| self._rates[key][1] += 1 | |
| class AttackerRolloutWorker(AsyncRolloutWorker): | |
| """TRL's rollout worker, spawning AttackerRolloutLoop instead of its own loop. Pass it to | |
| AsyncGRPOTrainer(rollout_worker=...); its keyword arguments go to AttackerRolloutLoop in the child.""" | |
| _loop_cls = AttackerRolloutLoop | |
Xet Storage Details
- Size:
- 10.9 kB
- Xet hash:
- 4303571b9b7833acec9c69df8a6f1721f0c77573f250722fe03e0c6df2de74d8
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.