From e482ab32a86144a23187ad9327a36478cedc57a9 Mon Sep 17 00:00:00 2001 From: Abhishek Enaguthi Date: Sun, 12 Jul 2026 22:15:19 -0700 Subject: [PATCH 1/2] feat(orchestrator): expose target_lag config (#2812) Promote the hardcoded dispatch-gate lag to orchestrator.target_lag (default 1) with docs clarifying it is unrelated to max_off_policy_steps. --- docs/training.md | 3 ++- .../src/prime_rl/configs/orchestrator.py | 11 ++++++++++- src/prime_rl/orchestrator/orchestrator.py | 7 ++++--- tests/unit/orchestrator/test_off_policy_lag.py | 16 ++++++++++++++++ 4 files changed, 32 insertions(+), 5 deletions(-) create mode 100644 tests/unit/orchestrator/test_off_policy_lag.py diff --git a/docs/training.md b/docs/training.md index 0372fd68c1..c74e9aa0de 100644 --- a/docs/training.md +++ b/docs/training.md @@ -58,7 +58,8 @@ A condensed view of the knobs you'll most often tune. For trainer-side paralleli |---|---| | `orchestrator.batch_size` | Tasks per trainer step. | | `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.max_off_policy_steps` | Per-rollout staleness cap (default 8). Lag = `(trainer_step - 1) - policy_version_at_generation`, including rollouts waiting in the train sink. Bump for throughput on long agentic rollouts; lower for tighter on-policyness. Watch `errored_rollouts` and `mismatch_kl/all/mean` when tuning. | +| `orchestrator.target_lag` | Dispatch gate: max batches the orchestrator may ship ahead of `policy.version` (default 1). Pauses rollout scheduling until the trainer publishes the next checkpoint — **not** the same knob as `max_off_policy_steps`. Watch `time/wait_for_policy` when tuning. | | `[orchestrator.algo]` | Training algorithm — its `type` names it (`grpo` default, `max_rl`, `opd`, `opsd`, `sft`, `echo`). See [Algorithms](#algorithms). | | `[[orchestrator.train.env]]` | Training environments. List multiple tables for multi-env training; weight them via `ratio`. See [Configuration § Environments](configuration.md#environments-orchestratortrainenv). | | `[[orchestrator.eval.env]]` + `orchestrator.eval.interval` | Eval environments and cadence (default every 100 steps). | 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 2c11a4ddb6..5e96e3b92e 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py @@ -529,7 +529,16 @@ class OrchestratorConfig(BaseConfig): """Maximum training steps. If None, runs indefinitely.""" max_off_policy_steps: int = Field(8, ge=0) - """Maximum policies allowed to generate a single rollout. Rollouts generated more than ``max_off_policy_steps`` ahead of training are discarded. Higher values yield better throughput at the cost of off-policy noise.""" + """Maximum trainer–policy lag allowed per rollout before it is discarded. + + Lag is measured as ``(trainer_step - 1) - policy_version_at_generation``, + covering both in-flight rollouts and rollouts waiting in the train sink. + This is **not** the orchestrator dispatch gate (see ``target_lag``).""" + + target_lag: int = Field(1, ge=1) + """Maximum batches the orchestrator may ship ahead of ``policy.version``. + Default ``1`` preserves the standard async trainer–inference pipeline. + Unrelated to ``max_off_policy_steps`` (per-rollout staleness).""" bench: bool = False """Benchmark mode. Sets ``max_steps`` to 5 and disables W&B.""" diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index c17ed9e2c7..fa132cd910 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -97,8 +97,9 @@ # Maximum batches the orchestrator may run ahead of the trainer. The # dispatcher is paused via ``update_dispatch_gate`` once this is exceeded; -# resumed when the watcher advances ``policy.version``. -TARGET_LAG = 1 +# resumed when the watcher advances ``policy.version``. Configurable via +# ``orchestrator.target_lag`` (default 1). Unrelated to ``max_off_policy_steps``. +DEFAULT_TARGET_LAG = 1 class Orchestrator: @@ -825,7 +826,7 @@ def update_dispatch_gate(self) -> None: ) gate = self.dispatcher.dispatch_allowed was_set = gate.is_set() - if lead > TARGET_LAG and not building_final_batch_nccl: + if lead > self.config.target_lag and not building_final_batch_nccl: if was_set: get_logger().info( "Pausing dispatcher to prevent orchestrator from racing from trainer. Waiting for new policy..." diff --git a/tests/unit/orchestrator/test_off_policy_lag.py b/tests/unit/orchestrator/test_off_policy_lag.py new file mode 100644 index 0000000000..b8abbcbf6a --- /dev/null +++ b/tests/unit/orchestrator/test_off_policy_lag.py @@ -0,0 +1,16 @@ +"""Tests for target_lag orchestrator config (#2812).""" + +import pytest +from pydantic import ValidationError + +from prime_rl.configs.orchestrator import OrchestratorConfig + + +def test_target_lag_defaults_to_one(): + cfg = OrchestratorConfig() + assert cfg.target_lag == 1 + + +def test_target_lag_rejects_zero(): + with pytest.raises(ValidationError): + OrchestratorConfig(target_lag=0) From 1cdd0813ed897347fd5f4261ab830ec5d8fde3fa Mon Sep 17 00:00:00 2001 From: Abhishek Enaguthi Date: Sun, 12 Jul 2026 22:21:42 -0700 Subject: [PATCH 2/2] fix: remove unused DEFAULT_TARGET_LAG constant Gate reads orchestrator.target_lag from config; module constant was dead. --- src/prime_rl/orchestrator/orchestrator.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index fa132cd910..5be3c7b8bc 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -95,11 +95,8 @@ # dataset; fail loudly instead of spinning MAX_CONSECUTIVE_EMPTY_BATCHES = 10 -# Maximum batches the orchestrator may run ahead of the trainer. The -# dispatcher is paused via ``update_dispatch_gate`` once this is exceeded; -# resumed when the watcher advances ``policy.version``. Configurable via -# ``orchestrator.target_lag`` (default 1). Unrelated to ``max_off_policy_steps``. -DEFAULT_TARGET_LAG = 1 +# Dispatch gate lag is configured via ``orchestrator.target_lag`` (default 1). +# Unrelated to ``max_off_policy_steps`` (per-rollout staleness). class Orchestrator: