diff --git a/README.md b/README.md index b479197149..2245657857 100644 --- a/README.md +++ b/README.md @@ -217,7 +217,7 @@ Check out the [docs](docs) directory for in-depth guides on how to use prime-rl. - [**Configuration**](docs/configuration.md) - TOML composition, CLI overrides, env vars, validation - [**Training**](docs/training.md) - RL, SFT, evals, checkpointing, observability, rules of thumb - [**Scaling**](docs/scaling.md) - Single-GPU through multi-node, FSDP/EP/CP, SLURM, benchmarking -- [**Algorithms**](docs/algorithms.md) - Async/off-policy training, the AIPO loss, advantage and filter plugins, trajectory merging +- [**Algorithms**](docs/algorithms.md) - Async/off-policy training, advantage plugins, rollout checks, trajectory merging - [**Advanced**](docs/advanced.md) - Custom modeling, multimodal training, LoRA - [**Development**](docs/development.md) - Test suite, pre-commit hooks, adding a new model diff --git a/configs/ci/nightly-fft/wiki-search.toml b/configs/ci/nightly-fft/wiki-search.toml index c3977c0adb..008cb2dd71 100644 --- a/configs/ci/nightly-fft/wiki-search.toml +++ b/configs/ci/nightly-fft/wiki-search.toml @@ -17,11 +17,7 @@ name = "Qwen/Qwen3-4B-Instruct-2507" [orchestrator] batch_size = 512 group_size = 16 -oversampling_factor = 2.0 - -[[orchestrator.pre_batch_filters]] -type = "zero_advantage" -enforce = true +max_inflight_episodes = 1024 [[orchestrator.train.source]] name = "wiki-search" diff --git a/configs/debug/algo/echo.toml b/configs/debug/algo/echo.toml index e9bdddccb3..722f961e59 100644 --- a/configs/debug/algo/echo.toml +++ b/configs/debug/algo/echo.toml @@ -44,12 +44,6 @@ type = "subprocess" [orchestrator.train.sampling] max_completion_tokens = 512 -# ECHO learns from observation tokens even when the GRPO advantage collapses -# to zero — keep zero-advantage rollouts in the batch. -[[orchestrator.post_batch_filters]] -type = "zero_advantage" -enforce = false - # Fine-tune inherits the PrimeIntellect Qwen3 template byte-for-byte. [orchestrator.renderer] name = "prime-qwen3" diff --git a/docs/algorithms.md b/docs/algorithms.md index e3012fa50f..02805bb79c 100644 --- a/docs/algorithms.md +++ b/docs/algorithms.md @@ -1,6 +1,6 @@ # Algorithms -This page covers the math and the configurable algorithmic components: the algorithm abstraction and its algorithms, how off-policy training works, the loss components and advantage functions, how to plug in your own, the filters applied between rollout and training, and how multi-turn rollouts get merged into training samples. +This page covers the math and the configurable algorithmic components: the algorithm abstraction and its algorithms, how off-policy training works, the loss components and advantage functions, how to plug in your own, the checks applied between rollout and training, and how multi-turn rollouts get merged into training samples. ## Table of Contents @@ -164,12 +164,12 @@ At runtime, each env's resolved config builds two objects: a `Sampler` (`prime_r | `hierarchical_grpo` | `HierarchicalGRPOAlgorithm` | `score_group`: GRPO baseline per episode for solvers, per group for the proposer | | `opd` | `OPDAlgorithm` | `score_rollout`: own-context prefill under the teacher | | `opsd` | `OPSDAlgorithm` | `score_rollout`: demo-conditioned prefill under the live policy | -| `sft` | `SFTDistillAlgorithm` | `score_group`: group-norm credit (feeds filters) | +| `sft` | `SFTDistillAlgorithm` | no credit — ce trains on every sampled token | Each class owns its hooks outright — reading one top to bottom reads the algorithm, and everything on the class is an override point. The two hooks are one scope-and-timing ladder — the wider scope is unlocked by a later barrier, so the two axes coincide. Each is handed the `Rollout` directly — the env's typed trace (`reward`, `nodes`, `num_turns`, ...) with `samples` attached, plus `assign_advantages` to write credit: -- `async score_rollout(rollout)` — one rollout, **on arrival** (as it's tokenized, before its group is complete): rollout-local credit (`rollout.assign_advantages(...)`, scalar broadcast or per-token), observation ce weights, **or** model I/O — query a reference pool (e.g. `self.teacher_pool`, connected in `setup()` via `self.connect(...)`, or the live `self.policy_pool` for opsd) and attach per-token results (e.g. teacher logprobs) with bounded concurrency. No siblings. `echo` weights observation tokens here, identifying env-provided observation nodes by their non-sampled status and source step role attribution, applying the optional user filter, and writing the `ce_weights` stream. Model I/O runs *before* the pre-batch filters, so it pays compute on rollouts that may then be filtered out. -- `score_group(group)` — the cohort, **before filtering** (filters read the streams), synchronous: group-relative credit (GRPO/MaxRL baselines). `group` is a list of `Rollout`. +- `async score_rollout(rollout)` — one rollout, **on arrival** (as it's tokenized, before its group is complete): rollout-local credit (`rollout.assign_advantages(...)`, scalar broadcast or per-token), observation ce weights, **or** model I/O — query a reference pool (e.g. `self.teacher_pool`, connected in `setup()` via `self.connect(...)`, or the live `self.policy_pool` for opsd) and attach per-token results (e.g. teacher logprobs) with bounded concurrency. No siblings. `echo` weights observation tokens here, identifying env-provided observation nodes by their non-sampled status and source step role attribution, applying the optional user filter, and writing the `ce_weights` stream. Model I/O runs *before* the zero-advantage drop, so it pays compute on rollouts that may then be dropped. +- `score_group(group)` — the cohort, **before the zero-advantage drop** (the check reads the streams), synchronous: group-relative credit (GRPO/MaxRL baselines). `group` is a list of `Rollout`. The pipeline drives the hooks through two non-virtual methods it never looks inside: `algorithm.finalize_rollout(rollout)` per arrival (rollout-local scoring + reference I/O) and `algorithm.finalize_group(rollouts)` per group (scoring + wire stamping; after this the records are frozen — groups die at stamping). Sample construction (interleaving) is pure pipeline — observation-token provenance is available through structural attribution (`node.sampled`, `node.is_content`) for any algorithm that trains on env-provided tokens. @@ -305,8 +305,8 @@ The per-token training signal is set by `algo.type` and the [algorithm](#the-alg | `rae` | `rl` | Reward minus a per-agent EMA baseline (SPIRAL's role-conditioned advantage estimation) — for multi-agent self-play envs. | | `hierarchical_grpo` | `rl` | GRPO for proposer-solver envs: solvers are compared within one proposed problem, while proposers are compared across proposals. | | `echo` | `rl` + `ce` | Group-norm on action tokens, plus weighted CE on env-provided tokens selected by message role (each role's `alpha` is its ECHO λ), optionally narrowed by a user filter. | -| `opd` | `ref_kl` | On-policy distillation: per-token reverse KL to a reference model (`teacher`, an inline frozen hosted model), evaluated in the trainer from shipped reference logprobs. No credit — rollouts keep `advantages = None` (advantage-based filters never fire) and ship no advantage stream; `group_size` only fans out sampling. | -| `opsd` | `ref_kl` | SDFT: per-token reverse KL to a demo-conditioned reference. No credit — rollouts keep `advantages = None` (advantage-based filters never fire) and ship no advantage stream. | +| `opd` | `ref_kl` | On-policy distillation: per-token reverse KL to a reference model (`teacher`, an inline frozen hosted model), evaluated in the trainer from shipped reference logprobs. No credit — rollouts keep `advantages = None` (the zero-advantage drop never fires) and ship no advantage stream; `group_size` only fans out sampling. | +| `opsd` | `ref_kl` | SDFT: per-token reverse KL to a demo-conditioned reference. No credit — rollouts keep `advantages = None` (the zero-advantage drop never fires) and ship no advantage stream. | | `sft` | `ce` | Cross-entropy on the sampled tokens. Assigns no advantage — trains on every sampled token. | ### Default Advantage @@ -382,7 +382,7 @@ id = "null" type = "subprocess" ``` -`group_size` controls how many problems are proposed from each source task. `env.n` controls how many solvers attempt each proposed problem. If a comparison contains only one trace—for example, a solver when `env.n = 1`—its advantage is zero and the zero-advantage filter removes it. +`group_size` controls how many problems are proposed from each source task. `env.n` controls how many solvers attempt each proposed problem. If a comparison contains only one trace—for example, a solver when `env.n = 1`—its advantage is zero and the zero-advantage drop removes it. This algorithm is accepted only for proposer-solver envs. Use the env's `train_proposer` and `train_solver` settings if you want to train only one role. @@ -439,7 +439,7 @@ class MyAlgorithm(Algorithm): Add a typed `MyAlgoConfig` to `prime_rl.configs.algorithm` and its discriminated union, then register `"my_algo": MyAlgorithm` in `ALGORITHM_CLASSES`. Pick the hook by *when* your signal is ready: `score_rollout` for per-arrival credit or credit that needs a model call (it's `async`), `score_group` for group-relative credit (GRPO/MaxRL). `assign_advantages` takes a scalar (broadcast over the rollout's trainable tokens — the common case) or a full-length per-token list aligned to the concatenated sample token_ids (process rewards, step-level credit; `0.0` off-mask). -Each per-token list must match the rollout's completion-token count exactly — validated loudly when the view writes it. Advantage-based filters and metrics derive from the streams (the zero-advantage filter checks for all-zero streams; logged distributions use per-rollout means). Signals that depend on the live policy's weights (like OPD's reverse KL) cannot be precomputed here; those are reference-scoring algorithms, evaluated in the trainer. +Each per-token list must match the rollout's completion-token count exactly — validated loudly when the view writes it. The zero-advantage drop and metrics derive from the streams (the drop checks for all-zero streams; logged distributions use per-rollout means). Signals that depend on the live policy's weights (like OPD's reverse KL) cannot be precomputed here; those are reference-scoring algorithms, evaluated in the trainer. ### Reference Scoring @@ -454,30 +454,21 @@ type = "opsd" demo_key = "demonstration" ``` -Scoring runs at arrival, *before* the pre-batch filters, so a rollout that is later filtered still cost its reference compute — accepted for the simpler one-rollout-at-a-time shape (advantage-based filters never fire for opd/opsd anyway, since neither assigns an advantage). +Scoring runs at arrival, *before* the zero-advantage drop, so a rollout that is later dropped still cost its reference compute — accepted for the simpler one-rollout-at-a-time shape (the drop never fires for opd/opsd anyway, since neither assigns an advantage). ## Filters -Filters drop rollouts between scoring and training. Built-ins (composable): +Between scoring and training the sink runs three hardcoded checks on every trainable rollout: -| Filter | Effect | +| Check | Effect | |---|---| -| `gibberish` | Drops rollouts whose mean log-prob fall below a threshold — usually a sign of degenerate output. | -| `repetition` | Drops rollouts with high n-gram repetition. | -| `zero_advantage` | Drops rollouts whose advantage is zero, so the trainer doesn't waste tokens on them. | +| `gibberish` | Detects rare tokens generated at high entropy — usually a sign of degenerate output. Monitor-only: tracked in metrics (`gibberish/mean`, next to the other per-agent trace verdicts), never dropped. | +| `repetition` | Detects long high-confidence loops. Monitor-only: tracked in metrics (`repetition/mean`), never dropped. | +| `zero_advantage` | A rollout whose advantage stream is all zero (its whole group earned the same reward) carries no learning signal — dropped before it consumes batch budget, so the trainer never wastes tokens on it. | -The default `[orchestrator]` config registers all three in both filter slots: `post_batch_filters` enforce by default (flagged rollouts are recorded but not shipped to the trainer), while `pre_batch_filters` run in monitor mode (`enforce = false`); flip `enforce = true` there to drop matching rollouts before they consume a slot in the batch. Setting a slot replaces its defaults wholesale: +Zero-advantage rollouts are exempt when the env's algorithm declares `trains_on_zero_advantage` (echo: the `ce` component trains observation tokens regardless of credit); algorithms that assign no advantage at all (opd/opsd) never match. `orchestrator.count_zero_advantage_in_batch = true` makes dropped rollouts still count toward `batch_size` — a fixed sampling budget per step, at the cost of a variable number of trained-on samples. -```toml -[[orchestrator.post_batch_filters]] -type = "zero_advantage" - -[[orchestrator.post_batch_filters]] -type = "repetition" -threshold = 0.4 -``` - -Filtered rollouts still appear in W&B distributions, just not in the trainer batch — useful for spotting whether filtering is doing its job. +Dropped rollouts still appear in W&B distributions and metrics (`is_filtered`, `zero_advantage`), just not in the trainer batch. ## Multi-Turn Trajectories diff --git a/docs/overview.md b/docs/overview.md index b33c8a0736..8d53a51fe0 100644 --- a/docs/overview.md +++ b/docs/overview.md @@ -40,6 +40,6 @@ The `rl` entrypoint reads `examples/basic/reverse-text/rl.toml`, splits it into - **[Training](training.md)** — Launch and observe RL and SFT runs. - **[Inference](inference.md)** — vLLM-backed server (or fleet) holding the current policy. - **[Scaling](scaling.md)** — Single-GPU through multi-node clusters via FSDP / EP / CP and SLURM. -- **[Algorithms](algorithms.md)** — Async semantics, loss / advantage / filter plugins, trajectory merging. +- **[Algorithms](algorithms.md)** — Async semantics, loss / advantage plugins, rollout checks, trajectory merging. - **[Advanced](advanced.md)** — Custom modeling, multimodal, LoRA, P/D inference. - **[Development](development.md)** — Test suite, pre-commit hooks, adding a new model. diff --git a/docs/training.md b/docs/training.md index af0b084f89..a2a105d36b 100644 --- a/docs/training.md +++ b/docs/training.md @@ -57,7 +57,8 @@ A condensed view of the knobs you'll most often tune. For trainer-side paralleli | Knob | What it does | |---|---| -| `orchestrator.batch_size` | Tasks per trainer step. | +| `orchestrator.batch_size` | Rollouts to train on per step. | +| `orchestrator.max_inflight_episodes` | Concurrent episodes kept in-flight. Defaults to `batch_size`; raise above it to oversample ahead of the next batch. | | `orchestrator.group_size` | Rollouts generated per task. | | `orchestrator.max_off_policy_steps` | How many distinct policies may have contributed to one rollout before it's discarded (default 8). The main off-policy dial on long agentic rollouts — bump for throughput, lower for tighter on-policyness. Watch `errored_rollouts` and `mismatch_kl/all/mean` when tuning. | | `[orchestrator.algo]` | Training algorithm — its `type` names it (`grpo` default, `max_rl`, `rae`, `hierarchical_grpo`, `opd`, `opsd`, `sft`, `echo`). See [Algorithms](#algorithms). | diff --git a/examples/advanced/glm-5.2/swe.toml b/examples/advanced/glm-5.2/swe.toml index 657a488c9b..43ef6b8f5a 100644 --- a/examples/advanced/glm-5.2/swe.toml +++ b/examples/advanced/glm-5.2/swe.toml @@ -66,7 +66,7 @@ weight_decay = 0.1 [orchestrator] batch_size = 4096 group_size = 16 -oversampling_factor = 3 +max_inflight_episodes = 12_288 max_off_policy_steps = 16 [orchestrator.model] @@ -88,10 +88,6 @@ id = "bash" type = "prime" labels = ["glm5-pd-disag", "swe-bench-verified"] -[[orchestrator.post_batch_filters]] -type = "gibberish" -enforce = true - [inference] # we need <0.85 bc glm5 layers are too large for 0.85 use_deep_gemm = true diff --git a/examples/advanced/intellect-3.1/rl.toml b/examples/advanced/intellect-3.1/rl.toml index 43f7515568..41bd2fe6ee 100644 --- a/examples/advanced/intellect-3.1/rl.toml +++ b/examples/advanced/intellect-3.1/rl.toml @@ -45,7 +45,7 @@ weight_decay = 0.01 [orchestrator] batch_size = 2048 -oversampling_factor = 2 +max_inflight_episodes = 4096 [[orchestrator.train.source]] name = "swe" @@ -119,10 +119,6 @@ id = "null" [orchestrator.train.source.env.agent.runtime] type = "subprocess" -[[orchestrator.pre_batch_filters]] -type = "zero_advantage" -enforce = true - [orchestrator.eval] interval = 25 diff --git a/examples/advanced/minimax-m2.5/swe.toml b/examples/advanced/minimax-m2.5/swe.toml index 381a6c5486..35a19d026a 100644 --- a/examples/advanced/minimax-m2.5/swe.toml +++ b/examples/advanced/minimax-m2.5/swe.toml @@ -45,7 +45,7 @@ weight_decay = 0.01 [orchestrator] batch_size = 2048 -oversampling_factor = 2 +max_inflight_episodes = 4096 max_off_policy_steps = 16 [[orchestrator.train.source]] diff --git a/examples/advanced/qwen3-30b-a3b/math.toml b/examples/advanced/qwen3-30b-a3b/math.toml index ff2ac4512a..c8e83526e7 100644 --- a/examples/advanced/qwen3-30b-a3b/math.toml +++ b/examples/advanced/qwen3-30b-a3b/math.toml @@ -42,7 +42,7 @@ lr = 1e-6 [orchestrator] batch_size = 512 -oversampling_factor = 2 +max_inflight_episodes = 1024 max_off_policy_steps = 8 [orchestrator.train.sampling] diff --git a/examples/advanced/qwen3-30b-a3b/swe.toml b/examples/advanced/qwen3-30b-a3b/swe.toml index 3bd60b1255..547bc88242 100644 --- a/examples/advanced/qwen3-30b-a3b/swe.toml +++ b/examples/advanced/qwen3-30b-a3b/swe.toml @@ -43,7 +43,7 @@ lr = 1e-6 [orchestrator] batch_size = 512 -oversampling_factor = 2 +max_inflight_episodes = 1024 max_off_policy_steps = 16 [[orchestrator.train.source]] diff --git a/examples/basic/wiki-search/rl.toml b/examples/basic/wiki-search/rl.toml index f9f256387b..936970f258 100644 --- a/examples/basic/wiki-search/rl.toml +++ b/examples/basic/wiki-search/rl.toml @@ -32,7 +32,7 @@ target_modules = [ [orchestrator] batch_size = 512 group_size = 16 -oversampling_factor = 2.0 +max_inflight_episodes = 1024 [orchestrator.model.lora] name = "qwen3-4b-wiki-search" @@ -52,10 +52,6 @@ id = "null" [orchestrator.train.source.env.agent.runtime] type = "subprocess" -[[orchestrator.pre_batch_filters]] -type = "zero_advantage" -enforce = true - [ckpt] # Checkpoint at the end of training [inference.vllm] diff --git a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py index 74af6c0de8..86e5d54a81 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py @@ -302,50 +302,6 @@ class CheckpointConfig(BaseConfig): """Skip loading the progress from checkpoint.""" -# Flags rare tokens generated at high entropy (Section 5.2, https://arxiv.org/abs/2510.02387). -class GibberishFilterConfig(BaseConfig): - type: Literal["gibberish"] = "gibberish" - - enforce: bool = False - """When True, skip detected rollouts entirely so they are not sent to the trainer. When False, only track detection metrics.""" - - token_id_threshold: int = 100_000 - """Token IDs above this are candidates for gibberish. BPE tokens are sorted by merge order.""" - - logprob_offset: float = 2.0 - """Offset from uniform-distribution logprob. Threshold = ``-log(vocab_size) - logprob_offset``.""" - - -# Flags rollouts stuck in a repetition loop: emits high-confidence tokens for an extended stretch. -# Flagged when `window` consecutive tokens are each sampled with probability above `prob_threshold`. -# (Section 3.2, https://arxiv.org/abs/2506.13585) -class RepetitionFilterConfig(BaseConfig): - type: Literal["repetition"] = "repetition" - - enforce: bool = False - """When True, skip detected rollouts entirely so they are not sent to the trainer. When False, only track detection metrics.""" - - window: int = Field(3_000, ge=1) - """Consecutive high-probability steps required to flag the rollout.""" - - prob_threshold: float = Field(0.99, gt=0, le=1) - """Tokens sampled with probability above this are considered repetitive. Consecutive such tokens count toward the window.""" - - -# Flags rollouts with zero advantage. -class ZeroAdvantageFilterConfig(BaseConfig): - type: Literal["zero_advantage"] = "zero_advantage" - - enforce: bool = True - """When True, skip detected rollouts entirely so they are not sent to the trainer. When False, only track detection metrics.""" - - -FilterConfig: TypeAlias = Annotated[ - GibberishFilterConfig | RepetitionFilterConfig | ZeroAdvantageFilterConfig, - Field(discriminator="type"), -] - - class FileSystemWeightBroadcastConfig(BaseConfig): type: Literal["filesystem"] = "filesystem" @@ -416,24 +372,6 @@ class OrchestratorConfig(BaseConfig): eval: EvalConfig | None = None """Evaluation configuration.""" - pre_batch_filters: list[FilterConfig] = [ - GibberishFilterConfig(enforce=False), - RepetitionFilterConfig(enforce=False), - ZeroAdvantageFilterConfig(enforce=False), - ] - """Filters applied *before* a rollout enters the training batch buffer. - All three filter types are registered in monitor mode by default; flip ``enforce=true`` per type - to drop matching rollouts before they consume a slot in the batch (e.g. a zero-advantage group - never makes it into a training batch).""" - - post_batch_filters: list[FilterConfig] = [ - GibberishFilterConfig(), - RepetitionFilterConfig(), - ZeroAdvantageFilterConfig(), - ] - """Filters applied *after* a batch has been assembled. Each filter annotates each rollout; - rollouts flagged by an enforcing filter are still recorded but not shipped to the trainer.""" - log: LogConfig = LogConfig() env_vars: EnvVars = {} @@ -470,17 +408,14 @@ class OrchestratorConfig(BaseConfig): env_server_base_port: int = Field(5000, ge=1, le=65535) """First port of the env-server port range: the source at position ``i`` (train, then eval) is served at ``tcp://127.0.0.1:``. Sources with an explicit ``serve.address`` keep it instead, without shifting the other sources' ports (indices stay positional). Give concurrent runs on one host distinct bases (e.g. one per multi-run orchestrator).""" - batch_size: int | None = Field(None, ge=1) - """Samples to train on per step (rollout-based batching). Set this OR ``token_batch_size``.""" + batch_size: int = Field(128, ge=1) + """Rollouts to train on per step. Must be divisible by ``group_size``.""" - token_batch_size: int | None = Field(None, ge=1) - """Tokens to train on per step (token-based batching). Set this OR ``batch_size``.""" - - oversampling_factor: float | None = Field(None, gt=0) - """Rollout-mode batching only. Multiplier used to derive ``max_inflight_episodes`` from ``batch_size`` when ``max_inflight_episodes`` is unset. Values below 1.0 intentionally cap in-flight episode capacity below ``batch_size``.""" + count_zero_advantage_in_batch: bool = False + """Count zero-advantage rollouts toward ``batch_size`` (they are still not shipped to the trainer). By default the batch fills with informative samples only, which keeps the trained-on batch predictable but makes the per-step sampling time vary with the zero-advantage rate. Opt in to recover a fixed sampling budget per step at the cost of a variable number of trained-on samples.""" max_inflight_episodes: int | None = Field(None, ge=1) - """Maximum number of episodes kept in-flight — one episode is one agent run at a time, whatever the env's agents are. Required for token-based batching. With ``batch_size`` set, defaults to ``batch_size * oversampling_factor`` (or ``batch_size`` when ``oversampling_factor`` is unset).""" + """Maximum number of episodes kept in-flight — one episode is one agent run at a time, whatever the env's agents are. Defaults to ``batch_size``; raise above it to oversample ahead of the next batch.""" group_size: int = Field(1, ge=1) """Output sequences returned per example during training.""" @@ -530,16 +465,6 @@ def auto_setup_prime_monitor_run_name(self): self.prime_monitor.run_name = self.wandb.name return self - @model_validator(mode="after") - def validate_unique_filter_types(self): - for slot_name in ("pre_batch_filters", "post_batch_filters"): - types = [f.type for f in getattr(self, slot_name)] - if len(types) != len(set(types)): - raise ValueError( - f"Duplicate filter types in {slot_name}: {types}. Each filter type may only appear once per slot." - ) - return self - @model_validator(mode="after") def inherit_env_algorithms(self): """Envs without their own algorithm inherit the top-level one. @@ -598,37 +523,11 @@ def validate_renderer_auto_resolves(self): @model_validator(mode="after") def resolve_batching(self): - has_rollout_batch = self.batch_size is not None - has_token_batch = self.token_batch_size is not None - - if has_rollout_batch and has_token_batch: - raise ValueError("Set exactly one of batch_size or token_batch_size") - - if not has_rollout_batch and not has_token_batch: - self.batch_size = 128 - - if has_token_batch: - if self.oversampling_factor is not None: - raise ValueError("oversampling_factor can only be set when batch_size is set") - if self.max_inflight_episodes is None: - raise ValueError("max_inflight_episodes must be set when token_batch_size is set") - else: - assert self.batch_size is not None - if self.batch_size % self.group_size != 0: - raise ValueError("Batch size must be divisible by the number of samples per problem") - oversampling_factor = self.oversampling_factor if self.oversampling_factor is not None else 1.0 - resolved_max_inflight_episodes = max( - self.group_size, - int(self.batch_size * oversampling_factor), - ) - if self.max_inflight_episodes is not None and self.oversampling_factor is not None: - expected_max_inflight_episodes = resolved_max_inflight_episodes - if self.max_inflight_episodes != expected_max_inflight_episodes: - raise ValueError("max_inflight_episodes conflicts with oversampling_factor * batch_size") - if self.max_inflight_episodes is None: - self.max_inflight_episodes = resolved_max_inflight_episodes - - if self.max_inflight_episodes is not None and self.max_inflight_episodes < self.group_size: + if self.batch_size % self.group_size != 0: + raise ValueError("batch_size must be divisible by group_size") + if self.max_inflight_episodes is None: + self.max_inflight_episodes = self.batch_size + if self.max_inflight_episodes < self.group_size: raise ValueError("max_inflight_episodes must be at least the number of rollouts per example") # Propagate the top-level ``group_size`` into each train env that didn't set its own. diff --git a/packages/prime-rl-configs/src/prime_rl/configs/rl.py b/packages/prime-rl-configs/src/prime_rl/configs/rl.py index 6d70c3fae1..9a23524c5a 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/rl.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/rl.py @@ -469,7 +469,7 @@ def auto_setup_bench(self): self.trainer.bench = BenchConfig() self.orchestrator.bench = True self.trainer.data.fake = FakeDataLoaderConfig( - batch_size=self.orchestrator.batch_size or 32, + batch_size=self.orchestrator.batch_size, ) trainer_bench_enabled = self.trainer.bench is not None diff --git a/skills/training/monitor-run/SKILL.md b/skills/training/monitor-run/SKILL.md index 084d1610dd..929a158911 100644 --- a/skills/training/monitor-run/SKILL.md +++ b/skills/training/monitor-run/SKILL.md @@ -104,7 +104,7 @@ All metrics print to the console log (and W&B when configured). | `train//effective//num_turns/mean` | avg turns for that agent alone (also token counts, `num_branches`) | | `train/agg/effective//is_truncated/mean` | fraction of that agent's rollouts truncated | | `train/agg/all//has_error/mean` | fraction of that agent's rollouts errored (per-type under `train/agg/all//error/`; also `dispatcher/errored/{train,eval}`) | -| `train/agg/all//is_trainable/mean` | fraction carrying a training signal — 0.0 for a frozen seat like a judge (also `is_filtered`, `filters/`) | +| `train/agg/all//is_trainable/mean` | fraction carrying a training signal — 0.0 for a frozen seat like a judge (also `is_filtered`, and the check verdicts `gibberish`, `repetition`, `zero_advantage`) | | `train//effective//metrics//mean` | env-specific metrics for that agent (e.g. pass rate) | | `train//effective//timing/agent/model/mean` | model vs harness share of that agent's phase | | `eval//effective//{avg@k,pass@k}` | eval scores for that agent, when configured | diff --git a/src/prime_rl/orchestrator/algo/base.py b/src/prime_rl/orchestrator/algo/base.py index 3540864277..167dff4889 100644 --- a/src/prime_rl/orchestrator/algo/base.py +++ b/src/prime_rl/orchestrator/algo/base.py @@ -19,8 +19,9 @@ I/O against another model — an inference pool the algorithm connected in ``setup()`` (a frozen teacher) or the live policy (opsd's self-distillation), queried with bounded concurrency. No siblings. -- ``score_group(group)`` — the cohort, on group completion, *before* filtering - (filters read the streams): group-relative credit (GRPO/MaxRL baselines). +- ``score_group(group)`` — the cohort, on group completion, *before* the + zero-advantage drop (the check reads the streams): group-relative credit + (GRPO/MaxRL baselines). How rollouts are *produced* is not the algorithm's concern: that is the env's :class:`~prime_rl.orchestrator.sampler.Sampler`, and sample construction @@ -95,21 +96,21 @@ class Algorithm: directly — read the trace, write credit via :meth:`Rollout.assign_advantages`. They are async so either stage may do I/O — e.g. a process-reward model or a - teacher at arrival, or a judge at group time whose signal a pre-batch - filter then reads; a hook that only does advantage math simply never - awaits. + teacher at arrival, or a judge at group time whose signal the + zero-advantage drop then reads; a hook that only does advantage math + simply never awaits. - :meth:`score_rollout` — one rollout, on arrival: rollout-local credit, observation ce weights, or per-token results from a model the algorithm connected in :meth:`setup` (e.g. teacher reference logprobs). Default: nothing. - - :meth:`score_group` — the cohort, *before* filtering (filters read the - streams): group-relative credit. Default: nothing — rollouts keep - ``advantages=None``, so advantage-based filters skip them. + - :meth:`score_group` — the cohort, *before* the zero-advantage drop + (the check reads the streams): group-relative credit. Default: + nothing — rollouts keep ``advantages=None``, so the drop skips them. Model I/O lives in :meth:`score_rollout`: it runs at arrival, *before* the - pre-batch filters, so it pays compute on rollouts that may then be filtered - out — accepted for the simpler one-rollout-at-a-time shape. + zero-advantage drop, so it pays compute on rollouts that may then be + dropped — accepted for the simpler one-rollout-at-a-time shape. Constructed with the algorithm config it interprets plus the live policy pool (``self.policy_pool`` — always available, never closed by the @@ -119,6 +120,12 @@ class Algorithm: action_loss_type: ClassVar[ActionLossType] = "rl" + trains_on_zero_advantage: ClassVar[bool] = False + """True when the algorithm still extracts training signal from a rollout + whose advantage stream is all zero (echo trains observation tokens through + the ``ce`` component regardless of credit). Such rollouts bypass the + pipeline's zero-advantage drop and ship to the trainer.""" + def __init__(self, config: AlgoConfig, policy_pool: InferencePool): self.policy_pool = policy_pool self.connected_pools: list[InferencePool] = [] # frozen pools connected in setup(); closed at shutdown @@ -145,8 +152,8 @@ async def score_rollout(self, rollout: Rollout) -> None: group stats.""" async def score_group(self, group: list[Rollout]) -> None: - """Group phase, the finalized cohort, before filtering: write - group-relative credit.""" + """Group phase, the finalized cohort, before the zero-advantage drop: + write group-relative credit.""" async def finalize_rollout(self, rollout: Rollout) -> None: """Arrival phase (non-virtual): rollout-local scoring as each rollout is diff --git a/src/prime_rl/orchestrator/algo/echo.py b/src/prime_rl/orchestrator/algo/echo.py index d4ecf74fa3..3f3a41e85c 100644 --- a/src/prime_rl/orchestrator/algo/echo.py +++ b/src/prime_rl/orchestrator/algo/echo.py @@ -23,6 +23,10 @@ class EchoAlgorithm(GRPOAlgorithm): mask and its denominator. An optional user filter narrows the selection per rollout (e.g. dropping tool-output warnings).""" + # The ce component trains observation tokens even when the GRPO advantage + # collapses to zero, so such rollouts still carry signal. + trains_on_zero_advantage = True + def __init__(self, config: EchoAlgoConfig, policy_pool: InferencePool): super().__init__(config, policy_pool) self.role_weights = { diff --git a/src/prime_rl/orchestrator/algo/max_rl.py b/src/prime_rl/orchestrator/algo/max_rl.py index 9a3978108d..9ee2529f23 100644 --- a/src/prime_rl/orchestrator/algo/max_rl.py +++ b/src/prime_rl/orchestrator/algo/max_rl.py @@ -20,8 +20,8 @@ class MaxRLAlgorithm(Algorithm): likelihood as it grows). Assumes non-negative (canonically binary) rewards; a group with mean reward - <= 0 carries no signal and gets zero advantages (the zero-advantage filter - drops it, matching the paper's no-success convention).""" + <= 0 carries no signal and gets zero advantages (the zero-advantage drop + removes it, matching the paper's no-success convention).""" async def score_group(self, group: list[Rollout]) -> None: rewards = torch.tensor([rollout.reward for rollout in group], dtype=torch.float32) diff --git a/src/prime_rl/orchestrator/algo/opd.py b/src/prime_rl/orchestrator/algo/opd.py index a4b3a06f72..e2812644fa 100644 --- a/src/prime_rl/orchestrator/algo/opd.py +++ b/src/prime_rl/orchestrator/algo/opd.py @@ -19,8 +19,8 @@ class OPDAlgorithm(Algorithm): The policy samples its own rollouts; at ship time each sample's full context is prefill-scored under the teacher (``ref_logprobs`` on the wire), and the trainer evaluates the KL against the live policy. No - credit is assigned — rollouts keep ``advantages=None`` (advantage-based - filters never fire) and samples ship no advantage stream; ``group_size`` + credit is assigned — rollouts keep ``advantages=None`` (the + zero-advantage drop never fires) and samples ship no advantage stream; ``group_size`` only fans out sampling.""" action_loss_type = "ref_kl" diff --git a/src/prime_rl/orchestrator/algo/sft.py b/src/prime_rl/orchestrator/algo/sft.py index c8c51221f3..744439c9bf 100644 --- a/src/prime_rl/orchestrator/algo/sft.py +++ b/src/prime_rl/orchestrator/algo/sft.py @@ -8,7 +8,6 @@ class SFTDistillAlgorithm(Algorithm): rollouts (``sampling.source``); the policy trains with CE on its tokens. Assigns no advantage — the ``ce`` loss ignores credit, and SFT trains on - every sampled token. Reward-based filtering, if wanted, is an explicit - filter, not smuggled through an unused advantage stream.""" + every sampled token (the zero-advantage drop never fires).""" action_loss_type = "ce" diff --git a/src/prime_rl/orchestrator/filters.py b/src/prime_rl/orchestrator/filters.py index ad023fd928..ddc88e1ea2 100644 --- a/src/prime_rl/orchestrator/filters.py +++ b/src/prime_rl/orchestrator/filters.py @@ -1,172 +1,71 @@ -"""Orchestrator-side rollout filters for detecting degenerate generations. +"""Hardcoded rollout checks between scoring and training. -Filters run after rollouts complete, inspecting token IDs and logprobs to -detect gibberish or repetition. Detection metrics are always tracked. -When enforce=True, detected rollouts are skipped entirely during training and -are not sent to the trainer. Reward is kept as-is for baseline calculation. +Gibberish and repetition detection runs on every trainable rollout and is +tracked in metrics only — a detection never drops a rollout. Zero-advantage +rollouts carry no learning signal (unless the env's algorithm says otherwise) +and are dropped before they enter the training batch. """ from __future__ import annotations import math -from dataclasses import dataclass -from typing import TYPE_CHECKING, Protocol - -from prime_rl.configs.orchestrator import FilterConfig -from prime_rl.utils.logger import get_logger +from typing import TYPE_CHECKING if TYPE_CHECKING: from prime_rl.orchestrator.types import Rollout - -@dataclass -class FilterResult: - detected: bool - - -class RolloutFilter(Protocol): - name: str - enforce: bool - - def check(self, rollout: Rollout) -> FilterResult: ... - - -@dataclass -class GibberishFilter: - """Flags rollouts containing rare tokens generated at high entropy. - - A token is flagged when both: - - id(token) > token_id_threshold (rare BPE token) - - logprob(token) < -log(vocab_size) - logprob_offset (high entropy) - - References: - Section 5.2, https://arxiv.org/abs/2510.02387 - """ - - name: str - token_id_threshold: int - logprob_threshold: float - enforce: bool = False - - def check(self, rollout: Rollout) -> FilterResult: - for branch in rollout.branches: - # branch.{token_ids,logprobs,sampled_mask} are flat and mutually aligned; the raw - # node arrays are not (node.logprobs covers only the sampled suffix, not the - # generation-prompt scaffold that token_ids/mask also span). - for token_id, logprob, sampled in zip(branch.token_ids, branch.logprobs, branch.sampled_mask): - if not sampled: - continue - if token_id > self.token_id_threshold and logprob < self.logprob_threshold: - return FilterResult(detected=True) - return FilterResult(detected=False) - - -@dataclass -class RepetitionFilter: - """Flags rollouts with pathological repetition loops. - - Counts consecutive tokens where logprob > log(prob_threshold), indicating - the model is generating with very high confidence. When the streak reaches - the window size, the rollout is flagged. - - References: - Section 3.2, https://arxiv.org/abs/2506.13585 - """ - - name: str - window: int - logprob_threshold: float - enforce: bool = False - - def check(self, rollout: Rollout) -> FilterResult: - for branch in rollout.branches: - # Aligned branch streams (see GibberishFilter), and reset the streak per branch: - # flat rollout.nodes interleaves distinct root->leaf paths (compaction/subagents), - # so a per-node walk would run a streak across a branch boundary. - consecutive = 0 - for logprob, sampled in zip(branch.logprobs, branch.sampled_mask): - if not sampled: - continue - if logprob > self.logprob_threshold: - consecutive += 1 - else: - consecutive = 0 - if consecutive >= self.window: - return FilterResult(detected=True) - return FilterResult(detected=False) - - -@dataclass -class ZeroAdvantageFilter: - """Flags rollouts whose advantage stream is all zero (e.g. all rollouts in - a GRPO group earned the same reward, so the centered advantage collapses).""" - - name: str - enforce: bool = True - - def check(self, rollout: Rollout) -> FilterResult: - if rollout.advantages is not None and all(a == 0.0 for a in rollout.advantages): - return FilterResult(detected=True) - return FilterResult(detected=False) - - -def setup_filter(config: FilterConfig, vocab_size: int) -> RolloutFilter: - """Create a RolloutFilter from a filter config.""" - if config.type == "gibberish": - return GibberishFilter( - name="gibberish", - token_id_threshold=config.token_id_threshold, - logprob_threshold=-math.log(vocab_size) - config.logprob_offset, - enforce=config.enforce, - ) - elif config.type == "repetition": - return RepetitionFilter( - name="repetition", - window=config.window, - logprob_threshold=math.log(config.prob_threshold), - enforce=config.enforce, - ) - elif config.type == "zero_advantage": - return ZeroAdvantageFilter( - name="zero_advantage", - enforce=config.enforce, - ) - raise ValueError(f"Unknown filter type: {config.type}") - - -def setup_filters(configs: list[FilterConfig], vocab_size: int, *, kind: str) -> list[RolloutFilter]: - """Create RolloutFilters from a list of filter configs.""" - filters = [setup_filter(config, vocab_size) for config in configs] - if filters: - get_logger().info(f"Configured {len(filters)} {kind} rollout filter(s):") - for config, filt in zip(configs, filters): - mode = "Enforcing" if filt.enforce else "Monitoring" - params = ", ".join(f"{k}={v}" for k, v in config.model_dump().items()) - get_logger().info(f" {mode} {filt.name} filter ({params})") - return filters - - -def apply_filters(filters: list[RolloutFilter], rollouts: list[Rollout]) -> None: - """Flag ``Rollout``\\ s in place with per-filter detection + drop decision. - - Each rollout's ``filter_results`` dict records per-filter detection bools; - ``is_filtered`` is True iff an enforcing filter detected it. First matching - filter wins per rollout (no double-counting). Reward and trajectory tokens - are left untouched so the rollout can still contribute to baseline - calculations and metric aggregation. - """ - for rollout in rollouts: - rollout.filter_results = {f.name: False for f in filters} - rollout.is_filtered = False - - if not filters: - return - - for rollout in rollouts: - for filt in filters: - result = filt.check(rollout) - if result.detected: - rollout.filter_results[filt.name] = True - if filt.enforce: - rollout.is_filtered = True - break +# Gibberish: rare tokens generated at high entropy (Section 5.2, +# https://arxiv.org/abs/2510.02387). A token is flagged when its id exceeds +# the threshold (rare BPE token, sorted by merge order) and its logprob is +# below ``-log(vocab_size) - offset`` (high entropy). +GIBBERISH_TOKEN_ID_THRESHOLD = 100_000 +GIBBERISH_LOGPROB_OFFSET = 2.0 + +# Repetition: pathological high-confidence loops (Section 3.2, +# https://arxiv.org/abs/2506.13585). Flagged when ``WINDOW`` consecutive +# tokens are each sampled with probability above ``PROB_THRESHOLD``. +REPETITION_WINDOW = 3_000 +REPETITION_PROB_THRESHOLD = 0.99 + + +def gibberish_logprob_threshold(vocab_size: int) -> float: + return -math.log(vocab_size) - GIBBERISH_LOGPROB_OFFSET + + +def detect_gibberish(rollout: Rollout, logprob_threshold: float) -> bool: + for branch in rollout.branches: + # branch.{token_ids,logprobs,sampled_mask} are flat and mutually aligned; the raw + # node arrays are not (node.logprobs covers only the sampled suffix, not the + # generation-prompt scaffold that token_ids/mask also span). + for token_id, logprob, sampled in zip(branch.token_ids, branch.logprobs, branch.sampled_mask): + if not sampled: + continue + if token_id > GIBBERISH_TOKEN_ID_THRESHOLD and logprob < logprob_threshold: + return True + return False + + +def detect_repetition(rollout: Rollout) -> bool: + logprob_threshold = math.log(REPETITION_PROB_THRESHOLD) + for branch in rollout.branches: + # Aligned branch streams (see detect_gibberish), and reset the streak per branch: + # flat rollout.nodes interleaves distinct root->leaf paths (compaction/subagents), + # so a per-node walk would run a streak across a branch boundary. + consecutive = 0 + for logprob, sampled in zip(branch.logprobs, branch.sampled_mask): + if not sampled: + continue + if logprob > logprob_threshold: + consecutive += 1 + else: + consecutive = 0 + if consecutive >= REPETITION_WINDOW: + return True + return False + + +def has_zero_advantage(rollout: Rollout) -> bool: + """True when the advantage stream is present but all zero (e.g. all + rollouts in a GRPO group earned the same reward, so the centered advantage + collapses). Algorithms that assign no advantage (opd/opsd) never match.""" + return rollout.advantages is not None and all(a == 0.0 for a in rollout.advantages) diff --git a/src/prime_rl/orchestrator/metrics.py b/src/prime_rl/orchestrator/metrics.py index 190a53601f..15dd2200e0 100644 --- a/src/prime_rl/orchestrator/metrics.py +++ b/src/prime_rl/orchestrator/metrics.py @@ -351,7 +351,7 @@ def to_wandb(self, *, prefix: str, subset: Subset) -> dict[str, float]: out[f"{p}/is_filtered/mean"] = sum(float(r.is_filtered) for r in rollouts) / len(rollouts) names = sorted({name for r in rollouts for name in r.filter_results}) out |= { - f"{p}/filters/{name}/mean": sum(1 for r in rollouts if r.filter_results.get(name)) / len(rollouts) + f"{p}/{name}/mean": sum(1 for r in rollouts if r.filter_results.get(name)) / len(rollouts) for name in names } return out diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index 4a4d99a179..4fc776928d 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -5,8 +5,8 @@ - ``RolloutDispatcher`` schedules rollouts; emits ``Rollout`` (train/eval discriminated by ``kind``) on its queue. -- ``TrainSink`` ingests train rollouts (tokenize → advantages → filters) - and returns a ``TrainBatch`` when the threshold is met. +- ``TrainSink`` ingests train rollouts (tokenize → advantages → zero-advantage + drop) and returns a ``TrainBatch`` when the batch is full. - ``EvalSink`` ingests eval rollouts and returns an ``EvalBatch`` (the full returned cohort) on epoch completion. - ``TrainRollouts`` / ``EvalRollouts`` carry the rollouts and build the per-step W&B metrics @@ -46,7 +46,6 @@ from prime_rl.orchestrator.envs import EvalEnvs, TrainEnvs from prime_rl.orchestrator.eval_sink import EvalSink from prime_rl.orchestrator.eval_source import EvalSource -from prime_rl.orchestrator.filters import setup_filters from prime_rl.orchestrator.inference_metrics import InferenceMetricsCollector from prime_rl.orchestrator.packing import BatchPacker from prime_rl.orchestrator.patches import ( @@ -95,9 +94,9 @@ # shutdown wedges (env-server ZMQ recv, vLLM admin aclose, etc) SHUTDOWN_TIMEOUT_S = 300 -# Abort after this many consecutive train batches drop all rollouts to -# post-batch filters — usually a misconfigured filter or homogeneous-reward -# dataset; fail loudly instead of spinning +# Abort after this many consecutive train batches ship no samples (possible +# when ``count_zero_advantage_in_batch`` lets zero-advantage rollouts fill the +# budget) — usually a homogeneous-reward dataset; fail loudly instead of spinning MAX_CONSECUTIVE_EMPTY_BATCHES = 10 # Maximum batches the orchestrator may run ahead of the trainer. The @@ -247,10 +246,6 @@ async def setup(self) -> None: if usage_base_url and usage_api_key: self.usage_reporter = UsageReporter() - # Filters apply to train rollouts only - pre_filters = setup_filters(config.pre_batch_filters, vocab_size=self.tokenizer.vocab_size, kind="pre-batch") - post_filters = setup_filters(config.post_batch_filters, vocab_size=self.tokenizer.vocab_size, kind="post-batch") - get_logger().info("Loading training environments") self.train_envs = TrainEnvs( config.train.source, @@ -396,7 +391,7 @@ async def setup(self) -> None: else None ) - assert config.max_inflight_episodes is not None, "max_inflight_episodes must be resolved before dispatcher init" + assert config.max_inflight_episodes is not None # resolved at config validation log_interval = config.log.interval wandb_enabled = config.wandb is not None self.dispatcher = RolloutDispatcher( @@ -415,10 +410,6 @@ async def setup(self) -> None: tokenizer=self.tokenizer, train_envs=self.train_envs, mm_token_type_ids_mapping=self.mm_token_type_ids_mapping, - batch_size=config.batch_size, - token_batch_size=config.token_batch_size, - pre_filters=pre_filters, - post_filters=post_filters, ) self.eval_sink = EvalSink(eval_envs=self.eval_envs) if self.eval_envs is not None else None self.watcher = WeightWatcher( @@ -601,8 +592,7 @@ async def finalize_train_batch(self, batch: TrainBatch) -> None: ) if self.consecutive_empty_batches >= MAX_CONSECUTIVE_EMPTY_BATCHES: raise RuntimeError( - f"{self.consecutive_empty_batches} consecutive empty train batches — " - "check filter config (pre_batch_filters / post_batch_filters) or task difficulty." + f"{self.consecutive_empty_batches} consecutive empty train batches — check task difficulty." ) return self.consecutive_empty_batches = 0 @@ -611,7 +601,7 @@ async def finalize_train_batch(self, batch: TrainBatch) -> None: if effective and n_trainable / len(effective) <= 0.1: get_logger().warning( f"Only {n_trainable}/{len(effective)} effective rollouts are trainable " - f"({n_trainable / len(effective):.1%}) — consider reviewing task difficulty / filter config" + f"({n_trainable / len(effective):.1%}) — consider reviewing task difficulty" ) # Ship batch ``step`` only once the trainer has published v{step-1-TARGET_LAG}. @@ -669,7 +659,7 @@ async def finalize_train_batch(self, batch: TrainBatch) -> None: for env_name, env_pool in pool.by_env().items(): metrics |= env_pool.metrics.to_wandb(prefix=f"train/{env_name}", subset=subset) - # Progress / timing / env-share / pre-filter accounting (assembled here, not in the metrics + # Progress / timing / env-share accounting (assembled here, not in the metrics # objects). ``num_tokens`` is over the full arrival window; the input/output breakdown is over # the effective (shipped) subset, summing the same ``vf.Trace`` token properties the metric # matrix reports. @@ -695,12 +685,6 @@ async def finalize_train_batch(self, batch: TrainBatch) -> None: } for env_name, env_pool in batch.rollouts.by_env().items(): metrics[f"batch/{env_name}"] = len(env_pool) / len(batch.rollouts) - if self.train_sink.pre_filter_seen > 0: - metrics["pre_filters/all/dropped_rate"] = ( - self.train_sink.pre_filter_dropped / self.train_sink.pre_filter_seen - ) - for name, count in self.train_sink.pre_filter_dropped_by_name.items(): - metrics[f"pre_filters/all/{name}/rate"] = count / self.train_sink.pre_filter_seen self.monitor.log(metrics, step=step) self.wait_for_policy_time = 0.0 self.monitor.log_samples(effective.rollouts, step=step) @@ -729,7 +713,6 @@ async def finalize_train_batch(self, batch: TrainBatch) -> None: self.log_train_batch(batch, step=step, step_time=step_time) - self.train_sink.reset_pre_filter_stats() self.maybe_trigger_eval(self.progress.step) trim_process_memory() @@ -769,7 +752,7 @@ def collect_pipeline_view(self) -> tuple[str, dict[str, float]]: inflight_by_env = self.dispatcher.inflight_by_env inflight_train = self.dispatcher.inflight_train_count inflight_eval = self.dispatcher.inflight_eval_count - train_batch, train_target, _train_unit = self.train_sink.batch_progress() + train_batch, train_target = self.train_sink.batch_progress() train_buffered = self.train_sink.buffered_count() train_batch_by_env = self.train_sink.pending_batch_by_env() eval_batches = self.eval_sink.batch_progress() if self.eval_sink is not None else [] diff --git a/src/prime_rl/orchestrator/train_sink.py b/src/prime_rl/orchestrator/train_sink.py index 5f052c14ed..5093635d8e 100644 --- a/src/prime_rl/orchestrator/train_sink.py +++ b/src/prime_rl/orchestrator/train_sink.py @@ -6,9 +6,10 @@ and untrainable rollouts skip this. 2. ``process_group`` — filters errored rollouts, hands the trainable survivors to the env algorithm's ``finalize_group`` (advantages + - per-sample wire stamping), runs the pre-batch filter pass. -3. ``process_batch`` — applies post-batch filter annotations and assembles - the trainer-bound ``TrainingSample`` list. Returns a ``TrainBatch``. + per-sample wire stamping), annotates degeneration detections, and drops + zero-advantage rollouts before they consume batch budget. +3. ``process_batch`` — pops a ``batch_size`` cohort and assembles the + trainer-bound ``TrainingSample`` list. Returns a ``TrainBatch``. ``add()`` takes one episode (``list[Rollout]``) and returns ``TrainBatch | None``; group accounting counts episodes, never loose traces. @@ -24,25 +25,22 @@ from prime_rl.configs.orchestrator import OrchestratorConfig from prime_rl.orchestrator.envs import TrainEnvs -from prime_rl.orchestrator.filters import RolloutFilter, apply_filters +from prime_rl.orchestrator.filters import ( + detect_gibberish, + detect_repetition, + gibberish_logprob_threshold, + has_zero_advantage, +) from prime_rl.orchestrator.metrics import TrainRollouts from prime_rl.orchestrator.trajectories import trace_to_samples from prime_rl.orchestrator.types import Rollout, TrainBatch from prime_rl.transport import TrainingSample from prime_rl.utils.logger import get_logger - -def payload_tokens(rollout: Rollout) -> int: - """Token cost of the rollout's trainer-bound payload — the samples built by - ``process_rollout``. This is what actually ships: forked traces can drop - branches with no trainable tokens, so ``Trace.num_total_tokens`` (which sums - over all branches) may overcount. For linear traces the two agree. - - Zero-payload rollouts (no trainable samples at all) fall back to the trace - total so they still advance token batching — a degenerate all-zero-payload - stream then ships empty batches and trips the orchestrator's - consecutive-empty-batch abort instead of stalling the readiness check.""" - return sum(len(sample.token_ids) for sample in rollout.samples) or rollout.num_total_tokens +# Warn every N consecutive finalized groups whose survivors were all dropped +# as zero-advantage — the batch isn't filling, usually a task-difficulty +# mismatch (rewards are homogeneous within every group) +ZERO_ADVANTAGE_STALL_WARN_GROUPS = 25 class TrainSink: @@ -55,22 +53,14 @@ def __init__( tokenizer, train_envs: TrainEnvs, mm_token_type_ids_mapping: dict[int, int] | None, - batch_size: int | None, - token_batch_size: int | None, - pre_filters: list[RolloutFilter], - post_filters: list[RolloutFilter], ) -> None: - assert (batch_size is None) != (token_batch_size is None), ( - "Exactly one of batch_size / token_batch_size must be set" - ) self.config = config self.tokenizer = tokenizer self.train_envs = train_envs self.mm_token_type_ids_mapping = mm_token_type_ids_mapping - self.batch_size = batch_size - self.token_batch_size = token_batch_size - self.pre_filters = pre_filters - self.post_filters = post_filters + self.batch_size = config.batch_size + self.count_zero_advantage_in_batch = config.count_zero_advantage_in_batch + self.gibberish_logprob_threshold = gibberish_logprob_threshold(tokenizer.vocab_size) # Observation window for the next shipped batch: rollouts of groups # finalized since the last ship (errored + filtered + survivors). @@ -84,28 +74,19 @@ def __init__( # add several traces to ``pending_groups`` but counts once here). self.pending_group_episodes: dict[uuid.UUID, int] = defaultdict(int) self.pending_batch: list[Rollout] = [] - # Running payload-token total of ``pending_batch`` (token-batched - # runs), kept in sync on append/pop so the readiness check never - # re-sums per arrival. - self.pending_tokens: int = 0 - - # Reset by the orchestrator after each ship via ``reset_pre_filter_stats`` - self.pre_filter_seen = 0 - self.pre_filter_dropped = 0 - self.pre_filter_dropped_by_name: dict[str, int] = {} + # Consecutive finalized groups that contributed nothing to + # ``pending_batch`` because every survivor was zero-advantage + self.consecutive_zero_advantage_groups = 0 def group_size_for(self, env_name: str) -> int: return self.train_envs.get(env_name).config.group_size - def batch_progress(self) -> tuple[int, int, str]: - """``(current, target, unit)`` for the train batch — counts only + def batch_progress(self) -> tuple[int, int]: + """``(current, target)`` rollouts for the train batch — counts only ``pending_batch`` (survivors of finalized groups, queued for the trainer), so it's an honest 0→target fill. Partial-group arrivals are reported separately by ``buffered_count()``.""" - if self.batch_size is not None: - return len(self.pending_batch), self.batch_size, "rollouts" - assert self.token_batch_size is not None - return self.pending_tokens, self.token_batch_size, "tokens" + return len(self.pending_batch), self.batch_size def buffered_count(self) -> int: """Episodes that have arrived but sit in not-yet-complete groups — @@ -137,12 +118,7 @@ async def add(self, episode: list[Rollout]) -> TrainBatch | None: # ``pending_batch`` only grows on group finalization, so readiness is # only re-checked here — the window of a shipped batch then always # contains at least the group that finalized it. - ready = ( - len(self.pending_batch) >= self.batch_size - if self.batch_size is not None - else self.pending_tokens >= (self.token_batch_size or 0) - ) - if ready: + if len(self.pending_batch) >= self.batch_size: return self.process_batch() return None @@ -167,7 +143,8 @@ async def process_rollout(self, rollout: Rollout) -> None: async def process_group(self, group_id: uuid.UUID) -> None: """Finalize one GRPO group: drop errored rollouts, assign advantages, - run pre-batch filters, append survivors to ``pending_batch``.""" + annotate detections, append the informative survivors to + ``pending_batch``.""" group = self.pending_groups.pop(group_id, []) self.pending_group_episodes.pop(group_id, None) if not group: @@ -205,65 +182,57 @@ async def process_group(self, group_id: uuid.UUID) -> None: for sample in r.samples: sample.temperatures = [temperature] * len(sample.token_ids) - if self.pre_filters: - apply_filters(self.pre_filters, survivors) - filtered_by_name: dict[str, int] = {} - num_filtered = 0 + # Degeneration detection is monitor-only (metrics); the zero-advantage + # check drops — a rollout whose advantage stream is all zero carries no + # learning signal, unless the env's algorithm trains without one (echo). + num_zero_advantage = 0 + appended = 0 for r in survivors: - self.pre_filter_seen += 1 + r.filter_results = { + "gibberish": detect_gibberish(r, self.gibberish_logprob_threshold), + "repetition": detect_repetition(r), + "zero_advantage": has_zero_advantage(r), + } + r.is_filtered = r.filter_results["zero_advantage"] and not env.algorithm.trains_on_zero_advantage if r.is_filtered: - self.pre_filter_dropped += 1 - num_filtered += 1 - for name, hit in r.filter_results.items(): - if hit: - self.pre_filter_dropped_by_name[name] = self.pre_filter_dropped_by_name.get(name, 0) + 1 - filtered_by_name[name] = filtered_by_name.get(name, 0) + 1 - continue - # Reset annotations so the post-batch filter pass starts clean - r.filter_results = {} - r.is_filtered = False + num_zero_advantage += 1 + # Opt-in: a dropped rollout still occupies a batch slot, so the + # per-step sampling effort stays fixed while the trained-on + # sample count varies with the zero-advantage rate. + if not self.count_zero_advantage_in_batch: + continue self.pending_batch.append(r) - if self.token_batch_size is not None: - self.pending_tokens += payload_tokens(r) + appended += 1 - # Per-group summary. One line per finalized group; per-filter - # detection breakdown lives at debug level in ``apply_filters`` + if appended: + self.consecutive_zero_advantage_groups = 0 + else: + self.consecutive_zero_advantage_groups += 1 + if self.consecutive_zero_advantage_groups % ZERO_ADVANTAGE_STALL_WARN_GROUPS == 0: + get_logger().warning( + f"{self.consecutive_zero_advantage_groups} consecutive groups dropped as zero-advantage — " + "the batch isn't filling; check task difficulty (rewards are homogeneous within every group)" + ) + + # Per-group summary. One line per finalized group. rewards = [r.reward for r in survivors] avg_reward = sum(rewards) / len(rewards) if rewards else 0.0 - filter_str = ", ".join(f"{n}={c}" for n, c in filtered_by_name.items()) if filtered_by_name else "—" get_logger().debug( f"Finished group | env={env_name} task_idx={task_idx} | " - f"rollouts={len(group)} (errored={num_errored}, filtered={num_filtered}) | " - f"reward={avg_reward:.4f} | filters: {filter_str}" + f"rollouts={len(group)} (errored={num_errored}, zero_advantage={num_zero_advantage}) | " + f"reward={avg_reward:.4f}" ) def process_batch(self) -> TrainBatch: - """Pop a cohort off ``pending_batch`` (by rollout count when - ``batch_size`` is set, by token count when ``token_batch_size`` is - set), apply post-batch filter annotations, and assemble the + """Pop a ``batch_size`` cohort off ``pending_batch`` and assemble the trainer-bound ``TrainingSample`` list. Overflow stays for the next batch.""" - if self.batch_size is not None: - cohort = self.pending_batch[: self.batch_size] - self.pending_batch = self.pending_batch[self.batch_size :] - else: - assert self.token_batch_size is not None - cut = 0 - running = 0 - for i, r in enumerate(self.pending_batch): - running += payload_tokens(r) - cut = i + 1 - if running >= self.token_batch_size: - break - cohort = self.pending_batch[:cut] - self.pending_batch = self.pending_batch[cut:] - self.pending_tokens -= running - - if self.post_filters: - apply_filters(self.post_filters, cohort) + cohort = self.pending_batch[: self.batch_size] + self.pending_batch = self.pending_batch[self.batch_size :] # Samples are pre-built by ``process_rollout``; ``process_group`` already stamped the - # advantage stream and loss routing on each sample. Filtered rollouts don't ship. + # advantage stream and loss routing on each sample. Zero-advantage rollouts kept in the + # budget by ``count_zero_advantage_in_batch`` don't ship. samples: list[TrainingSample] = [sample for r in cohort if not r.is_filtered for sample in r.samples] # ``rollouts`` is the observation window — every rollout of every group finalized since the @@ -276,8 +245,3 @@ def process_batch(self) -> TrainBatch: if samples: self.pending_rollouts = TrainRollouts() return TrainBatch(rollouts=rollouts, samples=samples) - - def reset_pre_filter_stats(self) -> None: - self.pre_filter_seen = 0 - self.pre_filter_dropped = 0 - self.pre_filter_dropped_by_name.clear() diff --git a/src/prime_rl/orchestrator/types.py b/src/prime_rl/orchestrator/types.py index b77de8ba31..3e8b5681f8 100644 --- a/src/prime_rl/orchestrator/types.py +++ b/src/prime_rl/orchestrator/types.py @@ -96,8 +96,8 @@ class Rollout(vf.Trace[DataT], Generic[DataT]): samples: list[TrainingSample] = Field(default_factory=list, exclude=True) # Per-token rl advantage stream, full-length-N (= len(token_ids)) per # sample, concatenated across the rollout's samples in order; 0.0 on - # non-trainable positions. None = no credit assigned (advantage-based - # filters skip it; the wire ships no advantage stream). + # non-trainable positions. None = no credit assigned (the zero-advantage + # drop skips it; the wire ships no advantage stream). advantages: list[float] | None = Field(default=None, exclude=True) is_filtered: bool = Field(default=False, exclude=True) filter_results: dict[str, bool] = Field(default_factory=dict, exclude=True) diff --git a/src/prime_rl/utils/monitor/prime.py b/src/prime_rl/utils/monitor/prime.py index eee454a9aa..8d57133321 100644 --- a/src/prime_rl/utils/monitor/prime.py +++ b/src/prime_rl/utils/monitor/prime.py @@ -164,8 +164,7 @@ def _register_run(self, config: PrimeMonitorConfig, run_config: OrchestratorConf "max_steps": (run_config.max_steps if run_config else None) or 0, } if run_config: - if run_config.batch_size is not None: - payload["batch_size"] = run_config.batch_size + payload["batch_size"] = run_config.batch_size payload["rollouts_per_example"] = run_config.group_size payload["seq_len"] = run_config.seq_len payload["environments"] = [{"id": env.env_id} for env in run_config.train.source] diff --git a/tests/unit/orchestrator/test_filters.py b/tests/unit/orchestrator/test_filters.py index 69ce76d029..04af5c9fb3 100644 --- a/tests/unit/orchestrator/test_filters.py +++ b/tests/unit/orchestrator/test_filters.py @@ -1,21 +1,21 @@ -import math import uuid import verifiers.v1 as vf -from prime_rl.configs.orchestrator import GibberishFilterConfig, RepetitionFilterConfig from prime_rl.orchestrator.filters import ( - GibberishFilter, - RepetitionFilter, - apply_filters, - setup_filter, - setup_filters, + REPETITION_WINDOW, + detect_gibberish, + detect_repetition, + gibberish_logprob_threshold, + has_zero_advantage, ) from prime_rl.orchestrator.types import Rollout +GIBBERISH_THRESHOLD = gibberish_logprob_threshold(vocab_size=128_000) + def _assistant_node(token_ids: list[int], logprobs: list[float]) -> vf.MessageNode: - """An assistant node whose tokens are all model-sampled (the filters read each node's + """An assistant node whose tokens are all model-sampled (the detectors read each node's masked-True tokens + logprobs).""" return vf.MessageNode( message=vf.AssistantMessage(content="x"), @@ -48,7 +48,7 @@ def _make_rollout( multi_step: bool = False, ) -> Rollout: """Build a ``Rollout`` (a message-graph trace) carrying the completion tokens — enough for - the filters to inspect each node's sampled tokens / logprobs.""" + the detectors to inspect each node's sampled tokens / logprobs.""" if multi_step: mid = len(completion_ids) // 2 nodes = [ @@ -68,69 +68,40 @@ def _make_rollout( return rollout -def _make_gibberish_filter(vocab_size=128_000, token_id_threshold=100_000, logprob_offset=2.0, enforce=False): - logprob_threshold = -math.log(vocab_size) - logprob_offset - return GibberishFilter( - name="gibberish", token_id_threshold=token_id_threshold, logprob_threshold=logprob_threshold, enforce=enforce - ) - - -def _make_repetition_filter(window=5, prob_threshold=0.99, enforce=False): - return RepetitionFilter( - name="repetition", window=window, logprob_threshold=math.log(prob_threshold), enforce=enforce - ) - - -# --- GibberishFilter tests --- +# --- detect_gibberish tests --- def test_gibberish_detects_rare_low_prob_token(): - gibberish_filter = _make_gibberish_filter() - - result = gibberish_filter.check( - _make_rollout( - completion_ids=[50, 120_000, 80], - completion_logprobs=[-1.0, gibberish_filter.logprob_threshold - 1.0, -0.5], - ) + rollout = _make_rollout( + completion_ids=[50, 120_000, 80], + completion_logprobs=[-1.0, GIBBERISH_THRESHOLD - 1.0, -0.5], ) - assert result.detected is True + assert detect_gibberish(rollout, GIBBERISH_THRESHOLD) is True def test_gibberish_ignores_normal_tokens(): - gibberish_filter = _make_gibberish_filter() - - result = gibberish_filter.check( - _make_rollout( - completion_ids=[10, 200, 5000], - completion_logprobs=[-1.0, -2.0, -3.0], - ) + rollout = _make_rollout( + completion_ids=[10, 200, 5000], + completion_logprobs=[-1.0, -2.0, -3.0], ) - assert result.detected is False + assert detect_gibberish(rollout, GIBBERISH_THRESHOLD) is False def test_gibberish_ignores_high_prob_rare_token(): - gibberish_filter = _make_gibberish_filter() - - result = gibberish_filter.check( - _make_rollout( - completion_ids=[120_000], - completion_logprobs=[-0.5], - ) + rollout = _make_rollout( + completion_ids=[120_000], + completion_logprobs=[-0.5], ) - assert result.detected is False + assert detect_gibberish(rollout, GIBBERISH_THRESHOLD) is False def test_gibberish_works_across_trajectory_steps(): - gibberish_filter = _make_gibberish_filter() - - result = gibberish_filter.check( - _make_rollout( - completion_ids=[50, 60, 120_000, 80], - completion_logprobs=[-1.0, -0.5, gibberish_filter.logprob_threshold - 1.0, -0.5], - multi_step=True, - ) + rollout = _make_rollout( + completion_ids=[50, 60, 120_000, 80], + completion_logprobs=[-1.0, -0.5, GIBBERISH_THRESHOLD - 1.0, -0.5], + multi_step=True, ) - assert result.detected is True + assert detect_gibberish(rollout, GIBBERISH_THRESHOLD) is True def test_gibberish_aligns_logprobs_under_generation_prompt_scaffold(): @@ -138,271 +109,58 @@ def test_gibberish_aligns_logprobs_under_generation_prompt_scaffold(): suffix-only logprobs, and the gibberish token is the LAST completion token. The old per-node ``zip(token_ids, logprobs, mask)`` truncated at len(logprobs) and never examined it; reading the aligned branch streams detects it.""" - gibberish_filter = _make_gibberish_filter() - rollout = Rollout[vf.TaskData]( task=vf.TraceTask(type="Task", data=vf.TaskData(idx=0, prompt="")), agent=vf.AgentInfo(config=vf.AgentConfig()), - nodes=[_scaffold_assistant_node([50, 80, 120_000], [-1.0, -0.5, gibberish_filter.logprob_threshold - 1.0])], + nodes=[_scaffold_assistant_node([50, 80, 120_000], [-1.0, -0.5, GIBBERISH_THRESHOLD - 1.0])], rewards={"reward": vf.Reward(score=1.0)}, ) - - result = gibberish_filter.check(rollout) - assert result.detected is True + assert detect_gibberish(rollout, GIBBERISH_THRESHOLD) is True -# --- RepetitionFilter tests --- +# --- detect_repetition tests --- def test_repetition_triggers_after_window(): - repetition_filter = _make_repetition_filter(window=5) - - result = repetition_filter.check( - _make_rollout( - completion_ids=list(range(5)), - completion_logprobs=[-0.001] * 5, - ) - ) - assert result.detected is True - - -def test_repetition_no_trigger_below_window(): - repetition_filter = _make_repetition_filter(window=5) - - result = repetition_filter.check( - _make_rollout( - completion_ids=list(range(4)), - completion_logprobs=[-0.001] * 4, - ) - ) - assert result.detected is False - - -def test_repetition_resets_on_low_prob(): - repetition_filter = _make_repetition_filter(window=5) - - logprobs = [-0.001] * 3 + [-2.0] + [-0.001] * 3 - result = repetition_filter.check( - _make_rollout( - completion_ids=list(range(7)), - completion_logprobs=logprobs, - ) - ) - assert result.detected is False - - -def test_repetition_varied_probs_no_trigger(): - repetition_filter = _make_repetition_filter(window=3) - - result = repetition_filter.check( - _make_rollout( - completion_ids=list(range(6)), - completion_logprobs=[-0.001, -3.0, -0.001, -3.0, -0.001, -3.0], - ) - ) - assert result.detected is False - - -# --- setup_filter / setup_filters tests --- - - -def test_setup_filter_gibberish(): - config = GibberishFilterConfig(token_id_threshold=100_000, logprob_offset=2.0) - gibberish_filter = setup_filter(config, vocab_size=128_000) - assert isinstance(gibberish_filter, GibberishFilter) - assert gibberish_filter.name == "gibberish" - assert gibberish_filter.token_id_threshold == 100_000 - assert abs(gibberish_filter.logprob_threshold - (-math.log(128_000) - 2.0)) < 1e-10 - assert gibberish_filter.enforce is False - - -def test_setup_filter_gibberish_enforce(): - config = GibberishFilterConfig(enforce=True) - gibberish_filter = setup_filter(config, vocab_size=128_000) - assert gibberish_filter.enforce is True - - -def test_setup_filter_repetition(): - config = RepetitionFilterConfig(window=3_000, prob_threshold=0.99) - repetition_filter = setup_filter(config, vocab_size=128_000) - assert isinstance(repetition_filter, RepetitionFilter) - assert repetition_filter.name == "repetition" - assert repetition_filter.window == 3_000 - assert abs(repetition_filter.logprob_threshold - math.log(0.99)) < 1e-10 - assert repetition_filter.enforce is False - - -def test_setup_filter_repetition_enforce(): - config = RepetitionFilterConfig(enforce=True) - repetition_filter = setup_filter(config, vocab_size=128_000) - assert repetition_filter.enforce is True - - -def test_setup_filters_multiple(): - configs = [ - GibberishFilterConfig(), - RepetitionFilterConfig(), - ] - filters = setup_filters(configs, vocab_size=128_000, kind="post-batch") - assert len(filters) == 2 - assert filters[0].name == "gibberish" - assert filters[1].name == "repetition" - - -# --- apply_filters tests (enforce=True) --- - - -def test_apply_filters_enforced_flags_rollout(): - gibberish_filter = _make_gibberish_filter(enforce=True) - - rollout = _make_rollout( - completion_ids=[120_000], - completion_logprobs=[gibberish_filter.logprob_threshold - 1.0], - reward=1.0, - ) - - apply_filters([gibberish_filter], [rollout]) - - assert rollout.reward == 1.0 - assert rollout.nodes[0].token_ids == [120_000] - assert rollout.nodes[0].mask == [True] - assert rollout.stop_condition is None - assert rollout.filter_results == {"gibberish": True} - assert rollout.is_filtered is True - - -def test_apply_filters_preserves_clean_rollouts(): - gibberish_filter = _make_gibberish_filter(enforce=True) - - rollout = _make_rollout( - completion_ids=[50, 60, 70], - completion_logprobs=[-1.0, -2.0, -1.5], - reward=1.0, - ) - - apply_filters([gibberish_filter], [rollout]) - - assert rollout.reward == 1.0 - assert rollout.nodes[0].token_ids == [50, 60, 70] - assert all(rollout.nodes[0].mask) - assert rollout.stop_condition is None - assert rollout.filter_results == {"gibberish": False} - assert rollout.is_filtered is False - - -def test_apply_filters_first_filter_wins(): - gibberish_filter = _make_gibberish_filter(enforce=True) - repetition_filter = _make_repetition_filter(window=2, enforce=True) - - rollout = _make_rollout( - completion_ids=[120_000, 1, 2], - completion_logprobs=[gibberish_filter.logprob_threshold - 1.0, -0.001, -0.001], - reward=1.0, - ) - - apply_filters([gibberish_filter, repetition_filter], [rollout]) - - assert rollout.stop_condition is None - assert rollout.filter_results == {"gibberish": True, "repetition": False} - assert rollout.is_filtered is True - - -def test_apply_filters_empty_list(): rollout = _make_rollout( - completion_ids=[1, 2, 3], - completion_logprobs=[-1.0, -1.0, -1.0], + completion_ids=list(range(REPETITION_WINDOW)), + completion_logprobs=[-0.001] * REPETITION_WINDOW, ) - apply_filters([], [rollout]) - assert rollout.filter_results == {} - assert rollout.is_filtered is False - assert rollout.reward == 1.0 + assert detect_repetition(rollout) is True -def test_apply_filters_mixed_batch(): - gibberish_filter = _make_gibberish_filter(enforce=True) - - clean = _make_rollout(completion_ids=[50], completion_logprobs=[-1.0], reward=1.0) - dirty = _make_rollout( - completion_ids=[120_000], completion_logprobs=[gibberish_filter.logprob_threshold - 1.0], reward=1.0 - ) - - apply_filters([gibberish_filter], [clean, dirty]) - - assert clean.reward == 1.0 - assert dirty.reward == 1.0 - assert clean.is_filtered is False - assert dirty.is_filtered is True - - -def test_apply_filters_enforced_preserves_rollout_tokens(): - gibberish_filter = _make_gibberish_filter(enforce=True) - +def test_repetition_no_trigger_below_window(): rollout = _make_rollout( - completion_ids=[10, 120_000, 30], - completion_logprobs=[-1.0, gibberish_filter.logprob_threshold - 1.0, -0.5], - reward=1.0, + completion_ids=list(range(REPETITION_WINDOW - 1)), + completion_logprobs=[-0.001] * (REPETITION_WINDOW - 1), ) + assert detect_repetition(rollout) is False - apply_filters([gibberish_filter], [rollout]) - - assert rollout.nodes[0].token_ids == [10, 120_000, 30] - assert rollout.nodes[0].logprobs == [ - -1.0, - gibberish_filter.logprob_threshold - 1.0, - -0.5, - ] - assert rollout.nodes[0].mask == [True, True, True] - assert rollout.is_filtered is True - - -def test_apply_filters_preserves_existing_stop_condition(): - gibberish_filter = _make_gibberish_filter(enforce=True) +def test_repetition_resets_on_low_prob(): + logprobs = [-0.001] * (REPETITION_WINDOW - 1) + [-2.0] + [-0.001] * (REPETITION_WINDOW - 1) rollout = _make_rollout( - completion_ids=[120_000], - completion_logprobs=[gibberish_filter.logprob_threshold - 1.0], - reward=1.0, + completion_ids=list(range(len(logprobs))), + completion_logprobs=logprobs, ) - rollout.stop_condition = "generation_truncated" - - apply_filters([gibberish_filter], [rollout]) - - assert rollout.stop_condition == "generation_truncated" - assert rollout.is_filtered is True + assert detect_repetition(rollout) is False -# --- apply_filters tests (monitor-only, enforce=False) --- +# --- has_zero_advantage tests --- -def test_apply_filters_monitor_only_tracks_detection(): - gibberish_filter = _make_gibberish_filter(enforce=False) +def test_zero_advantage_without_advantages(): + rollout = _make_rollout(completion_ids=[1, 2], completion_logprobs=[-1.0, -1.0]) + assert has_zero_advantage(rollout) is False - rollout = _make_rollout( - completion_ids=[120_000], - completion_logprobs=[gibberish_filter.logprob_threshold - 1.0], - reward=1.0, - ) - apply_filters([gibberish_filter], [rollout]) - - assert rollout.reward == 1.0 - assert all(rollout.nodes[0].mask) - assert rollout.stop_condition is None - assert rollout.filter_results == {"gibberish": True} - assert rollout.is_filtered is False - - -def test_apply_filters_monitor_only_mixed_batch(): - gibberish_filter = _make_gibberish_filter(enforce=False) - - clean = _make_rollout(completion_ids=[50], completion_logprobs=[-1.0], reward=1.0) - dirty = _make_rollout( - completion_ids=[120_000], completion_logprobs=[gibberish_filter.logprob_threshold - 1.0], reward=1.0 - ) +def test_zero_advantage_all_zero(): + rollout = _make_rollout(completion_ids=[1, 2], completion_logprobs=[-1.0, -1.0]) + rollout.advantages = [0.0, 0.0] + assert has_zero_advantage(rollout) is True - apply_filters([gibberish_filter], [clean, dirty]) - assert clean.reward == 1.0 - assert dirty.reward == 1.0 - assert clean.is_filtered is False - assert dirty.is_filtered is False +def test_zero_advantage_nonzero(): + rollout = _make_rollout(completion_ids=[1, 2], completion_logprobs=[-1.0, -1.0]) + rollout.advantages = [0.5, 0.0] + assert has_zero_advantage(rollout) is False diff --git a/tests/unit/orchestrator/test_metrics.py b/tests/unit/orchestrator/test_metrics.py index 546e594d9a..7afe4d3483 100644 --- a/tests/unit/orchestrator/test_metrics.py +++ b/tests/unit/orchestrator/test_metrics.py @@ -240,10 +240,10 @@ def test_train_only_metrics_absent_from_eval(): out = train_wandb(rollouts) assert out["train/agg/all/agent/is_trainable/mean"] == 0.5 assert out["train/agg/all/agent/is_filtered/mean"] == 0.5 - assert out["train/agg/all/agent/filters/gibberish/mean"] == 0.5 + assert out["train/agg/all/agent/gibberish/mean"] == 0.5 assert "train/agg/all/is_trainable/mean" not in out # pipeline verdicts are per-trace eval_out = EvalRollouts(rollouts).metrics.to_wandb(prefix="eval/x", subset="all") - assert not any("is_trainable" in k or "is_filtered" in k or "/filters/" in k for k in eval_out) + assert not any("is_trainable" in k or "is_filtered" in k or "gibberish" in k for k in eval_out) def test_eval_avg_at_k_and_pass_k():