From 817a5f71e839169fb740073b83367d55ffd0eff7 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Wed, 12 Aug 2026 18:59:44 +0000 Subject: [PATCH 1/4] Hardcode rollout filters and switch to token batching MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace the composable pre/post-batch filter config with hardcoded checks: gibberish and repetition detection are monitor-only metrics, and zero-advantage rollouts drop before they consume batch budget. Algorithms that train without credit (echo) declare trains_on_zero_advantage to keep their rollouts; count_zero_advantage_in_batch opts back into the fixed sampling budget of the old default. Remove rollout-based batching (batch_size, oversampling_factor): a rollout is an arbitrary unit of learning signal — the same batch size means very different step sizes across envs. token_batch_size is now the only batching mode. Checked-in configs are converted as batch_size x measured average tokens per rollout (from recent nightly/ablation W&B runs; estimated where no runs exist). Co-Authored-By: Claude Fable 5 --- README.md | 2 +- configs/basic/alphabet-sort/rl.toml | 2 +- configs/basic/hendrycks-sanity/rl.toml | 2 +- configs/basic/reverse-text/rl.toml | 2 +- configs/basic/wiki-search/rl.toml | 2 +- configs/basic/wordle/rl.toml | 2 +- configs/ci/integration/alphabet_sort.toml | 2 +- .../integration/reverse-text-lora/resume.toml | 2 +- .../integration/reverse-text-lora/start.toml | 2 +- .../integration/reverse-text-moe/start.toml | 2 +- .../reverse-text-rl-opd/start.toml | 2 +- .../reverse-text-rl-sft/start.toml | 2 +- .../ci/integration/reverse-text/resume.toml | 2 +- .../ci/integration/reverse-text/start.toml | 2 +- configs/ci/nightly-fft/alphabet-sort.toml | 3 +- configs/ci/nightly-fft/hendrycks-sanity.toml | 3 +- configs/ci/nightly-fft/reverse-text.toml | 2 +- configs/ci/nightly-fft/wiki-search.toml | 8 +- configs/ci/nightly-fft/wordle.toml | 3 +- .../ci/nightly/multimodal_color_codeword.toml | 3 +- configs/debug/algo/echo.toml | 9 +- configs/debug/algo/grpo.toml | 2 +- configs/debug/algo/max_rl.toml | 2 +- configs/debug/algo/mixed_grpo_opd.toml | 2 +- configs/debug/algo/opd.toml | 2 +- configs/debug/algo/opd_lora.toml | 2 +- configs/debug/algo/rae.toml | 3 +- configs/debug/algo/self_distill.toml | 3 +- configs/debug/algo/sft_distill.toml | 2 +- configs/debug/algo/sft_distill_lora.toml | 2 +- configs/debug/multi-env/rl.toml | 2 +- docs/algorithms.md | 41 +- docs/development.md | 2 +- docs/overview.md | 2 +- docs/training.md | 3 +- examples/advanced/glm-4.5-air/search.toml | 2 +- examples/advanced/glm-4.5-air/swe.toml | 2 +- examples/advanced/glm-4.5-air/terminal.toml | 2 +- examples/advanced/glm-5.2/swe-llmd.toml | 2 +- examples/advanced/glm-5.2/swe.toml | 8 +- examples/advanced/intellect-3.1/rl.toml | 8 +- examples/advanced/minimax-m2.5/swe.toml | 4 +- examples/advanced/nemotron-3-super/swe.toml | 2 +- examples/advanced/qwen3-30b-a3b/math.toml | 4 +- examples/advanced/qwen3-30b-a3b/swe.toml | 4 +- examples/advanced/qwen3-30b-a3b/tool.toml | 3 +- examples/basic/alphabet-sort/rl.toml | 3 +- examples/basic/hendrycks-sanity/rl.toml | 3 +- examples/basic/reverse-text/rl.toml | 2 +- examples/basic/wiki-search/rl.toml | 8 +- examples/basic/wordle/rl.toml | 3 +- k8s/prime-rl/examples/reverse-text/orch.toml | 2 +- .../src/prime_rl/configs/orchestrator.py | 119 +----- .../src/prime_rl/configs/rl.py | 2 +- src/prime_rl/orchestrator/algo/base.py | 31 +- src/prime_rl/orchestrator/algo/echo.py | 4 + src/prime_rl/orchestrator/algo/max_rl.py | 4 +- src/prime_rl/orchestrator/algo/opd.py | 4 +- src/prime_rl/orchestrator/algo/sft.py | 3 +- src/prime_rl/orchestrator/filters.py | 223 +++-------- src/prime_rl/orchestrator/orchestrator.py | 38 +- src/prime_rl/orchestrator/train_sink.py | 171 ++++----- src/prime_rl/orchestrator/types.py | 4 +- src/prime_rl/utils/monitor/prime.py | 2 - tests/unit/orchestrator/test_filters.py | 350 +++--------------- tests/unit/test_configs.py | 2 +- 66 files changed, 330 insertions(+), 818 deletions(-) diff --git a/README.md b/README.md index b479197149..0f843f384f 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, the AIPO loss, 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/basic/alphabet-sort/rl.toml b/configs/basic/alphabet-sort/rl.toml index bc70f7718d..aced4a54b6 100644 --- a/configs/basic/alphabet-sort/rl.toml +++ b/configs/basic/alphabet-sort/rl.toml @@ -26,7 +26,7 @@ alpha = 64 lr = 1e-5 [orchestrator] -batch_size = 128 +token_batch_size = 131072 # 128 rollouts x ~1k avg tokens group_size = 8 [orchestrator.train.sampling] diff --git a/configs/basic/hendrycks-sanity/rl.toml b/configs/basic/hendrycks-sanity/rl.toml index b16c7c7c4c..2ac600b6e9 100644 --- a/configs/basic/hendrycks-sanity/rl.toml +++ b/configs/basic/hendrycks-sanity/rl.toml @@ -15,7 +15,7 @@ project = "hendrycks-sanity" name = "hendrycks-sanity" [orchestrator] -batch_size = 128 +token_batch_size = 524288 # 128 rollouts x ~4k avg tokens group_size = 8 seq_len = 8192 diff --git a/configs/basic/reverse-text/rl.toml b/configs/basic/reverse-text/rl.toml index 3344560e99..82de868efc 100644 --- a/configs/basic/reverse-text/rl.toml +++ b/configs/basic/reverse-text/rl.toml @@ -16,7 +16,7 @@ project = "reverse-text" name = "reverse-text" [orchestrator] -batch_size = 128 +token_batch_size = 16384 # 128 rollouts x ~128 avg tokens group_size = 16 [orchestrator.train.sampling] diff --git a/configs/basic/wiki-search/rl.toml b/configs/basic/wiki-search/rl.toml index aa509e9049..9af64f9fe3 100644 --- a/configs/basic/wiki-search/rl.toml +++ b/configs/basic/wiki-search/rl.toml @@ -34,7 +34,7 @@ target_modules = [ ] [orchestrator] -batch_size = 128 +token_batch_size = 196608 # 128 rollouts x ~1.5k avg tokens group_size = 16 [orchestrator.train.sampling] diff --git a/configs/basic/wordle/rl.toml b/configs/basic/wordle/rl.toml index fff93ac571..0594dfb34e 100644 --- a/configs/basic/wordle/rl.toml +++ b/configs/basic/wordle/rl.toml @@ -18,7 +18,7 @@ name = "wordle" name = "PrimeIntellect/Qwen3-1.7B-Wordle-SFT" [orchestrator] -batch_size = 128 +token_batch_size = 393216 # 128 rollouts x ~3k avg tokens group_size = 8 [[orchestrator.train.source]] diff --git a/configs/ci/integration/alphabet_sort.toml b/configs/ci/integration/alphabet_sort.toml index 51472260e4..d38fe353ad 100644 --- a/configs/ci/integration/alphabet_sort.toml +++ b/configs/ci/integration/alphabet_sort.toml @@ -16,7 +16,7 @@ name = "Qwen/Qwen3-0.6B" lr = 1e-5 [orchestrator] -batch_size = 128 +token_batch_size = 131072 # 128 rollouts x ~1k avg tokens group_size = 8 [orchestrator.train.sampling] diff --git a/configs/ci/integration/reverse-text-lora/resume.toml b/configs/ci/integration/reverse-text-lora/resume.toml index ea6437df67..a02839f412 100644 --- a/configs/ci/integration/reverse-text-lora/resume.toml +++ b/configs/ci/integration/reverse-text-lora/resume.toml @@ -17,7 +17,7 @@ rank = 8 save_adapter_separately = true [orchestrator] -batch_size = 128 +token_batch_size = 16384 # 128 rollouts x ~128 avg tokens group_size = 16 [orchestrator.model.lora] diff --git a/configs/ci/integration/reverse-text-lora/start.toml b/configs/ci/integration/reverse-text-lora/start.toml index 74789ca4bc..ec27af4203 100644 --- a/configs/ci/integration/reverse-text-lora/start.toml +++ b/configs/ci/integration/reverse-text-lora/start.toml @@ -16,7 +16,7 @@ rank = 8 save_adapter_separately = true [orchestrator] -batch_size = 128 +token_batch_size = 16384 # 128 rollouts x ~128 avg tokens group_size = 16 [orchestrator.model.lora] diff --git a/configs/ci/integration/reverse-text-moe/start.toml b/configs/ci/integration/reverse-text-moe/start.toml index 8748737f55..19c9e6a280 100644 --- a/configs/ci/integration/reverse-text-moe/start.toml +++ b/configs/ci/integration/reverse-text-moe/start.toml @@ -19,7 +19,7 @@ lr = 3e-6 impl = "custom" [orchestrator] -batch_size = 128 +token_batch_size = 16384 # 128 rollouts x ~128 avg tokens group_size = 16 [orchestrator.train.sampling] diff --git a/configs/ci/integration/reverse-text-rl-opd/start.toml b/configs/ci/integration/reverse-text-rl-opd/start.toml index 80fdc80bba..2bd94e2457 100644 --- a/configs/ci/integration/reverse-text-rl-opd/start.toml +++ b/configs/ci/integration/reverse-text-rl-opd/start.toml @@ -21,7 +21,7 @@ project = "reverse-text-ci" name = "ci-rl-opd" [orchestrator] -batch_size = 128 +token_batch_size = 16384 # 128 rollouts x ~128 avg tokens group_size = 16 [orchestrator.algo] diff --git a/configs/ci/integration/reverse-text-rl-sft/start.toml b/configs/ci/integration/reverse-text-rl-sft/start.toml index dfd829b785..02668d1b02 100644 --- a/configs/ci/integration/reverse-text-rl-sft/start.toml +++ b/configs/ci/integration/reverse-text-rl-sft/start.toml @@ -21,7 +21,7 @@ project = "reverse-text-ci" name = "ci-rl-sft" [orchestrator] -batch_size = 128 +token_batch_size = 16384 # 128 rollouts x ~128 avg tokens group_size = 16 [orchestrator.algo] diff --git a/configs/ci/integration/reverse-text/resume.toml b/configs/ci/integration/reverse-text/resume.toml index 9b34f9835c..8f2aa964ab 100644 --- a/configs/ci/integration/reverse-text/resume.toml +++ b/configs/ci/integration/reverse-text/resume.toml @@ -17,7 +17,7 @@ name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT" lr = 3e-6 [orchestrator] -batch_size = 128 +token_batch_size = 16384 # 128 rollouts x ~128 avg tokens group_size = 16 [orchestrator.train.sampling] diff --git a/configs/ci/integration/reverse-text/start.toml b/configs/ci/integration/reverse-text/start.toml index 4c38c2e66e..14a2208e61 100644 --- a/configs/ci/integration/reverse-text/start.toml +++ b/configs/ci/integration/reverse-text/start.toml @@ -16,7 +16,7 @@ name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT" lr = 3e-6 [orchestrator] -batch_size = 128 +token_batch_size = 16384 # 128 rollouts x ~128 avg tokens group_size = 16 [orchestrator.train.sampling] diff --git a/configs/ci/nightly-fft/alphabet-sort.toml b/configs/ci/nightly-fft/alphabet-sort.toml index 5ccdf93587..a04983b726 100644 --- a/configs/ci/nightly-fft/alphabet-sort.toml +++ b/configs/ci/nightly-fft/alphabet-sort.toml @@ -15,7 +15,8 @@ name = "alphabet-sort" name = "Qwen/Qwen3-4B-Instruct-2507" [orchestrator] -batch_size = 512 +token_batch_size = 524288 # 512 rollouts x ~1k avg tokens +max_inflight_episodes = 512 group_size = 16 [[orchestrator.train.source]] diff --git a/configs/ci/nightly-fft/hendrycks-sanity.toml b/configs/ci/nightly-fft/hendrycks-sanity.toml index 8eb547efee..2b8bba3734 100644 --- a/configs/ci/nightly-fft/hendrycks-sanity.toml +++ b/configs/ci/nightly-fft/hendrycks-sanity.toml @@ -14,7 +14,8 @@ name = "hendrycks-sanity" name = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B" [orchestrator] -batch_size = 512 +token_batch_size = 2097152 # 512 rollouts x ~4k avg tokens +max_inflight_episodes = 512 group_size = 8 seq_len = 8192 diff --git a/configs/ci/nightly-fft/reverse-text.toml b/configs/ci/nightly-fft/reverse-text.toml index 50ddec8eff..106c663098 100644 --- a/configs/ci/nightly-fft/reverse-text.toml +++ b/configs/ci/nightly-fft/reverse-text.toml @@ -15,7 +15,7 @@ name = "reverse-text" name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT" [orchestrator] -batch_size = 128 +token_batch_size = 16384 # 128 rollouts x ~128 avg tokens group_size = 16 [[orchestrator.train.source]] diff --git a/configs/ci/nightly-fft/wiki-search.toml b/configs/ci/nightly-fft/wiki-search.toml index c3977c0adb..03f256324d 100644 --- a/configs/ci/nightly-fft/wiki-search.toml +++ b/configs/ci/nightly-fft/wiki-search.toml @@ -15,13 +15,9 @@ name = "wiki-search" name = "Qwen/Qwen3-4B-Instruct-2507" [orchestrator] -batch_size = 512 +token_batch_size = 786432 # 512 rollouts x ~1.5k avg tokens +max_inflight_episodes = 1024 group_size = 16 -oversampling_factor = 2.0 - -[[orchestrator.pre_batch_filters]] -type = "zero_advantage" -enforce = true [[orchestrator.train.source]] name = "wiki-search" diff --git a/configs/ci/nightly-fft/wordle.toml b/configs/ci/nightly-fft/wordle.toml index d31533836e..7d159b1916 100644 --- a/configs/ci/nightly-fft/wordle.toml +++ b/configs/ci/nightly-fft/wordle.toml @@ -15,7 +15,8 @@ name = "wordle" name = "PrimeIntellect/Qwen3-1.7B-Wordle-SFT" [orchestrator] -batch_size = 512 +token_batch_size = 1572864 # 512 rollouts x ~3k avg tokens +max_inflight_episodes = 512 group_size = 16 [[orchestrator.train.source]] diff --git a/configs/ci/nightly/multimodal_color_codeword.toml b/configs/ci/nightly/multimodal_color_codeword.toml index f4b97dd4b3..a870499d11 100644 --- a/configs/ci/nightly/multimodal_color_codeword.toml +++ b/configs/ci/nightly/multimodal_color_codeword.toml @@ -14,7 +14,8 @@ vision_encoder_attr = "model.visual" language_model_attr = "model.language_model" [orchestrator] -batch_size = 256 +token_batch_size = 262144 # 256 rollouts x ~1k avg tokens +max_inflight_episodes = 256 group_size = 16 [orchestrator.train.sampling] diff --git a/configs/debug/algo/echo.toml b/configs/debug/algo/echo.toml index e9bdddccb3..2080aa8fb5 100644 --- a/configs/debug/algo/echo.toml +++ b/configs/debug/algo/echo.toml @@ -13,7 +13,8 @@ project = "algorithms-debug" name = "debug-echo" [orchestrator] -batch_size = 32 +token_batch_size = 32768 # 32 rollouts x ~1k avg tokens +max_inflight_episodes = 32 group_size = 4 # alphabet-sort's feedback arrives as user messages, so train the user role @@ -44,12 +45,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/configs/debug/algo/grpo.toml b/configs/debug/algo/grpo.toml index 30dd90b660..f3d2701480 100644 --- a/configs/debug/algo/grpo.toml +++ b/configs/debug/algo/grpo.toml @@ -9,7 +9,7 @@ project = "algorithms-debug" name = "debug-rl" [orchestrator] -batch_size = 128 +token_batch_size = 16384 # 128 rollouts x ~128 avg tokens group_size = 16 [orchestrator.algo] diff --git a/configs/debug/algo/max_rl.toml b/configs/debug/algo/max_rl.toml index 15c5e83f19..c6a0b9c057 100644 --- a/configs/debug/algo/max_rl.toml +++ b/configs/debug/algo/max_rl.toml @@ -9,7 +9,7 @@ project = "algorithms-debug" name = "debug-max-rl" [orchestrator] -batch_size = 128 +token_batch_size = 16384 # 128 rollouts x ~128 avg tokens group_size = 16 [orchestrator.algo] diff --git a/configs/debug/algo/mixed_grpo_opd.toml b/configs/debug/algo/mixed_grpo_opd.toml index c56b37ad11..612f332821 100644 --- a/configs/debug/algo/mixed_grpo_opd.toml +++ b/configs/debug/algo/mixed_grpo_opd.toml @@ -19,7 +19,7 @@ project = "algorithms-debug" name = "debug-mixed-grpo-opd" [orchestrator] -batch_size = 128 +token_batch_size = 16384 # 128 rollouts x ~128 avg tokens group_size = 16 [orchestrator.algo] diff --git a/configs/debug/algo/opd.toml b/configs/debug/algo/opd.toml index 7b33bac34e..0fcebe1ef0 100644 --- a/configs/debug/algo/opd.toml +++ b/configs/debug/algo/opd.toml @@ -16,7 +16,7 @@ project = "algorithms-debug" name = "debug-opd" [orchestrator] -batch_size = 128 +token_batch_size = 16384 # 128 rollouts x ~128 avg tokens group_size = 16 [orchestrator.algo] diff --git a/configs/debug/algo/opd_lora.toml b/configs/debug/algo/opd_lora.toml index de66ace6d0..ec99c0701a 100644 --- a/configs/debug/algo/opd_lora.toml +++ b/configs/debug/algo/opd_lora.toml @@ -16,7 +16,7 @@ project = "algorithms-debug" name = "debug-opd-lora" [orchestrator] -batch_size = 128 +token_batch_size = 16384 # 128 rollouts x ~128 avg tokens group_size = 16 [orchestrator.algo] diff --git a/configs/debug/algo/rae.toml b/configs/debug/algo/rae.toml index 65b04cf9d1..3c27659381 100644 --- a/configs/debug/algo/rae.toml +++ b/configs/debug/algo/rae.toml @@ -9,7 +9,8 @@ project = "algorithms-debug" name = "debug-rae" [orchestrator] -batch_size = 32 +token_batch_size = 16384 # 32 rollouts x ~512 avg tokens +max_inflight_episodes = 32 group_size = 1 [orchestrator.algo] diff --git a/configs/debug/algo/self_distill.toml b/configs/debug/algo/self_distill.toml index 0b4dfbe869..1378e7d512 100644 --- a/configs/debug/algo/self_distill.toml +++ b/configs/debug/algo/self_distill.toml @@ -15,7 +15,8 @@ project = "algorithms-debug" name = "debug-self-distill" [orchestrator] -batch_size = 32 +token_batch_size = 4096 # 32 rollouts x ~128 avg tokens +max_inflight_episodes = 32 group_size = 1 # reverse-text's demo lives in the "answer" column. diff --git a/configs/debug/algo/sft_distill.toml b/configs/debug/algo/sft_distill.toml index cf19e32f37..dcb5e0eb94 100644 --- a/configs/debug/algo/sft_distill.toml +++ b/configs/debug/algo/sft_distill.toml @@ -16,7 +16,7 @@ project = "algorithms-debug" name = "debug-sft" [orchestrator] -batch_size = 128 +token_batch_size = 16384 # 128 rollouts x ~128 avg tokens group_size = 4 [orchestrator.algo] diff --git a/configs/debug/algo/sft_distill_lora.toml b/configs/debug/algo/sft_distill_lora.toml index deb639c3d8..ee9075bd0d 100644 --- a/configs/debug/algo/sft_distill_lora.toml +++ b/configs/debug/algo/sft_distill_lora.toml @@ -16,7 +16,7 @@ project = "algorithms-debug" name = "debug-sft-lora" [orchestrator] -batch_size = 128 +token_batch_size = 16384 # 128 rollouts x ~128 avg tokens group_size = 4 [orchestrator.algo] diff --git a/configs/debug/multi-env/rl.toml b/configs/debug/multi-env/rl.toml index ee77bbb1c6..9debdf5d8f 100644 --- a/configs/debug/multi-env/rl.toml +++ b/configs/debug/multi-env/rl.toml @@ -8,7 +8,7 @@ seq_len = 2048 name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT" [orchestrator] -batch_size = 128 +token_batch_size = 16384 # 128 rollouts x ~128 avg tokens group_size = 16 [orchestrator.renderer] diff --git a/docs/algorithms.md b/docs/algorithms.md index e3012fa50f..3de772d3ea 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 (`filters/gibberish`), never dropped. | +| `repetition` | Detects long high-confidence loops. Monitor-only: tracked in metrics (`filters/repetition`), 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 `token_batch_size` — a fixed sampling budget per step, at the cost of a variable number of trained-on tokens. -```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`, `filters/zero_advantage`), just not in the trainer batch. ## Multi-Turn Trajectories diff --git a/docs/development.md b/docs/development.md index 32b631500d..0e25f813c7 100644 --- a/docs/development.md +++ b/docs/development.md @@ -137,7 +137,7 @@ Before merging a new model, you need to ensure the following: - The model is correctly registered and defines and all the required methods - such as `convert_hf_layer_to_tt` and `convert_tt_layer_to_hf`. - The small smoke test passes. -In the PR that adds the new model, you also need to provide a table covering the KL mismatch across 20 steps on `math` environment with `batch_size=64`. All the entries in the table must lower than 0.015. If this is not met, the PR will not be merged (unless reasonable justification is provided). This is to ensure all our models are consistent and their implementations match the implementations in the inference framework. +In the PR that adds the new model, you also need to provide a table covering the KL mismatch across 20 steps on `math` environment with `token_batch_size=524288` (~64 rollouts at ~8k avg tokens). All the entries in the table must lower than 0.015. If this is not met, the PR will not be merged (unless reasonable justification is provided). This is to ensure all our models are consistent and their implementations match the implementations in the inference framework. ## Adding a Custom VLM Implementation 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..911585e472 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.token_batch_size` | Tokens to train on per step. Size it as (target rollouts per step) x (average tokens per rollout). | +| `orchestrator.max_inflight_episodes` | Concurrent episodes kept in-flight. Tune together with `token_batch_size`: roughly `token_batch_size / (average tokens per rollout)`, higher 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-4.5-air/search.toml b/examples/advanced/glm-4.5-air/search.toml index 054799f928..a1a9678b16 100644 --- a/examples/advanced/glm-4.5-air/search.toml +++ b/examples/advanced/glm-4.5-air/search.toml @@ -68,7 +68,7 @@ type = "muon" # --- Orchestrator --- [orchestrator] -batch_size = 256 +token_batch_size = 4194304 # 256 rollouts x ~16k avg tokens group_size = 16 max_off_policy_steps = 32 max_inflight_episodes = 512 diff --git a/examples/advanced/glm-4.5-air/swe.toml b/examples/advanced/glm-4.5-air/swe.toml index 696bec5442..197a13f575 100644 --- a/examples/advanced/glm-4.5-air/swe.toml +++ b/examples/advanced/glm-4.5-air/swe.toml @@ -79,7 +79,7 @@ type = "muon" # --- Orchestrator --- [orchestrator] -batch_size = 256 +token_batch_size = 5242880 # 256 rollouts x ~20k avg tokens group_size = 16 max_off_policy_steps = 32 max_inflight_episodes = 512 diff --git a/examples/advanced/glm-4.5-air/terminal.toml b/examples/advanced/glm-4.5-air/terminal.toml index c3bf09adbb..f130fabb36 100644 --- a/examples/advanced/glm-4.5-air/terminal.toml +++ b/examples/advanced/glm-4.5-air/terminal.toml @@ -69,7 +69,7 @@ type = "muon" # --- Orchestrator --- [orchestrator] -batch_size = 256 +token_batch_size = 4194304 # 256 rollouts x ~16k avg tokens group_size = 16 max_off_policy_steps = 32 max_inflight_episodes = 512 diff --git a/examples/advanced/glm-5.2/swe-llmd.toml b/examples/advanced/glm-5.2/swe-llmd.toml index 34edcb0989..4daebab432 100644 --- a/examples/advanced/glm-5.2/swe-llmd.toml +++ b/examples/advanced/glm-5.2/swe-llmd.toml @@ -109,7 +109,7 @@ lr = 1e-6 weight_decay = 0.0 [orchestrator] -batch_size = 256 +token_batch_size = 5242880 # 256 rollouts x ~20k avg tokens group_size = 16 max_inflight_episodes = 2048 max_off_policy_steps = 8 # 16/32 can be better, 8 is more CPU stable diff --git a/examples/advanced/glm-5.2/swe.toml b/examples/advanced/glm-5.2/swe.toml index 657a488c9b..dc10d26ce2 100644 --- a/examples/advanced/glm-5.2/swe.toml +++ b/examples/advanced/glm-5.2/swe.toml @@ -64,9 +64,9 @@ lr = 1e-6 weight_decay = 0.1 [orchestrator] -batch_size = 4096 +token_batch_size = 83886080 # 4096 rollouts x ~20k avg tokens +max_inflight_episodes = 12288 group_size = 16 -oversampling_factor = 3 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..93d08b9a88 100644 --- a/examples/advanced/intellect-3.1/rl.toml +++ b/examples/advanced/intellect-3.1/rl.toml @@ -44,8 +44,8 @@ lr = 1e-6 weight_decay = 0.01 [orchestrator] -batch_size = 2048 -oversampling_factor = 2 +token_batch_size = 33554432 # 2048 rollouts x ~16k avg tokens +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..652c7e4afe 100644 --- a/examples/advanced/minimax-m2.5/swe.toml +++ b/examples/advanced/minimax-m2.5/swe.toml @@ -44,8 +44,8 @@ lr = 1e-6 weight_decay = 0.01 [orchestrator] -batch_size = 2048 -oversampling_factor = 2 +token_batch_size = 41943040 # 2048 rollouts x ~20k avg tokens +max_inflight_episodes = 4096 max_off_policy_steps = 16 [[orchestrator.train.source]] diff --git a/examples/advanced/nemotron-3-super/swe.toml b/examples/advanced/nemotron-3-super/swe.toml index e7149b23cc..839458f329 100644 --- a/examples/advanced/nemotron-3-super/swe.toml +++ b/examples/advanced/nemotron-3-super/swe.toml @@ -63,7 +63,7 @@ skip_optimizer = true type = "adamw" [orchestrator] -batch_size = 256 +token_batch_size = 5242880 # 256 rollouts x ~20k avg tokens group_size = 16 max_off_policy_steps = 32 max_inflight_episodes = 1536 diff --git a/examples/advanced/qwen3-30b-a3b/math.toml b/examples/advanced/qwen3-30b-a3b/math.toml index ff2ac4512a..e236e91de1 100644 --- a/examples/advanced/qwen3-30b-a3b/math.toml +++ b/examples/advanced/qwen3-30b-a3b/math.toml @@ -41,8 +41,8 @@ type = "adamw" lr = 1e-6 [orchestrator] -batch_size = 512 -oversampling_factor = 2 +token_batch_size = 4194304 # 512 rollouts x ~8k avg tokens +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..39fdb17f3b 100644 --- a/examples/advanced/qwen3-30b-a3b/swe.toml +++ b/examples/advanced/qwen3-30b-a3b/swe.toml @@ -42,8 +42,8 @@ type = "adamw" lr = 1e-6 [orchestrator] -batch_size = 512 -oversampling_factor = 2 +token_batch_size = 10485760 # 512 rollouts x ~20k avg tokens +max_inflight_episodes = 1024 max_off_policy_steps = 16 [[orchestrator.train.source]] diff --git a/examples/advanced/qwen3-30b-a3b/tool.toml b/examples/advanced/qwen3-30b-a3b/tool.toml index 7732ff1a6a..6050960e46 100644 --- a/examples/advanced/qwen3-30b-a3b/tool.toml +++ b/examples/advanced/qwen3-30b-a3b/tool.toml @@ -31,7 +31,8 @@ freq = 1 [trainer.model.compile] [orchestrator] -batch_size = 512 +token_batch_size = 2097152 # 512 rollouts x ~4k avg tokens +max_inflight_episodes = 512 group_size = 16 max_off_policy_steps = 32 diff --git a/examples/basic/alphabet-sort/rl.toml b/examples/basic/alphabet-sort/rl.toml index fade9b0c4a..402d8c7e9d 100644 --- a/examples/basic/alphabet-sort/rl.toml +++ b/examples/basic/alphabet-sort/rl.toml @@ -25,7 +25,8 @@ alpha = 64 lr = 1e-5 [orchestrator] -batch_size = 256 +token_batch_size = 262144 # 256 rollouts x ~1k avg tokens +max_inflight_episodes = 256 group_size = 16 [orchestrator.train.sampling] diff --git a/examples/basic/hendrycks-sanity/rl.toml b/examples/basic/hendrycks-sanity/rl.toml index 7c0b80ca63..535b3f5317 100644 --- a/examples/basic/hendrycks-sanity/rl.toml +++ b/examples/basic/hendrycks-sanity/rl.toml @@ -12,7 +12,8 @@ num_infer_gpus = 4 name = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B" [orchestrator] -batch_size = 512 +token_batch_size = 2097152 # 512 rollouts x ~4k avg tokens +max_inflight_episodes = 512 group_size = 8 seq_len = 8192 diff --git a/examples/basic/reverse-text/rl.toml b/examples/basic/reverse-text/rl.toml index 184686ead1..e0f8eb6311 100644 --- a/examples/basic/reverse-text/rl.toml +++ b/examples/basic/reverse-text/rl.toml @@ -9,7 +9,7 @@ project = "reverse-text" name = "reverse-text" [orchestrator] -batch_size = 128 +token_batch_size = 16384 # 128 rollouts x ~128 avg tokens group_size = 16 [orchestrator.train.sampling] diff --git a/examples/basic/wiki-search/rl.toml b/examples/basic/wiki-search/rl.toml index f9f256387b..81791fd3b9 100644 --- a/examples/basic/wiki-search/rl.toml +++ b/examples/basic/wiki-search/rl.toml @@ -30,9 +30,9 @@ target_modules = [ ] [orchestrator] -batch_size = 512 +token_batch_size = 786432 # 512 rollouts x ~1.5k avg tokens +max_inflight_episodes = 1024 group_size = 16 -oversampling_factor = 2.0 [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/examples/basic/wordle/rl.toml b/examples/basic/wordle/rl.toml index 84606b3a30..57293cc4db 100644 --- a/examples/basic/wordle/rl.toml +++ b/examples/basic/wordle/rl.toml @@ -15,7 +15,8 @@ name = "wordle" name = "PrimeIntellect/Qwen3-1.7B-Wordle-SFT" [orchestrator] -batch_size = 512 +token_batch_size = 1572864 # 512 rollouts x ~3k avg tokens +max_inflight_episodes = 512 group_size = 16 [[orchestrator.train.source]] diff --git a/k8s/prime-rl/examples/reverse-text/orch.toml b/k8s/prime-rl/examples/reverse-text/orch.toml index a4361ab8db..09bd9a55ca 100644 --- a/k8s/prime-rl/examples/reverse-text/orch.toml +++ b/k8s/prime-rl/examples/reverse-text/orch.toml @@ -1,6 +1,6 @@ max_steps = 20 seq_len = 2048 -batch_size = 128 +token_batch_size = 16384 # 128 rollouts x ~128 avg tokens group_size = 16 [model] 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..73845898ca 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``.""" + token_batch_size: int = Field(131_072, ge=1) + """Tokens to train on per step. A batch ships once the pending rollouts' trainer-bound payload reaches this many tokens. Size it as (target rollouts per step) x (average tokens per rollout); the default matches the old 128-rollout default at ~1k tokens per rollout.""" - token_batch_size: int | None = Field(None, ge=1) - """Tokens to train on per step (token-based batching). Set this OR ``batch_size``.""" + count_zero_advantage_in_batch: bool = False + """Count zero-advantage rollouts toward ``token_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 tokens.""" - 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``.""" - - 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).""" + max_inflight_episodes: int = Field(128, ge=1) + """Maximum number of episodes kept in-flight — one episode is one agent run at a time, whatever the env's agents are. Tune together with ``token_batch_size``: roughly ``token_batch_size / (average tokens per rollout)``, higher 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,7 @@ 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.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..26a180fdcf 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=max(1, self.orchestrator.token_batch_size // self.orchestrator.seq_len), ) trainer_bench_enabled = self.trainer.bench is not None 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/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index 4a4d99a179..2b3b315111 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 token budget is met. - ``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,6 @@ async def setup(self) -> None: else None ) - assert config.max_inflight_episodes is not None, "max_inflight_episodes must be resolved before dispatcher init" log_interval = config.log.interval wandb_enabled = config.wandb is not None self.dispatcher = RolloutDispatcher( @@ -415,10 +409,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 +591,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 +600,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 +658,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 +684,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 +712,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 +751,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 [] @@ -779,7 +761,7 @@ def collect_pipeline_view(self) -> tuple[str, dict[str, float]]: # Train batch: finalized-group survivors only (0→target). Partial-group # arrivals are surfaced as a separate ``(+N buffered)`` addendum train_pct = train_batch / train_target if train_target else 0.0 - train_batch_part = f"Train batch {train_batch}/{train_target} ({train_pct:.1%})" + train_batch_part = f"Train batch {train_batch}/{train_target} tokens ({train_pct:.1%})" if multi_train: pairs = [(e.name, train_batch_by_env.get(e.name, 0)) for e in self.train_envs] train_batch_part += " (" + ", ".join(f"{n}={v}" for n, v in pairs) + ")" diff --git a/src/prime_rl/orchestrator/train_sink.py b/src/prime_rl/orchestrator/train_sink.py index 5f052c14ed..3982c22dae 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`` — 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,13 +25,23 @@ 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 +# 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 + def payload_tokens(rollout: Rollout) -> int: """Token cost of the rollout's trainer-bound payload — the samples built by @@ -55,22 +66,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.token_batch_size = config.token_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 +87,22 @@ 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. + # Running payload-token total of ``pending_batch``, 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 - ``pending_batch`` (survivors of finalized groups, queued for the + def batch_progress(self) -> tuple[int, int]: + """``(current, target)`` payload tokens 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 self.pending_tokens, self.token_batch_size def buffered_count(self) -> int: """Episodes that have arrived but sit in not-yet-complete groups — @@ -113,11 +110,11 @@ def buffered_count(self) -> int: return sum(self.pending_group_episodes.values()) def pending_batch_by_env(self) -> dict[str, int]: - """Per-env breakdown of ``batch_progress()`` (``pending_batch`` only); - values sum to the aggregate.""" + """Per-env payload-token breakdown of ``batch_progress()`` + (``pending_batch`` only); values sum to the aggregate.""" counts: dict[str, int] = defaultdict(int) for r in self.pending_batch: - counts[r.env_name] += 1 + counts[r.env_name] += payload_tokens(r) return dict(counts) async def add(self, episode: list[Rollout]) -> TrainBatch | None: @@ -137,12 +134,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 self.pending_tokens >= self.token_batch_size: return self.process_batch() return None @@ -167,7 +159,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 +198,66 @@ 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 consumes batch budget, so the + # per-step sampling effort stays fixed while the trained-on + # token 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) + 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 - 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) + """Pop the cohort whose payload tokens reach ``token_batch_size`` off + ``pending_batch`` and assemble the trainer-bound ``TrainingSample`` + list. Overflow stays for the next batch.""" + 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 # 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 +270,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..4450cd124f 100644 --- a/src/prime_rl/utils/monitor/prime.py +++ b/src/prime_rl/utils/monitor/prime.py @@ -164,8 +164,6 @@ 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["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/test_configs.py b/tests/unit/test_configs.py index abeac169e8..6c9526fe11 100644 --- a/tests/unit/test_configs.py +++ b/tests/unit/test_configs.py @@ -557,7 +557,7 @@ def test_shared_output_dir_propagates_through_cli(tmp_path): "seq_len": 128, "model": {"name": "Qwen/Qwen3-0.6B"}, "trainer": {}, - "orchestrator": {"batch_size": 16, "group_size": 1}, + "orchestrator": {"token_batch_size": 2048, "group_size": 1}, "inference": {}, }, ) From d11bee45e108224ec3f2a72202408fcdb44537ed Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Wed, 12 Aug 2026 19:08:19 +0000 Subject: [PATCH 2/4] Round token batch sizes to powers of two Drop the per-config derivation comments and write the large ints with underscore separators. Co-Authored-By: Claude Fable 5 --- configs/basic/alphabet-sort/rl.toml | 2 +- configs/basic/hendrycks-sanity/rl.toml | 2 +- configs/basic/reverse-text/rl.toml | 2 +- configs/basic/wiki-search/rl.toml | 2 +- configs/basic/wordle/rl.toml | 2 +- configs/ci/integration/alphabet_sort.toml | 2 +- configs/ci/integration/reverse-text-lora/resume.toml | 2 +- configs/ci/integration/reverse-text-lora/start.toml | 2 +- configs/ci/integration/reverse-text-moe/start.toml | 2 +- configs/ci/integration/reverse-text-rl-opd/start.toml | 2 +- configs/ci/integration/reverse-text-rl-sft/start.toml | 2 +- configs/ci/integration/reverse-text/resume.toml | 2 +- configs/ci/integration/reverse-text/start.toml | 2 +- configs/ci/nightly-fft/alphabet-sort.toml | 2 +- configs/ci/nightly-fft/hendrycks-sanity.toml | 2 +- configs/ci/nightly-fft/reverse-text.toml | 2 +- configs/ci/nightly-fft/wiki-search.toml | 2 +- configs/ci/nightly-fft/wordle.toml | 2 +- configs/ci/nightly/multimodal_color_codeword.toml | 2 +- configs/debug/algo/echo.toml | 2 +- configs/debug/algo/grpo.toml | 2 +- configs/debug/algo/max_rl.toml | 2 +- configs/debug/algo/mixed_grpo_opd.toml | 2 +- configs/debug/algo/opd.toml | 2 +- configs/debug/algo/opd_lora.toml | 2 +- configs/debug/algo/rae.toml | 2 +- configs/debug/algo/self_distill.toml | 2 +- configs/debug/algo/sft_distill.toml | 2 +- configs/debug/algo/sft_distill_lora.toml | 2 +- configs/debug/multi-env/rl.toml | 2 +- docs/development.md | 2 +- examples/advanced/glm-4.5-air/search.toml | 2 +- examples/advanced/glm-4.5-air/swe.toml | 2 +- examples/advanced/glm-4.5-air/terminal.toml | 2 +- examples/advanced/glm-5.2/swe-llmd.toml | 2 +- examples/advanced/glm-5.2/swe.toml | 4 ++-- examples/advanced/intellect-3.1/rl.toml | 2 +- examples/advanced/minimax-m2.5/swe.toml | 2 +- examples/advanced/nemotron-3-super/swe.toml | 2 +- examples/advanced/qwen3-30b-a3b/math.toml | 2 +- examples/advanced/qwen3-30b-a3b/swe.toml | 2 +- examples/advanced/qwen3-30b-a3b/tool.toml | 2 +- examples/basic/alphabet-sort/rl.toml | 2 +- examples/basic/hendrycks-sanity/rl.toml | 2 +- examples/basic/reverse-text/rl.toml | 2 +- examples/basic/wiki-search/rl.toml | 2 +- examples/basic/wordle/rl.toml | 2 +- k8s/prime-rl/examples/reverse-text/orch.toml | 2 +- 48 files changed, 49 insertions(+), 49 deletions(-) diff --git a/configs/basic/alphabet-sort/rl.toml b/configs/basic/alphabet-sort/rl.toml index aced4a54b6..f2a049a7f0 100644 --- a/configs/basic/alphabet-sort/rl.toml +++ b/configs/basic/alphabet-sort/rl.toml @@ -26,7 +26,7 @@ alpha = 64 lr = 1e-5 [orchestrator] -token_batch_size = 131072 # 128 rollouts x ~1k avg tokens +token_batch_size = 131_072 group_size = 8 [orchestrator.train.sampling] diff --git a/configs/basic/hendrycks-sanity/rl.toml b/configs/basic/hendrycks-sanity/rl.toml index 2ac600b6e9..2274deef56 100644 --- a/configs/basic/hendrycks-sanity/rl.toml +++ b/configs/basic/hendrycks-sanity/rl.toml @@ -15,7 +15,7 @@ project = "hendrycks-sanity" name = "hendrycks-sanity" [orchestrator] -token_batch_size = 524288 # 128 rollouts x ~4k avg tokens +token_batch_size = 524_288 group_size = 8 seq_len = 8192 diff --git a/configs/basic/reverse-text/rl.toml b/configs/basic/reverse-text/rl.toml index 82de868efc..0d38f8bbcd 100644 --- a/configs/basic/reverse-text/rl.toml +++ b/configs/basic/reverse-text/rl.toml @@ -16,7 +16,7 @@ project = "reverse-text" name = "reverse-text" [orchestrator] -token_batch_size = 16384 # 128 rollouts x ~128 avg tokens +token_batch_size = 16_384 group_size = 16 [orchestrator.train.sampling] diff --git a/configs/basic/wiki-search/rl.toml b/configs/basic/wiki-search/rl.toml index 9af64f9fe3..5c8ca7e64e 100644 --- a/configs/basic/wiki-search/rl.toml +++ b/configs/basic/wiki-search/rl.toml @@ -34,7 +34,7 @@ target_modules = [ ] [orchestrator] -token_batch_size = 196608 # 128 rollouts x ~1.5k avg tokens +token_batch_size = 262_144 group_size = 16 [orchestrator.train.sampling] diff --git a/configs/basic/wordle/rl.toml b/configs/basic/wordle/rl.toml index 0594dfb34e..5c63ca2719 100644 --- a/configs/basic/wordle/rl.toml +++ b/configs/basic/wordle/rl.toml @@ -18,7 +18,7 @@ name = "wordle" name = "PrimeIntellect/Qwen3-1.7B-Wordle-SFT" [orchestrator] -token_batch_size = 393216 # 128 rollouts x ~3k avg tokens +token_batch_size = 262_144 group_size = 8 [[orchestrator.train.source]] diff --git a/configs/ci/integration/alphabet_sort.toml b/configs/ci/integration/alphabet_sort.toml index d38fe353ad..46252bc63f 100644 --- a/configs/ci/integration/alphabet_sort.toml +++ b/configs/ci/integration/alphabet_sort.toml @@ -16,7 +16,7 @@ name = "Qwen/Qwen3-0.6B" lr = 1e-5 [orchestrator] -token_batch_size = 131072 # 128 rollouts x ~1k avg tokens +token_batch_size = 131_072 group_size = 8 [orchestrator.train.sampling] diff --git a/configs/ci/integration/reverse-text-lora/resume.toml b/configs/ci/integration/reverse-text-lora/resume.toml index a02839f412..c1a0af577f 100644 --- a/configs/ci/integration/reverse-text-lora/resume.toml +++ b/configs/ci/integration/reverse-text-lora/resume.toml @@ -17,7 +17,7 @@ rank = 8 save_adapter_separately = true [orchestrator] -token_batch_size = 16384 # 128 rollouts x ~128 avg tokens +token_batch_size = 16_384 group_size = 16 [orchestrator.model.lora] diff --git a/configs/ci/integration/reverse-text-lora/start.toml b/configs/ci/integration/reverse-text-lora/start.toml index ec27af4203..e7e3631c53 100644 --- a/configs/ci/integration/reverse-text-lora/start.toml +++ b/configs/ci/integration/reverse-text-lora/start.toml @@ -16,7 +16,7 @@ rank = 8 save_adapter_separately = true [orchestrator] -token_batch_size = 16384 # 128 rollouts x ~128 avg tokens +token_batch_size = 16_384 group_size = 16 [orchestrator.model.lora] diff --git a/configs/ci/integration/reverse-text-moe/start.toml b/configs/ci/integration/reverse-text-moe/start.toml index 19c9e6a280..b2cdc6f660 100644 --- a/configs/ci/integration/reverse-text-moe/start.toml +++ b/configs/ci/integration/reverse-text-moe/start.toml @@ -19,7 +19,7 @@ lr = 3e-6 impl = "custom" [orchestrator] -token_batch_size = 16384 # 128 rollouts x ~128 avg tokens +token_batch_size = 16_384 group_size = 16 [orchestrator.train.sampling] diff --git a/configs/ci/integration/reverse-text-rl-opd/start.toml b/configs/ci/integration/reverse-text-rl-opd/start.toml index 2bd94e2457..409a48fb8c 100644 --- a/configs/ci/integration/reverse-text-rl-opd/start.toml +++ b/configs/ci/integration/reverse-text-rl-opd/start.toml @@ -21,7 +21,7 @@ project = "reverse-text-ci" name = "ci-rl-opd" [orchestrator] -token_batch_size = 16384 # 128 rollouts x ~128 avg tokens +token_batch_size = 16_384 group_size = 16 [orchestrator.algo] diff --git a/configs/ci/integration/reverse-text-rl-sft/start.toml b/configs/ci/integration/reverse-text-rl-sft/start.toml index 02668d1b02..3e37c4a93a 100644 --- a/configs/ci/integration/reverse-text-rl-sft/start.toml +++ b/configs/ci/integration/reverse-text-rl-sft/start.toml @@ -21,7 +21,7 @@ project = "reverse-text-ci" name = "ci-rl-sft" [orchestrator] -token_batch_size = 16384 # 128 rollouts x ~128 avg tokens +token_batch_size = 16_384 group_size = 16 [orchestrator.algo] diff --git a/configs/ci/integration/reverse-text/resume.toml b/configs/ci/integration/reverse-text/resume.toml index 8f2aa964ab..8923d86149 100644 --- a/configs/ci/integration/reverse-text/resume.toml +++ b/configs/ci/integration/reverse-text/resume.toml @@ -17,7 +17,7 @@ name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT" lr = 3e-6 [orchestrator] -token_batch_size = 16384 # 128 rollouts x ~128 avg tokens +token_batch_size = 16_384 group_size = 16 [orchestrator.train.sampling] diff --git a/configs/ci/integration/reverse-text/start.toml b/configs/ci/integration/reverse-text/start.toml index 14a2208e61..f9e168dc7a 100644 --- a/configs/ci/integration/reverse-text/start.toml +++ b/configs/ci/integration/reverse-text/start.toml @@ -16,7 +16,7 @@ name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT" lr = 3e-6 [orchestrator] -token_batch_size = 16384 # 128 rollouts x ~128 avg tokens +token_batch_size = 16_384 group_size = 16 [orchestrator.train.sampling] diff --git a/configs/ci/nightly-fft/alphabet-sort.toml b/configs/ci/nightly-fft/alphabet-sort.toml index a04983b726..42b87569e9 100644 --- a/configs/ci/nightly-fft/alphabet-sort.toml +++ b/configs/ci/nightly-fft/alphabet-sort.toml @@ -15,7 +15,7 @@ name = "alphabet-sort" name = "Qwen/Qwen3-4B-Instruct-2507" [orchestrator] -token_batch_size = 524288 # 512 rollouts x ~1k avg tokens +token_batch_size = 524_288 max_inflight_episodes = 512 group_size = 16 diff --git a/configs/ci/nightly-fft/hendrycks-sanity.toml b/configs/ci/nightly-fft/hendrycks-sanity.toml index 2b8bba3734..08ccbe6014 100644 --- a/configs/ci/nightly-fft/hendrycks-sanity.toml +++ b/configs/ci/nightly-fft/hendrycks-sanity.toml @@ -14,7 +14,7 @@ name = "hendrycks-sanity" name = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B" [orchestrator] -token_batch_size = 2097152 # 512 rollouts x ~4k avg tokens +token_batch_size = 2_097_152 max_inflight_episodes = 512 group_size = 8 seq_len = 8192 diff --git a/configs/ci/nightly-fft/reverse-text.toml b/configs/ci/nightly-fft/reverse-text.toml index 106c663098..5bce53090a 100644 --- a/configs/ci/nightly-fft/reverse-text.toml +++ b/configs/ci/nightly-fft/reverse-text.toml @@ -15,7 +15,7 @@ name = "reverse-text" name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT" [orchestrator] -token_batch_size = 16384 # 128 rollouts x ~128 avg tokens +token_batch_size = 16_384 group_size = 16 [[orchestrator.train.source]] diff --git a/configs/ci/nightly-fft/wiki-search.toml b/configs/ci/nightly-fft/wiki-search.toml index 03f256324d..9b55a50dd5 100644 --- a/configs/ci/nightly-fft/wiki-search.toml +++ b/configs/ci/nightly-fft/wiki-search.toml @@ -15,7 +15,7 @@ name = "wiki-search" name = "Qwen/Qwen3-4B-Instruct-2507" [orchestrator] -token_batch_size = 786432 # 512 rollouts x ~1.5k avg tokens +token_batch_size = 1_048_576 max_inflight_episodes = 1024 group_size = 16 diff --git a/configs/ci/nightly-fft/wordle.toml b/configs/ci/nightly-fft/wordle.toml index 7d159b1916..e3f9c7db3b 100644 --- a/configs/ci/nightly-fft/wordle.toml +++ b/configs/ci/nightly-fft/wordle.toml @@ -15,7 +15,7 @@ name = "wordle" name = "PrimeIntellect/Qwen3-1.7B-Wordle-SFT" [orchestrator] -token_batch_size = 1572864 # 512 rollouts x ~3k avg tokens +token_batch_size = 1_048_576 max_inflight_episodes = 512 group_size = 16 diff --git a/configs/ci/nightly/multimodal_color_codeword.toml b/configs/ci/nightly/multimodal_color_codeword.toml index a870499d11..36539cf192 100644 --- a/configs/ci/nightly/multimodal_color_codeword.toml +++ b/configs/ci/nightly/multimodal_color_codeword.toml @@ -14,7 +14,7 @@ vision_encoder_attr = "model.visual" language_model_attr = "model.language_model" [orchestrator] -token_batch_size = 262144 # 256 rollouts x ~1k avg tokens +token_batch_size = 262_144 max_inflight_episodes = 256 group_size = 16 diff --git a/configs/debug/algo/echo.toml b/configs/debug/algo/echo.toml index 2080aa8fb5..a100053c41 100644 --- a/configs/debug/algo/echo.toml +++ b/configs/debug/algo/echo.toml @@ -13,7 +13,7 @@ project = "algorithms-debug" name = "debug-echo" [orchestrator] -token_batch_size = 32768 # 32 rollouts x ~1k avg tokens +token_batch_size = 32_768 max_inflight_episodes = 32 group_size = 4 diff --git a/configs/debug/algo/grpo.toml b/configs/debug/algo/grpo.toml index f3d2701480..207ba1812e 100644 --- a/configs/debug/algo/grpo.toml +++ b/configs/debug/algo/grpo.toml @@ -9,7 +9,7 @@ project = "algorithms-debug" name = "debug-rl" [orchestrator] -token_batch_size = 16384 # 128 rollouts x ~128 avg tokens +token_batch_size = 16_384 group_size = 16 [orchestrator.algo] diff --git a/configs/debug/algo/max_rl.toml b/configs/debug/algo/max_rl.toml index c6a0b9c057..8d00b71051 100644 --- a/configs/debug/algo/max_rl.toml +++ b/configs/debug/algo/max_rl.toml @@ -9,7 +9,7 @@ project = "algorithms-debug" name = "debug-max-rl" [orchestrator] -token_batch_size = 16384 # 128 rollouts x ~128 avg tokens +token_batch_size = 16_384 group_size = 16 [orchestrator.algo] diff --git a/configs/debug/algo/mixed_grpo_opd.toml b/configs/debug/algo/mixed_grpo_opd.toml index 612f332821..6b080ed7a0 100644 --- a/configs/debug/algo/mixed_grpo_opd.toml +++ b/configs/debug/algo/mixed_grpo_opd.toml @@ -19,7 +19,7 @@ project = "algorithms-debug" name = "debug-mixed-grpo-opd" [orchestrator] -token_batch_size = 16384 # 128 rollouts x ~128 avg tokens +token_batch_size = 16_384 group_size = 16 [orchestrator.algo] diff --git a/configs/debug/algo/opd.toml b/configs/debug/algo/opd.toml index 0fcebe1ef0..54e501f249 100644 --- a/configs/debug/algo/opd.toml +++ b/configs/debug/algo/opd.toml @@ -16,7 +16,7 @@ project = "algorithms-debug" name = "debug-opd" [orchestrator] -token_batch_size = 16384 # 128 rollouts x ~128 avg tokens +token_batch_size = 16_384 group_size = 16 [orchestrator.algo] diff --git a/configs/debug/algo/opd_lora.toml b/configs/debug/algo/opd_lora.toml index ec99c0701a..ec046dceab 100644 --- a/configs/debug/algo/opd_lora.toml +++ b/configs/debug/algo/opd_lora.toml @@ -16,7 +16,7 @@ project = "algorithms-debug" name = "debug-opd-lora" [orchestrator] -token_batch_size = 16384 # 128 rollouts x ~128 avg tokens +token_batch_size = 16_384 group_size = 16 [orchestrator.algo] diff --git a/configs/debug/algo/rae.toml b/configs/debug/algo/rae.toml index 3c27659381..1af70ba0ca 100644 --- a/configs/debug/algo/rae.toml +++ b/configs/debug/algo/rae.toml @@ -9,7 +9,7 @@ project = "algorithms-debug" name = "debug-rae" [orchestrator] -token_batch_size = 16384 # 32 rollouts x ~512 avg tokens +token_batch_size = 16_384 max_inflight_episodes = 32 group_size = 1 diff --git a/configs/debug/algo/self_distill.toml b/configs/debug/algo/self_distill.toml index 1378e7d512..deedea7cb0 100644 --- a/configs/debug/algo/self_distill.toml +++ b/configs/debug/algo/self_distill.toml @@ -15,7 +15,7 @@ project = "algorithms-debug" name = "debug-self-distill" [orchestrator] -token_batch_size = 4096 # 32 rollouts x ~128 avg tokens +token_batch_size = 4096 max_inflight_episodes = 32 group_size = 1 diff --git a/configs/debug/algo/sft_distill.toml b/configs/debug/algo/sft_distill.toml index dcb5e0eb94..3e35107d67 100644 --- a/configs/debug/algo/sft_distill.toml +++ b/configs/debug/algo/sft_distill.toml @@ -16,7 +16,7 @@ project = "algorithms-debug" name = "debug-sft" [orchestrator] -token_batch_size = 16384 # 128 rollouts x ~128 avg tokens +token_batch_size = 16_384 group_size = 4 [orchestrator.algo] diff --git a/configs/debug/algo/sft_distill_lora.toml b/configs/debug/algo/sft_distill_lora.toml index ee9075bd0d..7377d14daf 100644 --- a/configs/debug/algo/sft_distill_lora.toml +++ b/configs/debug/algo/sft_distill_lora.toml @@ -16,7 +16,7 @@ project = "algorithms-debug" name = "debug-sft-lora" [orchestrator] -token_batch_size = 16384 # 128 rollouts x ~128 avg tokens +token_batch_size = 16_384 group_size = 4 [orchestrator.algo] diff --git a/configs/debug/multi-env/rl.toml b/configs/debug/multi-env/rl.toml index 9debdf5d8f..67850894aa 100644 --- a/configs/debug/multi-env/rl.toml +++ b/configs/debug/multi-env/rl.toml @@ -8,7 +8,7 @@ seq_len = 2048 name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT" [orchestrator] -token_batch_size = 16384 # 128 rollouts x ~128 avg tokens +token_batch_size = 16_384 group_size = 16 [orchestrator.renderer] diff --git a/docs/development.md b/docs/development.md index 0e25f813c7..2378d8947f 100644 --- a/docs/development.md +++ b/docs/development.md @@ -137,7 +137,7 @@ Before merging a new model, you need to ensure the following: - The model is correctly registered and defines and all the required methods - such as `convert_hf_layer_to_tt` and `convert_tt_layer_to_hf`. - The small smoke test passes. -In the PR that adds the new model, you also need to provide a table covering the KL mismatch across 20 steps on `math` environment with `token_batch_size=524288` (~64 rollouts at ~8k avg tokens). All the entries in the table must lower than 0.015. If this is not met, the PR will not be merged (unless reasonable justification is provided). This is to ensure all our models are consistent and their implementations match the implementations in the inference framework. +In the PR that adds the new model, you also need to provide a table covering the KL mismatch across 20 steps on `math` environment with `token_batch_size=524_288` (~64 rollouts at ~8k avg tokens). All the entries in the table must lower than 0.015. If this is not met, the PR will not be merged (unless reasonable justification is provided). This is to ensure all our models are consistent and their implementations match the implementations in the inference framework. ## Adding a Custom VLM Implementation diff --git a/examples/advanced/glm-4.5-air/search.toml b/examples/advanced/glm-4.5-air/search.toml index a1a9678b16..53c8af7d11 100644 --- a/examples/advanced/glm-4.5-air/search.toml +++ b/examples/advanced/glm-4.5-air/search.toml @@ -68,7 +68,7 @@ type = "muon" # --- Orchestrator --- [orchestrator] -token_batch_size = 4194304 # 256 rollouts x ~16k avg tokens +token_batch_size = 4_194_304 group_size = 16 max_off_policy_steps = 32 max_inflight_episodes = 512 diff --git a/examples/advanced/glm-4.5-air/swe.toml b/examples/advanced/glm-4.5-air/swe.toml index 197a13f575..f7603c3b68 100644 --- a/examples/advanced/glm-4.5-air/swe.toml +++ b/examples/advanced/glm-4.5-air/swe.toml @@ -79,7 +79,7 @@ type = "muon" # --- Orchestrator --- [orchestrator] -token_batch_size = 5242880 # 256 rollouts x ~20k avg tokens +token_batch_size = 4_194_304 group_size = 16 max_off_policy_steps = 32 max_inflight_episodes = 512 diff --git a/examples/advanced/glm-4.5-air/terminal.toml b/examples/advanced/glm-4.5-air/terminal.toml index f130fabb36..bf2fca82b8 100644 --- a/examples/advanced/glm-4.5-air/terminal.toml +++ b/examples/advanced/glm-4.5-air/terminal.toml @@ -69,7 +69,7 @@ type = "muon" # --- Orchestrator --- [orchestrator] -token_batch_size = 4194304 # 256 rollouts x ~16k avg tokens +token_batch_size = 4_194_304 group_size = 16 max_off_policy_steps = 32 max_inflight_episodes = 512 diff --git a/examples/advanced/glm-5.2/swe-llmd.toml b/examples/advanced/glm-5.2/swe-llmd.toml index 4daebab432..acfa92d69f 100644 --- a/examples/advanced/glm-5.2/swe-llmd.toml +++ b/examples/advanced/glm-5.2/swe-llmd.toml @@ -109,7 +109,7 @@ lr = 1e-6 weight_decay = 0.0 [orchestrator] -token_batch_size = 5242880 # 256 rollouts x ~20k avg tokens +token_batch_size = 4_194_304 group_size = 16 max_inflight_episodes = 2048 max_off_policy_steps = 8 # 16/32 can be better, 8 is more CPU stable diff --git a/examples/advanced/glm-5.2/swe.toml b/examples/advanced/glm-5.2/swe.toml index dc10d26ce2..65306bf5da 100644 --- a/examples/advanced/glm-5.2/swe.toml +++ b/examples/advanced/glm-5.2/swe.toml @@ -64,8 +64,8 @@ lr = 1e-6 weight_decay = 0.1 [orchestrator] -token_batch_size = 83886080 # 4096 rollouts x ~20k avg tokens -max_inflight_episodes = 12288 +token_batch_size = 67_108_864 +max_inflight_episodes = 12_288 group_size = 16 max_off_policy_steps = 16 diff --git a/examples/advanced/intellect-3.1/rl.toml b/examples/advanced/intellect-3.1/rl.toml index 93d08b9a88..6ac1c3e741 100644 --- a/examples/advanced/intellect-3.1/rl.toml +++ b/examples/advanced/intellect-3.1/rl.toml @@ -44,7 +44,7 @@ lr = 1e-6 weight_decay = 0.01 [orchestrator] -token_batch_size = 33554432 # 2048 rollouts x ~16k avg tokens +token_batch_size = 33_554_432 max_inflight_episodes = 4096 [[orchestrator.train.source]] diff --git a/examples/advanced/minimax-m2.5/swe.toml b/examples/advanced/minimax-m2.5/swe.toml index 652c7e4afe..fb4c06b904 100644 --- a/examples/advanced/minimax-m2.5/swe.toml +++ b/examples/advanced/minimax-m2.5/swe.toml @@ -44,7 +44,7 @@ lr = 1e-6 weight_decay = 0.01 [orchestrator] -token_batch_size = 41943040 # 2048 rollouts x ~20k avg tokens +token_batch_size = 33_554_432 max_inflight_episodes = 4096 max_off_policy_steps = 16 diff --git a/examples/advanced/nemotron-3-super/swe.toml b/examples/advanced/nemotron-3-super/swe.toml index 839458f329..96b1853cdb 100644 --- a/examples/advanced/nemotron-3-super/swe.toml +++ b/examples/advanced/nemotron-3-super/swe.toml @@ -63,7 +63,7 @@ skip_optimizer = true type = "adamw" [orchestrator] -token_batch_size = 5242880 # 256 rollouts x ~20k avg tokens +token_batch_size = 4_194_304 group_size = 16 max_off_policy_steps = 32 max_inflight_episodes = 1536 diff --git a/examples/advanced/qwen3-30b-a3b/math.toml b/examples/advanced/qwen3-30b-a3b/math.toml index e236e91de1..da9b0254c9 100644 --- a/examples/advanced/qwen3-30b-a3b/math.toml +++ b/examples/advanced/qwen3-30b-a3b/math.toml @@ -41,7 +41,7 @@ type = "adamw" lr = 1e-6 [orchestrator] -token_batch_size = 4194304 # 512 rollouts x ~8k avg tokens +token_batch_size = 4_194_304 max_inflight_episodes = 1024 max_off_policy_steps = 8 diff --git a/examples/advanced/qwen3-30b-a3b/swe.toml b/examples/advanced/qwen3-30b-a3b/swe.toml index 39fdb17f3b..63a53d0148 100644 --- a/examples/advanced/qwen3-30b-a3b/swe.toml +++ b/examples/advanced/qwen3-30b-a3b/swe.toml @@ -42,7 +42,7 @@ type = "adamw" lr = 1e-6 [orchestrator] -token_batch_size = 10485760 # 512 rollouts x ~20k avg tokens +token_batch_size = 8_388_608 max_inflight_episodes = 1024 max_off_policy_steps = 16 diff --git a/examples/advanced/qwen3-30b-a3b/tool.toml b/examples/advanced/qwen3-30b-a3b/tool.toml index 6050960e46..74b9823a63 100644 --- a/examples/advanced/qwen3-30b-a3b/tool.toml +++ b/examples/advanced/qwen3-30b-a3b/tool.toml @@ -31,7 +31,7 @@ freq = 1 [trainer.model.compile] [orchestrator] -token_batch_size = 2097152 # 512 rollouts x ~4k avg tokens +token_batch_size = 2_097_152 max_inflight_episodes = 512 group_size = 16 max_off_policy_steps = 32 diff --git a/examples/basic/alphabet-sort/rl.toml b/examples/basic/alphabet-sort/rl.toml index 402d8c7e9d..f1962d8c1c 100644 --- a/examples/basic/alphabet-sort/rl.toml +++ b/examples/basic/alphabet-sort/rl.toml @@ -25,7 +25,7 @@ alpha = 64 lr = 1e-5 [orchestrator] -token_batch_size = 262144 # 256 rollouts x ~1k avg tokens +token_batch_size = 262_144 max_inflight_episodes = 256 group_size = 16 diff --git a/examples/basic/hendrycks-sanity/rl.toml b/examples/basic/hendrycks-sanity/rl.toml index 535b3f5317..9bcef7a7d2 100644 --- a/examples/basic/hendrycks-sanity/rl.toml +++ b/examples/basic/hendrycks-sanity/rl.toml @@ -12,7 +12,7 @@ num_infer_gpus = 4 name = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B" [orchestrator] -token_batch_size = 2097152 # 512 rollouts x ~4k avg tokens +token_batch_size = 2_097_152 max_inflight_episodes = 512 group_size = 8 seq_len = 8192 diff --git a/examples/basic/reverse-text/rl.toml b/examples/basic/reverse-text/rl.toml index e0f8eb6311..4f8a92b6d7 100644 --- a/examples/basic/reverse-text/rl.toml +++ b/examples/basic/reverse-text/rl.toml @@ -9,7 +9,7 @@ project = "reverse-text" name = "reverse-text" [orchestrator] -token_batch_size = 16384 # 128 rollouts x ~128 avg tokens +token_batch_size = 16_384 group_size = 16 [orchestrator.train.sampling] diff --git a/examples/basic/wiki-search/rl.toml b/examples/basic/wiki-search/rl.toml index 81791fd3b9..220a9a051a 100644 --- a/examples/basic/wiki-search/rl.toml +++ b/examples/basic/wiki-search/rl.toml @@ -30,7 +30,7 @@ target_modules = [ ] [orchestrator] -token_batch_size = 786432 # 512 rollouts x ~1.5k avg tokens +token_batch_size = 1_048_576 max_inflight_episodes = 1024 group_size = 16 diff --git a/examples/basic/wordle/rl.toml b/examples/basic/wordle/rl.toml index 57293cc4db..0cea3a65c1 100644 --- a/examples/basic/wordle/rl.toml +++ b/examples/basic/wordle/rl.toml @@ -15,7 +15,7 @@ name = "wordle" name = "PrimeIntellect/Qwen3-1.7B-Wordle-SFT" [orchestrator] -token_batch_size = 1572864 # 512 rollouts x ~3k avg tokens +token_batch_size = 1_048_576 max_inflight_episodes = 512 group_size = 16 diff --git a/k8s/prime-rl/examples/reverse-text/orch.toml b/k8s/prime-rl/examples/reverse-text/orch.toml index 09bd9a55ca..86962dff89 100644 --- a/k8s/prime-rl/examples/reverse-text/orch.toml +++ b/k8s/prime-rl/examples/reverse-text/orch.toml @@ -1,6 +1,6 @@ max_steps = 20 seq_len = 2048 -token_batch_size = 16384 # 128 rollouts x ~128 avg tokens +token_batch_size = 16_384 group_size = 16 [model] From cebea844f269c8f962b0c047bca292bd37fda70b Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Wed, 12 Aug 2026 19:47:24 +0000 Subject: [PATCH 3/4] Remove token batching, keep batch_size as the batching unit batch_size and group_size stay the only batching knobs. Also drop oversampling_factor: max_inflight_episodes defaults to batch_size, and configs that oversampled now set it explicitly. Co-Authored-By: Claude Fable 5 --- configs/basic/alphabet-sort/rl.toml | 2 +- configs/basic/hendrycks-sanity/rl.toml | 2 +- configs/basic/reverse-text/rl.toml | 2 +- configs/basic/wiki-search/rl.toml | 2 +- configs/basic/wordle/rl.toml | 2 +- configs/ci/integration/alphabet_sort.toml | 2 +- .../integration/reverse-text-lora/resume.toml | 2 +- .../integration/reverse-text-lora/start.toml | 2 +- .../integration/reverse-text-moe/start.toml | 2 +- .../reverse-text-rl-opd/start.toml | 2 +- .../reverse-text-rl-sft/start.toml | 2 +- .../ci/integration/reverse-text/resume.toml | 2 +- .../ci/integration/reverse-text/start.toml | 2 +- configs/ci/nightly-fft/alphabet-sort.toml | 3 +- configs/ci/nightly-fft/hendrycks-sanity.toml | 3 +- configs/ci/nightly-fft/reverse-text.toml | 2 +- configs/ci/nightly-fft/wiki-search.toml | 4 +- configs/ci/nightly-fft/wordle.toml | 3 +- .../ci/nightly/multimodal_color_codeword.toml | 3 +- configs/debug/algo/echo.toml | 3 +- configs/debug/algo/grpo.toml | 2 +- configs/debug/algo/max_rl.toml | 2 +- configs/debug/algo/mixed_grpo_opd.toml | 2 +- configs/debug/algo/opd.toml | 2 +- configs/debug/algo/opd_lora.toml | 2 +- configs/debug/algo/rae.toml | 3 +- configs/debug/algo/self_distill.toml | 3 +- configs/debug/algo/sft_distill.toml | 2 +- configs/debug/algo/sft_distill_lora.toml | 2 +- configs/debug/multi-env/rl.toml | 2 +- docs/algorithms.md | 2 +- docs/development.md | 2 +- docs/training.md | 4 +- examples/advanced/glm-4.5-air/search.toml | 2 +- examples/advanced/glm-4.5-air/swe.toml | 2 +- examples/advanced/glm-4.5-air/terminal.toml | 2 +- examples/advanced/glm-5.2/swe-llmd.toml | 2 +- examples/advanced/glm-5.2/swe.toml | 4 +- examples/advanced/intellect-3.1/rl.toml | 2 +- examples/advanced/minimax-m2.5/swe.toml | 2 +- examples/advanced/nemotron-3-super/swe.toml | 2 +- examples/advanced/qwen3-30b-a3b/math.toml | 2 +- examples/advanced/qwen3-30b-a3b/swe.toml | 2 +- examples/advanced/qwen3-30b-a3b/tool.toml | 3 +- examples/basic/alphabet-sort/rl.toml | 3 +- examples/basic/hendrycks-sanity/rl.toml | 3 +- examples/basic/reverse-text/rl.toml | 2 +- examples/basic/wiki-search/rl.toml | 4 +- examples/basic/wordle/rl.toml | 3 +- k8s/prime-rl/examples/reverse-text/orch.toml | 2 +- .../src/prime_rl/configs/orchestrator.py | 14 +++-- .../src/prime_rl/configs/rl.py | 2 +- src/prime_rl/orchestrator/orchestrator.py | 5 +- src/prime_rl/orchestrator/train_sink.py | 59 ++++++------------- src/prime_rl/utils/monitor/prime.py | 1 + tests/unit/test_configs.py | 2 +- 56 files changed, 86 insertions(+), 116 deletions(-) diff --git a/configs/basic/alphabet-sort/rl.toml b/configs/basic/alphabet-sort/rl.toml index f2a049a7f0..bc70f7718d 100644 --- a/configs/basic/alphabet-sort/rl.toml +++ b/configs/basic/alphabet-sort/rl.toml @@ -26,7 +26,7 @@ alpha = 64 lr = 1e-5 [orchestrator] -token_batch_size = 131_072 +batch_size = 128 group_size = 8 [orchestrator.train.sampling] diff --git a/configs/basic/hendrycks-sanity/rl.toml b/configs/basic/hendrycks-sanity/rl.toml index 2274deef56..b16c7c7c4c 100644 --- a/configs/basic/hendrycks-sanity/rl.toml +++ b/configs/basic/hendrycks-sanity/rl.toml @@ -15,7 +15,7 @@ project = "hendrycks-sanity" name = "hendrycks-sanity" [orchestrator] -token_batch_size = 524_288 +batch_size = 128 group_size = 8 seq_len = 8192 diff --git a/configs/basic/reverse-text/rl.toml b/configs/basic/reverse-text/rl.toml index 0d38f8bbcd..3344560e99 100644 --- a/configs/basic/reverse-text/rl.toml +++ b/configs/basic/reverse-text/rl.toml @@ -16,7 +16,7 @@ project = "reverse-text" name = "reverse-text" [orchestrator] -token_batch_size = 16_384 +batch_size = 128 group_size = 16 [orchestrator.train.sampling] diff --git a/configs/basic/wiki-search/rl.toml b/configs/basic/wiki-search/rl.toml index 5c8ca7e64e..aa509e9049 100644 --- a/configs/basic/wiki-search/rl.toml +++ b/configs/basic/wiki-search/rl.toml @@ -34,7 +34,7 @@ target_modules = [ ] [orchestrator] -token_batch_size = 262_144 +batch_size = 128 group_size = 16 [orchestrator.train.sampling] diff --git a/configs/basic/wordle/rl.toml b/configs/basic/wordle/rl.toml index 5c63ca2719..fff93ac571 100644 --- a/configs/basic/wordle/rl.toml +++ b/configs/basic/wordle/rl.toml @@ -18,7 +18,7 @@ name = "wordle" name = "PrimeIntellect/Qwen3-1.7B-Wordle-SFT" [orchestrator] -token_batch_size = 262_144 +batch_size = 128 group_size = 8 [[orchestrator.train.source]] diff --git a/configs/ci/integration/alphabet_sort.toml b/configs/ci/integration/alphabet_sort.toml index 46252bc63f..51472260e4 100644 --- a/configs/ci/integration/alphabet_sort.toml +++ b/configs/ci/integration/alphabet_sort.toml @@ -16,7 +16,7 @@ name = "Qwen/Qwen3-0.6B" lr = 1e-5 [orchestrator] -token_batch_size = 131_072 +batch_size = 128 group_size = 8 [orchestrator.train.sampling] diff --git a/configs/ci/integration/reverse-text-lora/resume.toml b/configs/ci/integration/reverse-text-lora/resume.toml index c1a0af577f..ea6437df67 100644 --- a/configs/ci/integration/reverse-text-lora/resume.toml +++ b/configs/ci/integration/reverse-text-lora/resume.toml @@ -17,7 +17,7 @@ rank = 8 save_adapter_separately = true [orchestrator] -token_batch_size = 16_384 +batch_size = 128 group_size = 16 [orchestrator.model.lora] diff --git a/configs/ci/integration/reverse-text-lora/start.toml b/configs/ci/integration/reverse-text-lora/start.toml index e7e3631c53..74789ca4bc 100644 --- a/configs/ci/integration/reverse-text-lora/start.toml +++ b/configs/ci/integration/reverse-text-lora/start.toml @@ -16,7 +16,7 @@ rank = 8 save_adapter_separately = true [orchestrator] -token_batch_size = 16_384 +batch_size = 128 group_size = 16 [orchestrator.model.lora] diff --git a/configs/ci/integration/reverse-text-moe/start.toml b/configs/ci/integration/reverse-text-moe/start.toml index b2cdc6f660..8748737f55 100644 --- a/configs/ci/integration/reverse-text-moe/start.toml +++ b/configs/ci/integration/reverse-text-moe/start.toml @@ -19,7 +19,7 @@ lr = 3e-6 impl = "custom" [orchestrator] -token_batch_size = 16_384 +batch_size = 128 group_size = 16 [orchestrator.train.sampling] diff --git a/configs/ci/integration/reverse-text-rl-opd/start.toml b/configs/ci/integration/reverse-text-rl-opd/start.toml index 409a48fb8c..80fdc80bba 100644 --- a/configs/ci/integration/reverse-text-rl-opd/start.toml +++ b/configs/ci/integration/reverse-text-rl-opd/start.toml @@ -21,7 +21,7 @@ project = "reverse-text-ci" name = "ci-rl-opd" [orchestrator] -token_batch_size = 16_384 +batch_size = 128 group_size = 16 [orchestrator.algo] diff --git a/configs/ci/integration/reverse-text-rl-sft/start.toml b/configs/ci/integration/reverse-text-rl-sft/start.toml index 3e37c4a93a..dfd829b785 100644 --- a/configs/ci/integration/reverse-text-rl-sft/start.toml +++ b/configs/ci/integration/reverse-text-rl-sft/start.toml @@ -21,7 +21,7 @@ project = "reverse-text-ci" name = "ci-rl-sft" [orchestrator] -token_batch_size = 16_384 +batch_size = 128 group_size = 16 [orchestrator.algo] diff --git a/configs/ci/integration/reverse-text/resume.toml b/configs/ci/integration/reverse-text/resume.toml index 8923d86149..9b34f9835c 100644 --- a/configs/ci/integration/reverse-text/resume.toml +++ b/configs/ci/integration/reverse-text/resume.toml @@ -17,7 +17,7 @@ name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT" lr = 3e-6 [orchestrator] -token_batch_size = 16_384 +batch_size = 128 group_size = 16 [orchestrator.train.sampling] diff --git a/configs/ci/integration/reverse-text/start.toml b/configs/ci/integration/reverse-text/start.toml index f9e168dc7a..4c38c2e66e 100644 --- a/configs/ci/integration/reverse-text/start.toml +++ b/configs/ci/integration/reverse-text/start.toml @@ -16,7 +16,7 @@ name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT" lr = 3e-6 [orchestrator] -token_batch_size = 16_384 +batch_size = 128 group_size = 16 [orchestrator.train.sampling] diff --git a/configs/ci/nightly-fft/alphabet-sort.toml b/configs/ci/nightly-fft/alphabet-sort.toml index 42b87569e9..5ccdf93587 100644 --- a/configs/ci/nightly-fft/alphabet-sort.toml +++ b/configs/ci/nightly-fft/alphabet-sort.toml @@ -15,8 +15,7 @@ name = "alphabet-sort" name = "Qwen/Qwen3-4B-Instruct-2507" [orchestrator] -token_batch_size = 524_288 -max_inflight_episodes = 512 +batch_size = 512 group_size = 16 [[orchestrator.train.source]] diff --git a/configs/ci/nightly-fft/hendrycks-sanity.toml b/configs/ci/nightly-fft/hendrycks-sanity.toml index 08ccbe6014..8eb547efee 100644 --- a/configs/ci/nightly-fft/hendrycks-sanity.toml +++ b/configs/ci/nightly-fft/hendrycks-sanity.toml @@ -14,8 +14,7 @@ name = "hendrycks-sanity" name = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B" [orchestrator] -token_batch_size = 2_097_152 -max_inflight_episodes = 512 +batch_size = 512 group_size = 8 seq_len = 8192 diff --git a/configs/ci/nightly-fft/reverse-text.toml b/configs/ci/nightly-fft/reverse-text.toml index 5bce53090a..50ddec8eff 100644 --- a/configs/ci/nightly-fft/reverse-text.toml +++ b/configs/ci/nightly-fft/reverse-text.toml @@ -15,7 +15,7 @@ name = "reverse-text" name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT" [orchestrator] -token_batch_size = 16_384 +batch_size = 128 group_size = 16 [[orchestrator.train.source]] diff --git a/configs/ci/nightly-fft/wiki-search.toml b/configs/ci/nightly-fft/wiki-search.toml index 9b55a50dd5..008cb2dd71 100644 --- a/configs/ci/nightly-fft/wiki-search.toml +++ b/configs/ci/nightly-fft/wiki-search.toml @@ -15,9 +15,9 @@ name = "wiki-search" name = "Qwen/Qwen3-4B-Instruct-2507" [orchestrator] -token_batch_size = 1_048_576 -max_inflight_episodes = 1024 +batch_size = 512 group_size = 16 +max_inflight_episodes = 1024 [[orchestrator.train.source]] name = "wiki-search" diff --git a/configs/ci/nightly-fft/wordle.toml b/configs/ci/nightly-fft/wordle.toml index e3f9c7db3b..d31533836e 100644 --- a/configs/ci/nightly-fft/wordle.toml +++ b/configs/ci/nightly-fft/wordle.toml @@ -15,8 +15,7 @@ name = "wordle" name = "PrimeIntellect/Qwen3-1.7B-Wordle-SFT" [orchestrator] -token_batch_size = 1_048_576 -max_inflight_episodes = 512 +batch_size = 512 group_size = 16 [[orchestrator.train.source]] diff --git a/configs/ci/nightly/multimodal_color_codeword.toml b/configs/ci/nightly/multimodal_color_codeword.toml index 36539cf192..f4b97dd4b3 100644 --- a/configs/ci/nightly/multimodal_color_codeword.toml +++ b/configs/ci/nightly/multimodal_color_codeword.toml @@ -14,8 +14,7 @@ vision_encoder_attr = "model.visual" language_model_attr = "model.language_model" [orchestrator] -token_batch_size = 262_144 -max_inflight_episodes = 256 +batch_size = 256 group_size = 16 [orchestrator.train.sampling] diff --git a/configs/debug/algo/echo.toml b/configs/debug/algo/echo.toml index a100053c41..722f961e59 100644 --- a/configs/debug/algo/echo.toml +++ b/configs/debug/algo/echo.toml @@ -13,8 +13,7 @@ project = "algorithms-debug" name = "debug-echo" [orchestrator] -token_batch_size = 32_768 -max_inflight_episodes = 32 +batch_size = 32 group_size = 4 # alphabet-sort's feedback arrives as user messages, so train the user role diff --git a/configs/debug/algo/grpo.toml b/configs/debug/algo/grpo.toml index 207ba1812e..30dd90b660 100644 --- a/configs/debug/algo/grpo.toml +++ b/configs/debug/algo/grpo.toml @@ -9,7 +9,7 @@ project = "algorithms-debug" name = "debug-rl" [orchestrator] -token_batch_size = 16_384 +batch_size = 128 group_size = 16 [orchestrator.algo] diff --git a/configs/debug/algo/max_rl.toml b/configs/debug/algo/max_rl.toml index 8d00b71051..15c5e83f19 100644 --- a/configs/debug/algo/max_rl.toml +++ b/configs/debug/algo/max_rl.toml @@ -9,7 +9,7 @@ project = "algorithms-debug" name = "debug-max-rl" [orchestrator] -token_batch_size = 16_384 +batch_size = 128 group_size = 16 [orchestrator.algo] diff --git a/configs/debug/algo/mixed_grpo_opd.toml b/configs/debug/algo/mixed_grpo_opd.toml index 6b080ed7a0..c56b37ad11 100644 --- a/configs/debug/algo/mixed_grpo_opd.toml +++ b/configs/debug/algo/mixed_grpo_opd.toml @@ -19,7 +19,7 @@ project = "algorithms-debug" name = "debug-mixed-grpo-opd" [orchestrator] -token_batch_size = 16_384 +batch_size = 128 group_size = 16 [orchestrator.algo] diff --git a/configs/debug/algo/opd.toml b/configs/debug/algo/opd.toml index 54e501f249..7b33bac34e 100644 --- a/configs/debug/algo/opd.toml +++ b/configs/debug/algo/opd.toml @@ -16,7 +16,7 @@ project = "algorithms-debug" name = "debug-opd" [orchestrator] -token_batch_size = 16_384 +batch_size = 128 group_size = 16 [orchestrator.algo] diff --git a/configs/debug/algo/opd_lora.toml b/configs/debug/algo/opd_lora.toml index ec046dceab..de66ace6d0 100644 --- a/configs/debug/algo/opd_lora.toml +++ b/configs/debug/algo/opd_lora.toml @@ -16,7 +16,7 @@ project = "algorithms-debug" name = "debug-opd-lora" [orchestrator] -token_batch_size = 16_384 +batch_size = 128 group_size = 16 [orchestrator.algo] diff --git a/configs/debug/algo/rae.toml b/configs/debug/algo/rae.toml index 1af70ba0ca..65b04cf9d1 100644 --- a/configs/debug/algo/rae.toml +++ b/configs/debug/algo/rae.toml @@ -9,8 +9,7 @@ project = "algorithms-debug" name = "debug-rae" [orchestrator] -token_batch_size = 16_384 -max_inflight_episodes = 32 +batch_size = 32 group_size = 1 [orchestrator.algo] diff --git a/configs/debug/algo/self_distill.toml b/configs/debug/algo/self_distill.toml index deedea7cb0..0b4dfbe869 100644 --- a/configs/debug/algo/self_distill.toml +++ b/configs/debug/algo/self_distill.toml @@ -15,8 +15,7 @@ project = "algorithms-debug" name = "debug-self-distill" [orchestrator] -token_batch_size = 4096 -max_inflight_episodes = 32 +batch_size = 32 group_size = 1 # reverse-text's demo lives in the "answer" column. diff --git a/configs/debug/algo/sft_distill.toml b/configs/debug/algo/sft_distill.toml index 3e35107d67..cf19e32f37 100644 --- a/configs/debug/algo/sft_distill.toml +++ b/configs/debug/algo/sft_distill.toml @@ -16,7 +16,7 @@ project = "algorithms-debug" name = "debug-sft" [orchestrator] -token_batch_size = 16_384 +batch_size = 128 group_size = 4 [orchestrator.algo] diff --git a/configs/debug/algo/sft_distill_lora.toml b/configs/debug/algo/sft_distill_lora.toml index 7377d14daf..deb639c3d8 100644 --- a/configs/debug/algo/sft_distill_lora.toml +++ b/configs/debug/algo/sft_distill_lora.toml @@ -16,7 +16,7 @@ project = "algorithms-debug" name = "debug-sft-lora" [orchestrator] -token_batch_size = 16_384 +batch_size = 128 group_size = 4 [orchestrator.algo] diff --git a/configs/debug/multi-env/rl.toml b/configs/debug/multi-env/rl.toml index 67850894aa..ee77bbb1c6 100644 --- a/configs/debug/multi-env/rl.toml +++ b/configs/debug/multi-env/rl.toml @@ -8,7 +8,7 @@ seq_len = 2048 name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT" [orchestrator] -token_batch_size = 16_384 +batch_size = 128 group_size = 16 [orchestrator.renderer] diff --git a/docs/algorithms.md b/docs/algorithms.md index 3de772d3ea..c38b19f194 100644 --- a/docs/algorithms.md +++ b/docs/algorithms.md @@ -466,7 +466,7 @@ Between scoring and training the sink runs three hardcoded checks on every train | `repetition` | Detects long high-confidence loops. Monitor-only: tracked in metrics (`filters/repetition`), 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. | -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 `token_batch_size` — a fixed sampling budget per step, at the cost of a variable number of trained-on tokens. +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. Dropped rollouts still appear in W&B distributions and metrics (`is_filtered`, `filters/zero_advantage`), just not in the trainer batch. diff --git a/docs/development.md b/docs/development.md index 2378d8947f..32b631500d 100644 --- a/docs/development.md +++ b/docs/development.md @@ -137,7 +137,7 @@ Before merging a new model, you need to ensure the following: - The model is correctly registered and defines and all the required methods - such as `convert_hf_layer_to_tt` and `convert_tt_layer_to_hf`. - The small smoke test passes. -In the PR that adds the new model, you also need to provide a table covering the KL mismatch across 20 steps on `math` environment with `token_batch_size=524_288` (~64 rollouts at ~8k avg tokens). All the entries in the table must lower than 0.015. If this is not met, the PR will not be merged (unless reasonable justification is provided). This is to ensure all our models are consistent and their implementations match the implementations in the inference framework. +In the PR that adds the new model, you also need to provide a table covering the KL mismatch across 20 steps on `math` environment with `batch_size=64`. All the entries in the table must lower than 0.015. If this is not met, the PR will not be merged (unless reasonable justification is provided). This is to ensure all our models are consistent and their implementations match the implementations in the inference framework. ## Adding a Custom VLM Implementation diff --git a/docs/training.md b/docs/training.md index 911585e472..a2a105d36b 100644 --- a/docs/training.md +++ b/docs/training.md @@ -57,8 +57,8 @@ A condensed view of the knobs you'll most often tune. For trainer-side paralleli | Knob | What it does | |---|---| -| `orchestrator.token_batch_size` | Tokens to train on per step. Size it as (target rollouts per step) x (average tokens per rollout). | -| `orchestrator.max_inflight_episodes` | Concurrent episodes kept in-flight. Tune together with `token_batch_size`: roughly `token_batch_size / (average tokens per rollout)`, higher to oversample ahead of the next batch. | +| `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-4.5-air/search.toml b/examples/advanced/glm-4.5-air/search.toml index 53c8af7d11..054799f928 100644 --- a/examples/advanced/glm-4.5-air/search.toml +++ b/examples/advanced/glm-4.5-air/search.toml @@ -68,7 +68,7 @@ type = "muon" # --- Orchestrator --- [orchestrator] -token_batch_size = 4_194_304 +batch_size = 256 group_size = 16 max_off_policy_steps = 32 max_inflight_episodes = 512 diff --git a/examples/advanced/glm-4.5-air/swe.toml b/examples/advanced/glm-4.5-air/swe.toml index f7603c3b68..696bec5442 100644 --- a/examples/advanced/glm-4.5-air/swe.toml +++ b/examples/advanced/glm-4.5-air/swe.toml @@ -79,7 +79,7 @@ type = "muon" # --- Orchestrator --- [orchestrator] -token_batch_size = 4_194_304 +batch_size = 256 group_size = 16 max_off_policy_steps = 32 max_inflight_episodes = 512 diff --git a/examples/advanced/glm-4.5-air/terminal.toml b/examples/advanced/glm-4.5-air/terminal.toml index bf2fca82b8..c3bf09adbb 100644 --- a/examples/advanced/glm-4.5-air/terminal.toml +++ b/examples/advanced/glm-4.5-air/terminal.toml @@ -69,7 +69,7 @@ type = "muon" # --- Orchestrator --- [orchestrator] -token_batch_size = 4_194_304 +batch_size = 256 group_size = 16 max_off_policy_steps = 32 max_inflight_episodes = 512 diff --git a/examples/advanced/glm-5.2/swe-llmd.toml b/examples/advanced/glm-5.2/swe-llmd.toml index acfa92d69f..34edcb0989 100644 --- a/examples/advanced/glm-5.2/swe-llmd.toml +++ b/examples/advanced/glm-5.2/swe-llmd.toml @@ -109,7 +109,7 @@ lr = 1e-6 weight_decay = 0.0 [orchestrator] -token_batch_size = 4_194_304 +batch_size = 256 group_size = 16 max_inflight_episodes = 2048 max_off_policy_steps = 8 # 16/32 can be better, 8 is more CPU stable diff --git a/examples/advanced/glm-5.2/swe.toml b/examples/advanced/glm-5.2/swe.toml index 65306bf5da..43ef6b8f5a 100644 --- a/examples/advanced/glm-5.2/swe.toml +++ b/examples/advanced/glm-5.2/swe.toml @@ -64,9 +64,9 @@ lr = 1e-6 weight_decay = 0.1 [orchestrator] -token_batch_size = 67_108_864 -max_inflight_episodes = 12_288 +batch_size = 4096 group_size = 16 +max_inflight_episodes = 12_288 max_off_policy_steps = 16 [orchestrator.model] diff --git a/examples/advanced/intellect-3.1/rl.toml b/examples/advanced/intellect-3.1/rl.toml index 6ac1c3e741..41bd2fe6ee 100644 --- a/examples/advanced/intellect-3.1/rl.toml +++ b/examples/advanced/intellect-3.1/rl.toml @@ -44,7 +44,7 @@ lr = 1e-6 weight_decay = 0.01 [orchestrator] -token_batch_size = 33_554_432 +batch_size = 2048 max_inflight_episodes = 4096 [[orchestrator.train.source]] diff --git a/examples/advanced/minimax-m2.5/swe.toml b/examples/advanced/minimax-m2.5/swe.toml index fb4c06b904..35a19d026a 100644 --- a/examples/advanced/minimax-m2.5/swe.toml +++ b/examples/advanced/minimax-m2.5/swe.toml @@ -44,7 +44,7 @@ lr = 1e-6 weight_decay = 0.01 [orchestrator] -token_batch_size = 33_554_432 +batch_size = 2048 max_inflight_episodes = 4096 max_off_policy_steps = 16 diff --git a/examples/advanced/nemotron-3-super/swe.toml b/examples/advanced/nemotron-3-super/swe.toml index 96b1853cdb..e7149b23cc 100644 --- a/examples/advanced/nemotron-3-super/swe.toml +++ b/examples/advanced/nemotron-3-super/swe.toml @@ -63,7 +63,7 @@ skip_optimizer = true type = "adamw" [orchestrator] -token_batch_size = 4_194_304 +batch_size = 256 group_size = 16 max_off_policy_steps = 32 max_inflight_episodes = 1536 diff --git a/examples/advanced/qwen3-30b-a3b/math.toml b/examples/advanced/qwen3-30b-a3b/math.toml index da9b0254c9..c8e83526e7 100644 --- a/examples/advanced/qwen3-30b-a3b/math.toml +++ b/examples/advanced/qwen3-30b-a3b/math.toml @@ -41,7 +41,7 @@ type = "adamw" lr = 1e-6 [orchestrator] -token_batch_size = 4_194_304 +batch_size = 512 max_inflight_episodes = 1024 max_off_policy_steps = 8 diff --git a/examples/advanced/qwen3-30b-a3b/swe.toml b/examples/advanced/qwen3-30b-a3b/swe.toml index 63a53d0148..547bc88242 100644 --- a/examples/advanced/qwen3-30b-a3b/swe.toml +++ b/examples/advanced/qwen3-30b-a3b/swe.toml @@ -42,7 +42,7 @@ type = "adamw" lr = 1e-6 [orchestrator] -token_batch_size = 8_388_608 +batch_size = 512 max_inflight_episodes = 1024 max_off_policy_steps = 16 diff --git a/examples/advanced/qwen3-30b-a3b/tool.toml b/examples/advanced/qwen3-30b-a3b/tool.toml index 74b9823a63..7732ff1a6a 100644 --- a/examples/advanced/qwen3-30b-a3b/tool.toml +++ b/examples/advanced/qwen3-30b-a3b/tool.toml @@ -31,8 +31,7 @@ freq = 1 [trainer.model.compile] [orchestrator] -token_batch_size = 2_097_152 -max_inflight_episodes = 512 +batch_size = 512 group_size = 16 max_off_policy_steps = 32 diff --git a/examples/basic/alphabet-sort/rl.toml b/examples/basic/alphabet-sort/rl.toml index f1962d8c1c..fade9b0c4a 100644 --- a/examples/basic/alphabet-sort/rl.toml +++ b/examples/basic/alphabet-sort/rl.toml @@ -25,8 +25,7 @@ alpha = 64 lr = 1e-5 [orchestrator] -token_batch_size = 262_144 -max_inflight_episodes = 256 +batch_size = 256 group_size = 16 [orchestrator.train.sampling] diff --git a/examples/basic/hendrycks-sanity/rl.toml b/examples/basic/hendrycks-sanity/rl.toml index 9bcef7a7d2..7c0b80ca63 100644 --- a/examples/basic/hendrycks-sanity/rl.toml +++ b/examples/basic/hendrycks-sanity/rl.toml @@ -12,8 +12,7 @@ num_infer_gpus = 4 name = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B" [orchestrator] -token_batch_size = 2_097_152 -max_inflight_episodes = 512 +batch_size = 512 group_size = 8 seq_len = 8192 diff --git a/examples/basic/reverse-text/rl.toml b/examples/basic/reverse-text/rl.toml index 4f8a92b6d7..184686ead1 100644 --- a/examples/basic/reverse-text/rl.toml +++ b/examples/basic/reverse-text/rl.toml @@ -9,7 +9,7 @@ project = "reverse-text" name = "reverse-text" [orchestrator] -token_batch_size = 16_384 +batch_size = 128 group_size = 16 [orchestrator.train.sampling] diff --git a/examples/basic/wiki-search/rl.toml b/examples/basic/wiki-search/rl.toml index 220a9a051a..936970f258 100644 --- a/examples/basic/wiki-search/rl.toml +++ b/examples/basic/wiki-search/rl.toml @@ -30,9 +30,9 @@ target_modules = [ ] [orchestrator] -token_batch_size = 1_048_576 -max_inflight_episodes = 1024 +batch_size = 512 group_size = 16 +max_inflight_episodes = 1024 [orchestrator.model.lora] name = "qwen3-4b-wiki-search" diff --git a/examples/basic/wordle/rl.toml b/examples/basic/wordle/rl.toml index 0cea3a65c1..84606b3a30 100644 --- a/examples/basic/wordle/rl.toml +++ b/examples/basic/wordle/rl.toml @@ -15,8 +15,7 @@ name = "wordle" name = "PrimeIntellect/Qwen3-1.7B-Wordle-SFT" [orchestrator] -token_batch_size = 1_048_576 -max_inflight_episodes = 512 +batch_size = 512 group_size = 16 [[orchestrator.train.source]] diff --git a/k8s/prime-rl/examples/reverse-text/orch.toml b/k8s/prime-rl/examples/reverse-text/orch.toml index 86962dff89..a4361ab8db 100644 --- a/k8s/prime-rl/examples/reverse-text/orch.toml +++ b/k8s/prime-rl/examples/reverse-text/orch.toml @@ -1,6 +1,6 @@ max_steps = 20 seq_len = 2048 -token_batch_size = 16_384 +batch_size = 128 group_size = 16 [model] 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 73845898ca..86e5d54a81 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py @@ -408,14 +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).""" - token_batch_size: int = Field(131_072, ge=1) - """Tokens to train on per step. A batch ships once the pending rollouts' trainer-bound payload reaches this many tokens. Size it as (target rollouts per step) x (average tokens per rollout); the default matches the old 128-rollout default at ~1k tokens per rollout.""" + batch_size: int = Field(128, ge=1) + """Rollouts to train on per step. Must be divisible by ``group_size``.""" count_zero_advantage_in_batch: bool = False - """Count zero-advantage rollouts toward ``token_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 tokens.""" + """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 = Field(128, ge=1) - """Maximum number of episodes kept in-flight — one episode is one agent run at a time, whatever the env's agents are. Tune together with ``token_batch_size``: roughly ``token_batch_size / (average tokens per rollout)``, higher to oversample ahead of the next batch.""" + 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. 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.""" @@ -523,6 +523,10 @@ def validate_renderer_auto_resolves(self): @model_validator(mode="after") def resolve_batching(self): + 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") 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 26a180fdcf..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=max(1, self.orchestrator.token_batch_size // self.orchestrator.seq_len), + batch_size=self.orchestrator.batch_size, ) trainer_bench_enabled = self.trainer.bench is not None diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index 2b3b315111..4fc776928d 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -6,7 +6,7 @@ - ``RolloutDispatcher`` schedules rollouts; emits ``Rollout`` (train/eval discriminated by ``kind``) on its queue. - ``TrainSink`` ingests train rollouts (tokenize → advantages → zero-advantage - drop) and returns a ``TrainBatch`` when the token budget is met. + 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 @@ -391,6 +391,7 @@ async def setup(self) -> None: else None ) + 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( @@ -761,7 +762,7 @@ def collect_pipeline_view(self) -> tuple[str, dict[str, float]]: # Train batch: finalized-group survivors only (0→target). Partial-group # arrivals are surfaced as a separate ``(+N buffered)`` addendum train_pct = train_batch / train_target if train_target else 0.0 - train_batch_part = f"Train batch {train_batch}/{train_target} tokens ({train_pct:.1%})" + train_batch_part = f"Train batch {train_batch}/{train_target} ({train_pct:.1%})" if multi_train: pairs = [(e.name, train_batch_by_env.get(e.name, 0)) for e in self.train_envs] train_batch_part += " (" + ", ".join(f"{n}={v}" for n, v in pairs) + ")" diff --git a/src/prime_rl/orchestrator/train_sink.py b/src/prime_rl/orchestrator/train_sink.py index 3982c22dae..5093635d8e 100644 --- a/src/prime_rl/orchestrator/train_sink.py +++ b/src/prime_rl/orchestrator/train_sink.py @@ -8,8 +8,8 @@ survivors to the env algorithm's ``finalize_group`` (advantages + per-sample wire stamping), annotates degeneration detections, and drops zero-advantage rollouts before they consume batch budget. -3. ``process_batch`` — assembles the trainer-bound ``TrainingSample`` list. - Returns a ``TrainBatch``. +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. @@ -43,19 +43,6 @@ ZERO_ADVANTAGE_STALL_WARN_GROUPS = 25 -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 - - class TrainSink: """Three-level train sink. Constructed once, fed via ``add(rollout)``.""" @@ -71,7 +58,7 @@ def __init__( self.tokenizer = tokenizer self.train_envs = train_envs self.mm_token_type_ids_mapping = mm_token_type_ids_mapping - self.token_batch_size = config.token_batch_size + 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) @@ -87,9 +74,6 @@ 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``, kept in sync on - # append/pop so the readiness check never re-sums per arrival. - self.pending_tokens: int = 0 # Consecutive finalized groups that contributed nothing to # ``pending_batch`` because every survivor was zero-advantage self.consecutive_zero_advantage_groups = 0 @@ -98,11 +82,11 @@ 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]: - """``(current, target)`` payload tokens for the train batch — counts - only ``pending_batch`` (survivors of finalized groups, queued for the + """``(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()``.""" - return self.pending_tokens, self.token_batch_size + return len(self.pending_batch), self.batch_size def buffered_count(self) -> int: """Episodes that have arrived but sit in not-yet-complete groups — @@ -110,11 +94,11 @@ def buffered_count(self) -> int: return sum(self.pending_group_episodes.values()) def pending_batch_by_env(self) -> dict[str, int]: - """Per-env payload-token breakdown of ``batch_progress()`` - (``pending_batch`` only); values sum to the aggregate.""" + """Per-env breakdown of ``batch_progress()`` (``pending_batch`` only); + values sum to the aggregate.""" counts: dict[str, int] = defaultdict(int) for r in self.pending_batch: - counts[r.env_name] += payload_tokens(r) + counts[r.env_name] += 1 return dict(counts) async def add(self, episode: list[Rollout]) -> TrainBatch | None: @@ -134,7 +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. - if self.pending_tokens >= self.token_batch_size: + if len(self.pending_batch) >= self.batch_size: return self.process_batch() return None @@ -212,13 +196,12 @@ async def process_group(self, group_id: uuid.UUID) -> None: r.is_filtered = r.filter_results["zero_advantage"] and not env.algorithm.trains_on_zero_advantage if r.is_filtered: num_zero_advantage += 1 - # Opt-in: a dropped rollout still consumes batch budget, so the + # Opt-in: a dropped rollout still occupies a batch slot, so the # per-step sampling effort stays fixed while the trained-on - # token count varies with the zero-advantage rate. + # sample count varies with the zero-advantage rate. if not self.count_zero_advantage_in_batch: continue self.pending_batch.append(r) - self.pending_tokens += payload_tokens(r) appended += 1 if appended: @@ -241,19 +224,11 @@ async def process_group(self, group_id: uuid.UUID) -> None: ) def process_batch(self) -> TrainBatch: - """Pop the cohort whose payload tokens reach ``token_batch_size`` off - ``pending_batch`` and assemble the trainer-bound ``TrainingSample`` - list. Overflow stays for the next batch.""" - 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 + """Pop a ``batch_size`` cohort off ``pending_batch`` and assemble the + trainer-bound ``TrainingSample`` list. Overflow stays for the next + batch.""" + 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. Zero-advantage rollouts kept in the diff --git a/src/prime_rl/utils/monitor/prime.py b/src/prime_rl/utils/monitor/prime.py index 4450cd124f..8d57133321 100644 --- a/src/prime_rl/utils/monitor/prime.py +++ b/src/prime_rl/utils/monitor/prime.py @@ -164,6 +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: + 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/test_configs.py b/tests/unit/test_configs.py index 6c9526fe11..abeac169e8 100644 --- a/tests/unit/test_configs.py +++ b/tests/unit/test_configs.py @@ -557,7 +557,7 @@ def test_shared_output_dir_propagates_through_cli(tmp_path): "seq_len": 128, "model": {"name": "Qwen/Qwen3-0.6B"}, "trainer": {}, - "orchestrator": {"token_batch_size": 2048, "group_size": 1}, + "orchestrator": {"batch_size": 16, "group_size": 1}, "inference": {}, }, ) From fc70759ca0df94e019d7ddec7fe397bcaca464b2 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Wed, 12 Aug 2026 19:56:47 +0000 Subject: [PATCH 4/4] Drop the filters/ metric key segment and AIPO README mention The check verdicts log as {scope}/{subset}/{agent}//mean, next to the other per-trace verdicts (is_trainable, is_filtered). Co-Authored-By: Claude Fable 5 --- README.md | 2 +- docs/algorithms.md | 6 +++--- skills/training/monitor-run/SKILL.md | 2 +- src/prime_rl/orchestrator/metrics.py | 2 +- tests/unit/orchestrator/test_metrics.py | 4 ++-- 5 files changed, 8 insertions(+), 8 deletions(-) diff --git a/README.md b/README.md index 0f843f384f..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 plugins, rollout checks, 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/docs/algorithms.md b/docs/algorithms.md index c38b19f194..02805bb79c 100644 --- a/docs/algorithms.md +++ b/docs/algorithms.md @@ -462,13 +462,13 @@ Between scoring and training the sink runs three hardcoded checks on every train | Check | Effect | |---|---| -| `gibberish` | Detects rare tokens generated at high entropy — usually a sign of degenerate output. Monitor-only: tracked in metrics (`filters/gibberish`), never dropped. | -| `repetition` | Detects long high-confidence loops. Monitor-only: tracked in metrics (`filters/repetition`), never dropped. | +| `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. | 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. -Dropped rollouts still appear in W&B distributions and metrics (`is_filtered`, `filters/zero_advantage`), just not in the trainer batch. +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/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/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/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():