From 80089622089aa7d15731dac7577bda7589f0f69b Mon Sep 17 00:00:00 2001 From: Timothy Kostolansky <39891386+tim0120@users.noreply.github.com> Date: Mon, 11 May 2026 17:46:55 +0000 Subject: [PATCH 01/47] feat: add student eval inference pool for SFT hard distillation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When teacher_rollout_model is configured for SFT distillation, the orchestrator now supports a separate student inference pool for online evals and weight sync. Previously, configuring [inference] alongside teacher_rollout_model was forbidden — evals either ran on the teacher or were skipped entirely. Changes: - Relax RLConfig validator to allow [inference] + teacher_rollout_model - Create eval_inference_pool from config.client when teacher_rollout_model is set, pointing at the student vLLM server - Route eval calls and weight updates to the student pool - Add eval_inference_pool param to Scheduler for weight sync targeting - All existing RL/soft-distill paths are unchanged (eval_inference_pool defaults to inference_pool when no teacher_rollout_model is configured) Co-Authored-By: Claude Opus 4.6 (1M context) --- .../src/prime_rl/configs/rl.py | 15 ------- src/prime_rl/orchestrator/orchestrator.py | 41 +++++++++++++++---- src/prime_rl/orchestrator/scheduler.py | 10 +++-- 3 files changed, 41 insertions(+), 25 deletions(-) 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 1990b46876..3e689d9b0b 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/rl.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/rl.py @@ -433,21 +433,6 @@ def validate_teacher_model(self): ) return self - @model_validator(mode="after") - def validate_external_rollout_inference(self): - """Forbid configuring a local inference server when rollouts come from an external teacher. - - Orchestrator-only invariants (``use_sft_loss`` paired with ``teacher_rollout_model``, - and ``use_token_client`` coupling) live on ``OrchestratorConfig`` so the hosted - orchestrator entrypoint also enforces them. - """ - if self.orchestrator.teacher_rollout_model is not None and self.inference is not None: - raise ValueError( - "inference must be omitted when orchestrator.teacher_rollout_model is configured. " - "External rollout mode does not use the local inference server." - ) - return self - ### Auto-setup and validate shared configs @model_validator(mode="after") diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index bc1128ebc7..ff7f48026b 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -156,6 +156,24 @@ async def orchestrate(config: OrchestratorConfig): logger=logger, ) + # When using an external teacher for rollouts, set up a separate student + # inference pool for online evals and weight sync. The student URL comes + # from the orchestrator's default client config (config.client). + if config.teacher_rollout_model is not None: + student_model_name = config.model.name + logger.info( + f"Initializing student eval inference pool (base_url={', '.join(config.client.base_url)}, " + f"model={student_model_name})" + ) + eval_inference_pool = await setup_inference_pool( + config.client, + model_name=student_model_name, + eval_client_type="openai_chat_completions", + ) + enable_policy_updates = True + else: + eval_inference_pool = inference_pool + # Setup monitor (may register the run and set RUN_ID in the environment) logger.info(f"Initializing monitor (wandb={config.wandb}, prime_monitor={config.prime_monitor})") monitor = setup_monitor( @@ -236,6 +254,7 @@ async def orchestrate(config: OrchestratorConfig): train_envs=train_envs, buffer=buffer, inference_pool=inference_pool, + eval_inference_pool=eval_inference_pool, max_inflight_rollouts=config.max_inflight_rollouts, max_async_level=config.max_async_level, max_off_policy_steps=config.max_off_policy_steps, @@ -254,9 +273,15 @@ async def orchestrate(config: OrchestratorConfig): # Check health of the inference pool logger.info("Waiting for inference pool to be ready") await inference_pool.wait_for_ready(rollout_model_name) - logger.success("Inference pool ready") + # Check health of student eval inference pool if separate from rollout pool + if eval_inference_pool is not inference_pool: + student_model_name = config.model.name + logger.info("Waiting for student eval inference pool to be ready") + await eval_inference_pool.wait_for_ready(student_model_name) + logger.success("Student eval inference pool ready") + # Start inference metrics collector (requires W&B) inference_metrics_collector = None if config.wandb is not None and config.collect_inference_metrics: @@ -274,7 +299,7 @@ async def orchestrate(config: OrchestratorConfig): logger.info(f"Initializing weight broadcast ({config.weight_broadcast})") if config.weight_broadcast.type == "nccl": await init_nccl_broadcast( - inference_pool.admin_clients, + eval_inference_pool.admin_clients, config.weight_broadcast.host, config.weight_broadcast.port, config.weight_broadcast.timeout, @@ -316,7 +341,7 @@ async def orchestrate(config: OrchestratorConfig): config.output_dir, scheduler.ckpt_step, check_exists=check_exists, wait_timeout=wait_timeout ) lora_name = config.model.lora.name if config.model.lora else None - await inference_pool.update_weights(weights_path, lora_name=lora_name, step=scheduler.ckpt_step) + await eval_inference_pool.update_weights(weights_path, lora_name=lora_name, step=scheduler.ckpt_step) else: logger.info("Training from scratch") @@ -390,8 +415,8 @@ async def orchestrate(config: OrchestratorConfig): eval_results = await asyncio.gather( *[ eval_env.evaluate( - model_name=scheduler.model_name, - get_client=inference_pool.get_eval_client, + model_name=eval_inference_pool.model_name, + get_client=eval_inference_pool.get_eval_client, ckpt_step=ckpt_step, step=progress.step, cache_salt=str(ckpt_step), @@ -809,8 +834,8 @@ def compute_solve_rates(df): eval_results = await asyncio.gather( *[ eval_env.evaluate( - model_name=scheduler.model_name, - get_client=inference_pool.get_eval_client, + model_name=eval_inference_pool.model_name, + get_client=eval_inference_pool.get_eval_client, ckpt_step=ckpt_step, step=progress.step, cache_salt=str(ckpt_step), @@ -847,6 +872,8 @@ async def _graceful_shutdown() -> None: if inference_metrics_collector is not None: await inference_metrics_collector.stop() await inference_pool.stop() + if eval_inference_pool is not inference_pool: + await eval_inference_pool.stop() if teacher_inference_pool is not None: await teacher_inference_pool.stop() event_loop_lag_monitor_task.cancel() diff --git a/src/prime_rl/orchestrator/scheduler.py b/src/prime_rl/orchestrator/scheduler.py index a8427b69f7..d3527cff12 100644 --- a/src/prime_rl/orchestrator/scheduler.py +++ b/src/prime_rl/orchestrator/scheduler.py @@ -68,6 +68,7 @@ def __init__( tasks_per_minute: int | None, enable_policy_updates: bool = True, lora_name: str | None = None, + eval_inference_pool: InferencePool | None = None, ): self.logger = get_logger() if tasks_per_minute is not None: @@ -89,8 +90,11 @@ def __init__( self.model_name = self.config.model.name self.json_logging = config.log.json_logging - # Inference pool - used for admin operations (adapter sync) and metrics + # Inference pool - used for rollout client selection and metrics self.inference_pool = inference_pool + # Eval inference pool - receives weight updates and serves evals. + # Defaults to inference_pool (standard RL where one pool does both). + self.eval_inference_pool = eval_inference_pool or inference_pool group_scoring_envs = [env.name for env in train_envs if env.requires_group_scoring] if group_scoring_envs: @@ -303,14 +307,14 @@ async def _apply_policy_update(self, next_ckpt_step: int) -> None: update_weights_start_time = time.perf_counter() weights_path = get_step_path(get_broadcast_dir(self.config.output_dir), next_ckpt_step) - await self.inference_pool.update_weights(weights_path, lora_name=self.lora_name, step=next_ckpt_step) + await self.eval_inference_pool.update_weights(weights_path, lora_name=self.lora_name, step=next_ckpt_step) self.update_weights_time = time.perf_counter() - update_weights_start_time self.logger.debug(f"Updated weights to step {next_ckpt_step} in {self.update_weights_time:.2f}s") self.ckpt_step = next_ckpt_step if self.lora_name is not None: self.model_name = self.lora_name - self.inference_pool.update_model_name(self.model_name) + self.eval_inference_pool.update_model_name(self.model_name) self.checkpoint_ready.set() await self._update_off_policy() From 0e14a44685a2b200099c30726fee96685dc62a76 Mon Sep 17 00:00:00 2001 From: Timothy Kostolansky <39891386+tim0120@users.noreply.github.com> Date: Tue, 12 May 2026 05:22:15 +0000 Subject: [PATCH 02/47] fix: keep hard distill student eval pool optional --- .../src/prime_rl/configs/orchestrator.py | 11 +++++++ .../src/prime_rl/configs/rl.py | 11 +++++++ skills/config/SKILL.md | 4 ++- src/prime_rl/orchestrator/orchestrator.py | 7 ++--- src/prime_rl/orchestrator/scheduler.py | 5 +-- tests/unit/orchestrator/test_scheduler.py | 31 +++++++++++++++++++ tests/unit/test_configs.py | 20 ++++++++++++ 7 files changed, 82 insertions(+), 7 deletions(-) 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 3911816227..61a905f9ac 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py @@ -950,6 +950,17 @@ class OrchestratorConfig(BaseConfig): ), ] = None + use_student_eval_inference_pool: Annotated[ + bool, + Field( + description=( + "When teacher_rollout_model is set, use orchestrator.client as a separate student inference pool " + "for online evals and weight updates. The RL config enables this automatically when [inference] " + "is configured." + ), + ), + ] = False + # When True, trainer uses SFT loss instead of RL loss (per-run override for hosted multi-tenant training) use_sft_loss: Annotated[ bool, 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 3e689d9b0b..c6788d4514 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/rl.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/rl.py @@ -900,6 +900,17 @@ def auto_setup_dp_rank_count(self): ) return self + @model_validator(mode="after") + def auto_setup_student_eval_inference_pool(self): + """Enable student eval inference for hard distill only when [inference] exists.""" + if ( + self.orchestrator.teacher_rollout_model is not None + and self.inference is not None + and "use_student_eval_inference_pool" not in self.orchestrator.model_fields_set + ): + self.orchestrator.use_student_eval_inference_pool = True + return self + @model_validator(mode="after") def auto_setup_teacher_inference(self): """Auto-configure teacher inference server and orchestrator teacher_model client.""" diff --git a/skills/config/SKILL.md b/skills/config/SKILL.md index e8dc13216c..d18e0ee032 100644 --- a/skills/config/SKILL.md +++ b/skills/config/SKILL.md @@ -155,7 +155,9 @@ If you wish to configure values of the default variant, you don't need to set th ### SFT hard distill override -For hosted multi-tenant runs where the trainer image's `trainer.loss.type` is fixed, the orchestrator exposes a per-run override that forces SFT loss on every micro-batch without rebuilding the trainer. Set `orchestrator.use_sft_loss = true` alongside `orchestrator.teacher_rollout_model`; both must be configured together (the orchestrator validator enforces this). The orchestrator stamps each `TrainingSample.sft_loss = True`, which the trainer's `compute_loss` honors by dispatching to `sft_loss_fn` per batch — independent of the trainer's configured default loss. +For hosted multi-tenant runs where the trainer image's `trainer.loss.type` is fixed, the orchestrator exposes a per-run override that forces SFT loss on every micro-batch without rebuilding the trainer. Set `orchestrator.use_sft_loss = true` alongside `orchestrator.teacher_rollout_model`; both must be configured together (the orchestrator validator enforces this). The orchestrator stamps each `TrainingSample.sft_loss = True`, which the trainer's `compute_loss` honors by dispatching to `sft_loss_fn` per batch, independent of the trainer's configured default loss. + +When hard distill also needs online evals or policy weight sync against the student model, configure `[inference]` in the RL config. `RLConfig` then enables `orchestrator.use_student_eval_inference_pool` automatically, and `orchestrator.client` points evals and weight updates at the student inference server while rollouts keep using `orchestrator.teacher_rollout_model`. If `[inference]` is omitted, hard distill keeps the teacher-only rollout behavior and does not wait for a local student inference server. ### Model fields diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index ff7f48026b..5dd67da423 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -156,10 +156,9 @@ async def orchestrate(config: OrchestratorConfig): logger=logger, ) - # When using an external teacher for rollouts, set up a separate student - # inference pool for online evals and weight sync. The student URL comes - # from the orchestrator's default client config (config.client). - if config.teacher_rollout_model is not None: + # When configured for external teacher rollouts, a separate student pool can + # serve online evals and receive policy weight updates. + if config.teacher_rollout_model is not None and config.use_student_eval_inference_pool: student_model_name = config.model.name logger.info( f"Initializing student eval inference pool (base_url={', '.join(config.client.base_url)}, " diff --git a/src/prime_rl/orchestrator/scheduler.py b/src/prime_rl/orchestrator/scheduler.py index d3527cff12..ab5d8e0b3f 100644 --- a/src/prime_rl/orchestrator/scheduler.py +++ b/src/prime_rl/orchestrator/scheduler.py @@ -313,8 +313,9 @@ async def _apply_policy_update(self, next_ckpt_step: int) -> None: self.ckpt_step = next_ckpt_step if self.lora_name is not None: - self.model_name = self.lora_name - self.eval_inference_pool.update_model_name(self.model_name) + self.eval_inference_pool.update_model_name(self.lora_name) + if self.eval_inference_pool is self.inference_pool: + self.model_name = self.lora_name self.checkpoint_ready.set() await self._update_off_policy() diff --git a/tests/unit/orchestrator/test_scheduler.py b/tests/unit/orchestrator/test_scheduler.py index 9e73b5207d..86129e9606 100644 --- a/tests/unit/orchestrator/test_scheduler.py +++ b/tests/unit/orchestrator/test_scheduler.py @@ -106,6 +106,7 @@ async def update_weights(weight_dir, lora_name=None, step=0) -> None: update_weights=update_weights, update_model_name=MagicMock(), ) + scheduler.eval_inference_pool = scheduler.inference_pool scheduler._update_off_policy = AsyncMock() with ( @@ -146,6 +147,7 @@ async def update_weights(weight_dir, lora_name=None, step=0) -> None: update_weights=update_weights, update_model_name=MagicMock(), ) + scheduler.eval_inference_pool = scheduler.inference_pool scheduler._update_off_policy = AsyncMock() with ( @@ -174,3 +176,32 @@ def test_client_identity_distinguishes_base_url_and_dp_rank(): ) assert Scheduler._client_identity(client_a) != Scheduler._client_identity(client_b) + + +def test_lora_policy_update_keeps_teacher_rollout_model_name_with_separate_eval_pool(): + async def run() -> None: + scheduler = make_scheduler() + scheduler.model_name = "teacher-model" + scheduler.lora_name = "student-lora" + + scheduler.inference_pool = SimpleNamespace( + update_model_name=MagicMock(), + ) + scheduler.eval_inference_pool = SimpleNamespace( + update_weights=AsyncMock(), + update_model_name=MagicMock(), + ) + scheduler._update_off_policy = AsyncMock() + + with ( + patch("prime_rl.orchestrator.scheduler.get_latest_ckpt_step", return_value=8), + patch("prime_rl.orchestrator.scheduler.wait_for_path", new=AsyncMock()), + ): + await scheduler.maybe_update_policy() + + scheduler.eval_inference_pool.update_weights.assert_awaited_once() + scheduler.eval_inference_pool.update_model_name.assert_called_once_with("student-lora") + scheduler.inference_pool.update_model_name.assert_not_called() + assert scheduler.model_name == "teacher-model" + + asyncio.run(run()) diff --git a/tests/unit/test_configs.py b/tests/unit/test_configs.py index ffcc18f270..0f2f095063 100644 --- a/tests/unit/test_configs.py +++ b/tests/unit/test_configs.py @@ -159,3 +159,23 @@ def test_removed_fused_lm_head_chunk_size_field_is_rejected(): def test_selective_activation_checkpointing_requires_custom_impl(): with pytest.raises(ValidationError, match="Selective activation checkpointing requires model.impl='custom'"): TrainerModelConfig.model_validate({"impl": "hf", "ac": {"mode": "selective"}}) + + +def test_teacher_rollout_student_eval_pool_requires_inference_config(): + base_config = { + "trainer": {}, + "orchestrator": { + "use_sft_loss": True, + "use_token_client": False, + "teacher_rollout_model": { + "client": {"base_url": ["http://teacher.example/v1"]}, + "model": {"name": "teacher-model"}, + }, + }, + } + + config = RLConfig.model_validate(base_config) + assert not config.orchestrator.use_student_eval_inference_pool + + config = RLConfig.model_validate({**base_config, "inference": {}}) + assert config.orchestrator.use_student_eval_inference_pool From 703b991e6433ec3b3e9956f93ea0c0da2e4c5393 Mon Sep 17 00:00:00 2001 From: Timothy Kostolansky <39891386+tim0120@users.noreply.github.com> Date: Tue, 12 May 2026 06:53:44 +0000 Subject: [PATCH 03/47] fix: scope teacher rollout routing --- .../src/prime_rl/configs/orchestrator.py | 20 ++--- .../src/prime_rl/configs/rl.py | 7 +- skills/config/SKILL.md | 2 +- src/prime_rl/orchestrator/orchestrator.py | 75 +++++++++---------- src/prime_rl/orchestrator/scheduler.py | 33 ++++---- src/prime_rl/orchestrator/utils.py | 5 +- tests/unit/orchestrator/test_scheduler.py | 59 ++++++++++++--- tests/unit/test_configs.py | 6 +- 8 files changed, 121 insertions(+), 86 deletions(-) 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 61a905f9ac..f70d90756f 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py @@ -950,27 +950,27 @@ class OrchestratorConfig(BaseConfig): ), ] = None - use_student_eval_inference_pool: Annotated[ + # When True, trainer uses SFT loss instead of RL loss (per-run override for hosted multi-tenant training) + use_sft_loss: Annotated[ bool, Field( description=( - "When teacher_rollout_model is set, use orchestrator.client as a separate student inference pool " - "for online evals and weight updates. The RL config enables this automatically when [inference] " - "is configured." + "When True, use SFT masked NLL loss instead of the trainer's configured RL loss. " + "Requires a teacher_rollout_model to be configured." ), ), ] = False - # When True, trainer uses SFT loss instead of RL loss (per-run override for hosted multi-tenant training) - use_sft_loss: Annotated[ - bool, + enable_policy_updates: Annotated[ + bool | None, Field( description=( - "When True, use SFT masked NLL loss instead of the trainer's configured RL loss. " - "Requires a teacher_rollout_model to be configured." + "Whether the orchestrator should receive policy weight updates. " + "Defaults to False for teacher_rollout_model runs and True otherwise; the RL config enables it " + "for teacher_rollout_model runs when [inference] is configured." ), ), - ] = False + ] = None # The evaluation configuration eval: EvalConfig | None = None 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 c6788d4514..6d62924832 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/rl.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/rl.py @@ -901,14 +901,13 @@ def auto_setup_dp_rank_count(self): return self @model_validator(mode="after") - def auto_setup_student_eval_inference_pool(self): - """Enable student eval inference for hard distill only when [inference] exists.""" + def auto_setup_teacher_rollout_policy_updates(self): if ( self.orchestrator.teacher_rollout_model is not None and self.inference is not None - and "use_student_eval_inference_pool" not in self.orchestrator.model_fields_set + and "enable_policy_updates" not in self.orchestrator.model_fields_set ): - self.orchestrator.use_student_eval_inference_pool = True + self.orchestrator.enable_policy_updates = True return self @model_validator(mode="after") diff --git a/skills/config/SKILL.md b/skills/config/SKILL.md index d18e0ee032..6f173f503b 100644 --- a/skills/config/SKILL.md +++ b/skills/config/SKILL.md @@ -157,7 +157,7 @@ If you wish to configure values of the default variant, you don't need to set th For hosted multi-tenant runs where the trainer image's `trainer.loss.type` is fixed, the orchestrator exposes a per-run override that forces SFT loss on every micro-batch without rebuilding the trainer. Set `orchestrator.use_sft_loss = true` alongside `orchestrator.teacher_rollout_model`; both must be configured together (the orchestrator validator enforces this). The orchestrator stamps each `TrainingSample.sft_loss = True`, which the trainer's `compute_loss` honors by dispatching to `sft_loss_fn` per batch, independent of the trainer's configured default loss. -When hard distill also needs online evals or policy weight sync against the student model, configure `[inference]` in the RL config. `RLConfig` then enables `orchestrator.use_student_eval_inference_pool` automatically, and `orchestrator.client` points evals and weight updates at the student inference server while rollouts keep using `orchestrator.teacher_rollout_model`. If `[inference]` is omitted, hard distill keeps the teacher-only rollout behavior and does not wait for a local student inference server. +When hard distill also needs online evals or policy weight sync against the student model, configure `[inference]` in the RL config. `RLConfig` then points `orchestrator.client` at the standard student inference server for evals and weight updates, while `orchestrator.teacher_rollout_model` is scoped to training rollout generation only. If `[inference]` is omitted, hard distill keeps the teacher-only rollout behavior and does not wait for a local student inference server. ### Model fields diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index 5dd67da423..7d3dd07e32 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -56,6 +56,7 @@ from prime_rl.trainer.model import setup_tokenizer from prime_rl.utils.client import ( init_nccl_broadcast, + setup_clients, setup_inference_pool, ) from prime_rl.utils.config import cli @@ -148,30 +149,30 @@ async def orchestrate(config: OrchestratorConfig): config.model.name, trust_remote_code=config.model.trust_remote_code, use_fast=True ) - renderer, inference_pool = await setup_rollout_inference_pool( - config=config, - rollout_client_config=rollout_client_config, - rollout_model_name=rollout_model_name, - tokenizer=tokenizer, - logger=logger, - ) - - # When configured for external teacher rollouts, a separate student pool can - # serve online evals and receive policy weight updates. - if config.teacher_rollout_model is not None and config.use_student_eval_inference_pool: - student_model_name = config.model.name - logger.info( - f"Initializing student eval inference pool (base_url={', '.join(config.client.base_url)}, " - f"model={student_model_name})" + teacher_rollout_clients = None + teacher_rollout_model_name = None + use_teacher_rollout_override = config.teacher_rollout_model is not None and enable_policy_updates + if use_teacher_rollout_override: + logger.info("Using external rollout model (MITO) without renderer client") + teacher_rollout_clients = setup_clients( + rollout_client_config, + client_type="openai_chat_completions", ) - eval_inference_pool = await setup_inference_pool( + teacher_rollout_model_name = rollout_model_name + renderer = None + inference_pool = await setup_inference_pool( config.client, - model_name=student_model_name, + model_name=config.model.name, eval_client_type="openai_chat_completions", ) - enable_policy_updates = True else: - eval_inference_pool = inference_pool + renderer, inference_pool = await setup_rollout_inference_pool( + config=config, + rollout_client_config=rollout_client_config, + rollout_model_name=rollout_model_name, + tokenizer=tokenizer, + logger=logger, + ) # Setup monitor (may register the run and set RUN_ID in the environment) logger.info(f"Initializing monitor (wandb={config.wandb}, prime_monitor={config.prime_monitor})") @@ -253,7 +254,8 @@ async def orchestrate(config: OrchestratorConfig): train_envs=train_envs, buffer=buffer, inference_pool=inference_pool, - eval_inference_pool=eval_inference_pool, + teacher_rollout_clients=teacher_rollout_clients, + teacher_rollout_model_name=teacher_rollout_model_name, max_inflight_rollouts=config.max_inflight_rollouts, max_async_level=config.max_async_level, max_off_policy_steps=config.max_off_policy_steps, @@ -263,24 +265,14 @@ async def orchestrate(config: OrchestratorConfig): lora_name=config.model.lora.name if config.model.lora else None, config=config, ) - scheduler.model_name = rollout_model_name - - if checkpoint_step is not None and config.model.lora is not None and enable_policy_updates: - assert config.model.lora.name is not None - scheduler.model_name = config.model.lora.name + scheduler.model_name = config.model.name if use_teacher_rollout_override else rollout_model_name # Check health of the inference pool logger.info("Waiting for inference pool to be ready") - await inference_pool.wait_for_ready(rollout_model_name) + inference_model_name = config.model.name if use_teacher_rollout_override else rollout_model_name + await inference_pool.wait_for_ready(inference_model_name) logger.success("Inference pool ready") - # Check health of student eval inference pool if separate from rollout pool - if eval_inference_pool is not inference_pool: - student_model_name = config.model.name - logger.info("Waiting for student eval inference pool to be ready") - await eval_inference_pool.wait_for_ready(student_model_name) - logger.success("Student eval inference pool ready") - # Start inference metrics collector (requires W&B) inference_metrics_collector = None if config.wandb is not None and config.collect_inference_metrics: @@ -298,7 +290,7 @@ async def orchestrate(config: OrchestratorConfig): logger.info(f"Initializing weight broadcast ({config.weight_broadcast})") if config.weight_broadcast.type == "nccl": await init_nccl_broadcast( - eval_inference_pool.admin_clients, + inference_pool.admin_clients, config.weight_broadcast.host, config.weight_broadcast.port, config.weight_broadcast.timeout, @@ -340,7 +332,10 @@ async def orchestrate(config: OrchestratorConfig): config.output_dir, scheduler.ckpt_step, check_exists=check_exists, wait_timeout=wait_timeout ) lora_name = config.model.lora.name if config.model.lora else None - await eval_inference_pool.update_weights(weights_path, lora_name=lora_name, step=scheduler.ckpt_step) + await inference_pool.update_weights(weights_path, lora_name=lora_name, step=scheduler.ckpt_step) + if lora_name is not None: + inference_pool.update_model_name(lora_name) + scheduler.model_name = lora_name else: logger.info("Training from scratch") @@ -414,8 +409,8 @@ async def orchestrate(config: OrchestratorConfig): eval_results = await asyncio.gather( *[ eval_env.evaluate( - model_name=eval_inference_pool.model_name, - get_client=eval_inference_pool.get_eval_client, + model_name=inference_pool.model_name, + get_client=inference_pool.get_eval_client, ckpt_step=ckpt_step, step=progress.step, cache_salt=str(ckpt_step), @@ -833,8 +828,8 @@ def compute_solve_rates(df): eval_results = await asyncio.gather( *[ eval_env.evaluate( - model_name=eval_inference_pool.model_name, - get_client=eval_inference_pool.get_eval_client, + model_name=inference_pool.model_name, + get_client=inference_pool.get_eval_client, ckpt_step=ckpt_step, step=progress.step, cache_salt=str(ckpt_step), @@ -871,8 +866,6 @@ async def _graceful_shutdown() -> None: if inference_metrics_collector is not None: await inference_metrics_collector.stop() await inference_pool.stop() - if eval_inference_pool is not inference_pool: - await eval_inference_pool.stop() if teacher_inference_pool is not None: await teacher_inference_pool.stop() event_loop_lag_monitor_task.cancel() diff --git a/src/prime_rl/orchestrator/scheduler.py b/src/prime_rl/orchestrator/scheduler.py index ab5d8e0b3f..e67e44b768 100644 --- a/src/prime_rl/orchestrator/scheduler.py +++ b/src/prime_rl/orchestrator/scheduler.py @@ -68,7 +68,8 @@ def __init__( tasks_per_minute: int | None, enable_policy_updates: bool = True, lora_name: str | None = None, - eval_inference_pool: InferencePool | None = None, + teacher_rollout_clients: list[vf.ClientConfig] | None = None, + teacher_rollout_model_name: str | None = None, ): self.logger = get_logger() if tasks_per_minute is not None: @@ -90,11 +91,9 @@ def __init__( self.model_name = self.config.model.name self.json_logging = config.log.json_logging - # Inference pool - used for rollout client selection and metrics self.inference_pool = inference_pool - # Eval inference pool - receives weight updates and serves evals. - # Defaults to inference_pool (standard RL where one pool does both). - self.eval_inference_pool = eval_inference_pool or inference_pool + self.teacher_rollout_clients = teacher_rollout_clients + self.teacher_rollout_model_name = teacher_rollout_model_name group_scoring_envs = [env.name for env in train_envs if env.requires_group_scoring] if group_scoring_envs: @@ -170,6 +169,14 @@ async def _select_least_loaded_client(self) -> vf.ClientConfig: inflight = Counter(self._client_identity(info.client_config) for info in self.inflight_requests.values()) return min(clients, key=lambda c: inflight[self._client_identity(c)]) + def _resolve_rollout_request_target(self, client_config: vf.ClientConfig) -> tuple[vf.ClientConfig, str]: + if self.teacher_rollout_clients is None: + return client_config, self.model_name + + teacher_client = self.teacher_rollout_clients[client_config.client_idx % len(self.teacher_rollout_clients)] + assert self.teacher_rollout_model_name is not None + return teacher_client, self.teacher_rollout_model_name + async def drop_group(self, group_id: int) -> int: """Drop a group and cancel any remaining in-flight rollouts for it. Returns the number of cancelled rollouts.""" tasks_to_cancel = [] @@ -203,15 +210,16 @@ async def schedule_rollout(self, group_id: int): env_name = group.example["env_name"] env = self.train_envs.get(env_name) + request_client_config, request_model_name = self._resolve_rollout_request_target(client_config) cache_salt = str(self.ckpt_step) if env.requires_group_scoring: rollout_count = group.rollouts_to_schedule group.rollouts_to_schedule = 0 task = asyncio.create_task( env.run_group( - client=client_config, + client=request_client_config, example=group.example, - model_name=self.model_name, + model_name=request_model_name, rollouts_per_example=rollout_count, cache_salt=cache_salt, ) @@ -221,9 +229,9 @@ async def schedule_rollout(self, group_id: int): group.rollouts_to_schedule -= 1 task = asyncio.create_task( env.run_rollout( - client=client_config, + client=request_client_config, example=group.example, - model_name=self.model_name, + model_name=request_model_name, cache_salt=cache_salt, ) ) @@ -307,15 +315,14 @@ async def _apply_policy_update(self, next_ckpt_step: int) -> None: update_weights_start_time = time.perf_counter() weights_path = get_step_path(get_broadcast_dir(self.config.output_dir), next_ckpt_step) - await self.eval_inference_pool.update_weights(weights_path, lora_name=self.lora_name, step=next_ckpt_step) + await self.inference_pool.update_weights(weights_path, lora_name=self.lora_name, step=next_ckpt_step) self.update_weights_time = time.perf_counter() - update_weights_start_time self.logger.debug(f"Updated weights to step {next_ckpt_step} in {self.update_weights_time:.2f}s") self.ckpt_step = next_ckpt_step if self.lora_name is not None: - self.eval_inference_pool.update_model_name(self.lora_name) - if self.eval_inference_pool is self.inference_pool: - self.model_name = self.lora_name + self.inference_pool.update_model_name(self.lora_name) + self.model_name = self.lora_name self.checkpoint_ready.set() await self._update_off_policy() diff --git a/src/prime_rl/orchestrator/utils.py b/src/prime_rl/orchestrator/utils.py index 524f607741..abf778bf4c 100644 --- a/src/prime_rl/orchestrator/utils.py +++ b/src/prime_rl/orchestrator/utils.py @@ -155,12 +155,13 @@ def setup_external_rollout_model(config: OrchestratorConfig, logger) -> tuple[An """Resolve rollout client/model and whether policy updates should be enabled.""" rollout_client_config = config.client rollout_model_name = config.model.name - enable_policy_updates = True + enable_policy_updates = config.enable_policy_updates if config.enable_policy_updates is not None else True if config.teacher_rollout_model is not None: rollout_client_config = config.teacher_rollout_model.client rollout_model_name = config.teacher_rollout_model.model.name - enable_policy_updates = False + if config.enable_policy_updates is None: + enable_policy_updates = False logger.info( f"Using external teacher rollout model (base_url={', '.join(rollout_client_config.base_url)}, " f"model={rollout_model_name})" diff --git a/tests/unit/orchestrator/test_scheduler.py b/tests/unit/orchestrator/test_scheduler.py index 86129e9606..3af71616b7 100644 --- a/tests/unit/orchestrator/test_scheduler.py +++ b/tests/unit/orchestrator/test_scheduler.py @@ -5,7 +5,7 @@ import verifiers as vf -from prime_rl.orchestrator.scheduler import InflightRequest, Scheduler +from prime_rl.orchestrator.scheduler import GroupState, InflightRequest, Scheduler from prime_rl.utils.async_utils import safe_cancel @@ -31,6 +31,9 @@ def make_scheduler() -> Scheduler: scheduler.inflight_policy_update_task = None scheduler.update_policy_task = None scheduler.enable_policy_updates = True + scheduler.rate_limiter = None + scheduler.teacher_rollout_clients = None + scheduler.teacher_rollout_model_name = None return scheduler @@ -106,7 +109,6 @@ async def update_weights(weight_dir, lora_name=None, step=0) -> None: update_weights=update_weights, update_model_name=MagicMock(), ) - scheduler.eval_inference_pool = scheduler.inference_pool scheduler._update_off_policy = AsyncMock() with ( @@ -147,7 +149,6 @@ async def update_weights(weight_dir, lora_name=None, step=0) -> None: update_weights=update_weights, update_model_name=MagicMock(), ) - scheduler.eval_inference_pool = scheduler.inference_pool scheduler._update_off_policy = AsyncMock() with ( @@ -178,16 +179,14 @@ def test_client_identity_distinguishes_base_url_and_dp_rank(): assert Scheduler._client_identity(client_a) != Scheduler._client_identity(client_b) -def test_lora_policy_update_keeps_teacher_rollout_model_name_with_separate_eval_pool(): +def test_lora_policy_update_keeps_student_model_name_with_teacher_rollout_override(): async def run() -> None: scheduler = make_scheduler() - scheduler.model_name = "teacher-model" + scheduler.model_name = "student-model" scheduler.lora_name = "student-lora" + scheduler.teacher_rollout_model_name = "teacher-model" scheduler.inference_pool = SimpleNamespace( - update_model_name=MagicMock(), - ) - scheduler.eval_inference_pool = SimpleNamespace( update_weights=AsyncMock(), update_model_name=MagicMock(), ) @@ -199,9 +198,45 @@ async def run() -> None: ): await scheduler.maybe_update_policy() - scheduler.eval_inference_pool.update_weights.assert_awaited_once() - scheduler.eval_inference_pool.update_model_name.assert_called_once_with("student-lora") - scheduler.inference_pool.update_model_name.assert_not_called() - assert scheduler.model_name == "teacher-model" + scheduler.inference_pool.update_weights.assert_awaited_once() + scheduler.inference_pool.update_model_name.assert_called_once_with("student-lora") + assert scheduler.model_name == "student-lora" + assert scheduler.teacher_rollout_model_name == "teacher-model" + + asyncio.run(run()) + + +def test_schedule_rollout_applies_teacher_override_at_request_submission(): + async def run() -> None: + scheduler = make_scheduler() + student_client = vf.ClientConfig(api_base_url="http://student.example/v1") + teacher_client = vf.ClientConfig(api_base_url="http://teacher.example/v1") + env = SimpleNamespace( + requires_group_scoring=False, + run_rollout=AsyncMock(return_value=[]), + ) + scheduler.inference_pool = SimpleNamespace(train_clients=[student_client]) + scheduler.teacher_rollout_clients = [teacher_client] + scheduler.teacher_rollout_model_name = "teacher-model" + scheduler.train_envs = SimpleNamespace(get=MagicMock(return_value=env)) + scheduler.groups = { + 0: GroupState( + example={"env_name": "math", "example_id": "ex-1"}, + rollouts_to_schedule=1, + ) + } + + await scheduler.schedule_rollout(group_id=0) + await asyncio.gather(*scheduler.inflight_requests) + + env.run_rollout.assert_awaited_once_with( + client=teacher_client, + example={"env_name": "math", "example_id": "ex-1"}, + model_name="teacher-model", + cache_salt="7", + ) + assert scheduler.groups[0].pinned_client is student_client + [info] = scheduler.inflight_requests.values() + assert info.client_config is student_client asyncio.run(run()) diff --git a/tests/unit/test_configs.py b/tests/unit/test_configs.py index 0f2f095063..8b84197066 100644 --- a/tests/unit/test_configs.py +++ b/tests/unit/test_configs.py @@ -161,7 +161,7 @@ def test_selective_activation_checkpointing_requires_custom_impl(): TrainerModelConfig.model_validate({"impl": "hf", "ac": {"mode": "selective"}}) -def test_teacher_rollout_student_eval_pool_requires_inference_config(): +def test_teacher_rollout_policy_updates_require_inference_config(): base_config = { "trainer": {}, "orchestrator": { @@ -175,7 +175,7 @@ def test_teacher_rollout_student_eval_pool_requires_inference_config(): } config = RLConfig.model_validate(base_config) - assert not config.orchestrator.use_student_eval_inference_pool + assert config.orchestrator.enable_policy_updates is None config = RLConfig.model_validate({**base_config, "inference": {}}) - assert config.orchestrator.use_student_eval_inference_pool + assert config.orchestrator.enable_policy_updates From 5f6ccf5ef612baabc2a80aabdd8093ca3017ea67 Mon Sep 17 00:00:00 2001 From: Timothy Kostolansky <39891386+tim0120@users.noreply.github.com> Date: Fri, 15 May 2026 21:22:58 +0000 Subject: [PATCH 04/47] fix: address teacher rollout review comments --- .../prime-rl-configs/src/prime_rl/configs/orchestrator.py | 5 +++++ src/prime_rl/orchestrator/orchestrator.py | 4 ++-- src/prime_rl/orchestrator/scheduler.py | 2 ++ tests/unit/test_configs.py | 2 ++ 4 files changed, 11 insertions(+), 2 deletions(-) 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 f70d90756f..0f2c04eddd 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py @@ -897,6 +897,11 @@ class TeacherModelConfig(BaseConfig): class TeacherRolloutModelConfig(BaseConfig): """Configures an external teacher model used to generate rollout text.""" + client_type: Annotated[ + str, + Field(description="The verifiers client type to use for rollout generation."), + ] = "openai_chat_completions" + client: Annotated[ ClientConfig, Field(description="The OAI client configuration for rollout generation."), diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index 7d3dd07e32..4e8da4d562 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -153,10 +153,10 @@ async def orchestrate(config: OrchestratorConfig): teacher_rollout_model_name = None use_teacher_rollout_override = config.teacher_rollout_model is not None and enable_policy_updates if use_teacher_rollout_override: - logger.info("Using external rollout model (MITO) without renderer client") + logger.info(f"Using external rollout model ({config.teacher_rollout_model.client_type})") teacher_rollout_clients = setup_clients( rollout_client_config, - client_type="openai_chat_completions", + client_type=config.teacher_rollout_model.client_type, ) teacher_rollout_model_name = rollout_model_name renderer = None diff --git a/src/prime_rl/orchestrator/scheduler.py b/src/prime_rl/orchestrator/scheduler.py index e67e44b768..41498af099 100644 --- a/src/prime_rl/orchestrator/scheduler.py +++ b/src/prime_rl/orchestrator/scheduler.py @@ -173,6 +173,8 @@ def _resolve_rollout_request_target(self, client_config: vf.ClientConfig) -> tup if self.teacher_rollout_clients is None: return client_config, self.model_name + # Multiple teacher URLs are uncommon, but when configured we spread requests + # deterministically using the selected student client's index. teacher_client = self.teacher_rollout_clients[client_config.client_idx % len(self.teacher_rollout_clients)] assert self.teacher_rollout_model_name is not None return teacher_client, self.teacher_rollout_model_name diff --git a/tests/unit/test_configs.py b/tests/unit/test_configs.py index 8b84197066..8f57d6426a 100644 --- a/tests/unit/test_configs.py +++ b/tests/unit/test_configs.py @@ -168,6 +168,7 @@ def test_teacher_rollout_policy_updates_require_inference_config(): "use_sft_loss": True, "use_token_client": False, "teacher_rollout_model": { + "client_type": "custom_chat_client", "client": {"base_url": ["http://teacher.example/v1"]}, "model": {"name": "teacher-model"}, }, @@ -176,6 +177,7 @@ def test_teacher_rollout_policy_updates_require_inference_config(): config = RLConfig.model_validate(base_config) assert config.orchestrator.enable_policy_updates is None + assert config.orchestrator.teacher_rollout_model.client_type == "custom_chat_client" config = RLConfig.model_validate({**base_config, "inference": {}}) assert config.orchestrator.enable_policy_updates From 6cd5ecd77bd33595a94045c87ccf90d298ceb4fe Mon Sep 17 00:00:00 2001 From: Timothy Kostolansky <39891386+tim0120@users.noreply.github.com> Date: Fri, 15 May 2026 21:40:22 +0000 Subject: [PATCH 05/47] fix: clarify teacher rollout client routing --- .../src/prime_rl/configs/orchestrator.py | 2 +- src/prime_rl/orchestrator/orchestrator.py | 14 ++++++++++---- src/prime_rl/orchestrator/scheduler.py | 5 +++-- 3 files changed, 14 insertions(+), 7 deletions(-) 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 0f2c04eddd..85dd1ca0ef 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py @@ -1223,7 +1223,7 @@ def validate_sft_distill_mode(self): if has_teacher and self.use_renderer: raise ValueError( "orchestrator.use_renderer must be false when orchestrator.teacher_rollout_model is configured " - "(external rollout uses MITO)." + "(teacher rollout uses teacher_rollout_model.client_type)." ) return self diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index 4e8da4d562..b7da3228e3 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -153,10 +153,16 @@ async def orchestrate(config: OrchestratorConfig): teacher_rollout_model_name = None use_teacher_rollout_override = config.teacher_rollout_model is not None and enable_policy_updates if use_teacher_rollout_override: - logger.info(f"Using external rollout model ({config.teacher_rollout_model.client_type})") + teacher_client_type = config.teacher_rollout_model.client_type + teacher_client_label = ( + "MITO (openai_chat_completions)" + if teacher_client_type == "openai_chat_completions" + else teacher_client_type + ) + logger.info(f"Using teacher rollout override ({teacher_client_label}, model={rollout_model_name})") teacher_rollout_clients = setup_clients( rollout_client_config, - client_type=config.teacher_rollout_model.client_type, + client_type=teacher_client_type, ) teacher_rollout_model_name = rollout_model_name renderer = None @@ -919,8 +925,8 @@ async def setup_rollout_inference_pool( ``config.use_renderer`` (mutually exclusive — config-level validators block both being True): - - external teacher rollout → MITO (``openai_chat_completions``), - forced regardless of the toggles (config-level validator + - external teacher rollout → configured teacher rollout client type, + selected independently of the toggles (config-level validator rejects ``use_token_client`` / ``use_renderer`` in that case) - ``use_renderer=True`` → renderer client (``/v1/generate``). Not allowed for VLMs (validated at config time). diff --git a/src/prime_rl/orchestrator/scheduler.py b/src/prime_rl/orchestrator/scheduler.py index 41498af099..020c8ef23d 100644 --- a/src/prime_rl/orchestrator/scheduler.py +++ b/src/prime_rl/orchestrator/scheduler.py @@ -173,8 +173,9 @@ def _resolve_rollout_request_target(self, client_config: vf.ClientConfig) -> tup if self.teacher_rollout_clients is None: return client_config, self.model_name - # Multiple teacher URLs are uncommon, but when configured we spread requests - # deterministically using the selected student client's index. + # The scheduler pins/load-balances against the student inference pool. + # Map that selected logical client onto the teacher client set, which may + # have a different size due to different base_url or dp_rank_count settings. teacher_client = self.teacher_rollout_clients[client_config.client_idx % len(self.teacher_rollout_clients)] assert self.teacher_rollout_model_name is not None return teacher_client, self.teacher_rollout_model_name From 634e0812866135b1ede03944106de63bdbf88d8c Mon Sep 17 00:00:00 2001 From: Timothy Kostolansky <39891386+tim0120@users.noreply.github.com> Date: Fri, 15 May 2026 22:07:58 +0000 Subject: [PATCH 06/47] fix: honor teacher rollout client type --- src/prime_rl/orchestrator/orchestrator.py | 12 +++++++++--- src/prime_rl/utils/client.py | 5 +++++ tests/unit/orchestrator/test_orchestrator_setup.py | 13 ++++++++++--- 3 files changed, 24 insertions(+), 6 deletions(-) diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index b7da3228e3..b0f69e72fb 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -936,12 +936,18 @@ async def setup_rollout_inference_pool( - both False → MITO (``openai_chat_completions``). """ if config.teacher_rollout_model is not None: - logger.info("Using external rollout model (MITO) without renderer client") + teacher_client_type = config.teacher_rollout_model.client_type + teacher_client_label = ( + "MITO (openai_chat_completions)" + if teacher_client_type == "openai_chat_completions" + else teacher_client_type + ) + logger.info(f"Using external rollout model ({teacher_client_label}) without renderer client") inference_pool = await setup_inference_pool( rollout_client_config, model_name=rollout_model_name, - train_client_type="openai_chat_completions", - eval_client_type="openai_chat_completions", + train_client_type=teacher_client_type, + eval_client_type=teacher_client_type, ) return None, inference_pool diff --git a/src/prime_rl/utils/client.py b/src/prime_rl/utils/client.py index 21659dfc46..a7ed299fdf 100644 --- a/src/prime_rl/utils/client.py +++ b/src/prime_rl/utils/client.py @@ -20,6 +20,11 @@ class InferencePool(Protocol): """Protocol for inference pools (static or elastic).""" + @property + def model_name(self) -> str: + """Get current model name for inference requests.""" + ... + @property def train_clients(self) -> list[vf.ClientConfig]: """Get inference clients.""" diff --git a/tests/unit/orchestrator/test_orchestrator_setup.py b/tests/unit/orchestrator/test_orchestrator_setup.py index ff9bb5b79f..bc009280bb 100644 --- a/tests/unit/orchestrator/test_orchestrator_setup.py +++ b/tests/unit/orchestrator/test_orchestrator_setup.py @@ -9,7 +9,7 @@ def test_setup_rollout_inference_pool_uses_plain_client_for_external_teacher_rol async def run() -> None: tokenizer = object() config = SimpleNamespace( - teacher_rollout_model=SimpleNamespace(), + teacher_rollout_model=SimpleNamespace(client_type="custom_chat_client"), model=SimpleNamespace(renderer="auto", name="student-model"), ) rollout_client_config = SimpleNamespace(base_url=["https://api.pinference.ai/api/v1"]) @@ -18,8 +18,9 @@ async def run() -> None: with ( patch( - "prime_rl.orchestrator.orchestrator.setup_inference_pool", new=AsyncMock(return_value=inference_pool) - ), + "prime_rl.orchestrator.orchestrator.setup_inference_pool", + new=AsyncMock(return_value=inference_pool), + ) as setup_pool_mock, patch("prime_rl.orchestrator.orchestrator.create_renderer") as create_renderer_mock, ): renderer, returned_pool = await setup_rollout_inference_pool( @@ -33,6 +34,12 @@ async def run() -> None: assert renderer is None assert returned_pool is inference_pool create_renderer_mock.assert_not_called() + setup_pool_mock.assert_awaited_once_with( + rollout_client_config, + model_name="teacher-model", + train_client_type="custom_chat_client", + eval_client_type="custom_chat_client", + ) asyncio.run(run()) From e8066ccd08c55a0af1d5006cfe8ad31275f647c7 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Mon, 18 May 2026 21:36:43 +0000 Subject: [PATCH 07/47] refactor(orchestrator): rename model/teacher_model to student/teacher MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Consolidates 3 training modes (rl, opd, sft) under a unified config structure: `student` (always present) and `teacher` (optional, role determined by `training_mode`). Removes scattered flags (`use_sft_loss`, `update_student_inference_weights`) in favour of a single discriminator. Backward-compat aliases (`model` → `student`, `teacher_model` → `teacher`) keep existing TOML configs parsing without changes. Also adds `configs/reverse_text/debug_{rl,opd,sft}.toml` for local mode debugging. Co-Authored-By: Claude Sonnet 4.6 --- .../integration/reverse_text_lora/resume.toml | 2 +- .../integration/reverse_text_lora/start.toml | 2 +- .../reverse_text_multi_run/orchestrator.toml | 2 +- configs/elastic/rl.toml | 2 +- configs/reverse_text/debug_opd.toml | 43 +++++++ configs/reverse_text/debug_rl.toml | 32 +++++ configs/reverse_text/debug_sft.toml | 42 ++++++ docs/on_policy_distillation.md | 8 +- examples/alphabet_sort/sft_distill_hard.toml | 6 +- examples/wiki_search/rl.toml | 2 +- .../src/prime_rl/configs/orchestrator.py | 120 ++++++------------ .../src/prime_rl/configs/rl.py | 111 ++++++++-------- .../src/prime_rl/utils/validation.py | 8 +- skills/config/SKILL.md | 4 +- src/prime_rl/entrypoints/rl.py | 18 ++- src/prime_rl/orchestrator/orchestrator.py | 64 +++++----- src/prime_rl/orchestrator/scheduler.py | 18 +-- src/prime_rl/orchestrator/utils.py | 31 +++-- .../trainer/rl/broadcast/filesystem.py | 2 +- src/prime_rl/trainer/runs.py | 16 ++- .../orchestrator/test_orchestrator_setup.py | 34 ++--- tests/unit/orchestrator/test_scheduler.py | 12 +- tests/unit/test_configs.py | 42 +++--- 23 files changed, 349 insertions(+), 272 deletions(-) create mode 100644 configs/reverse_text/debug_opd.toml create mode 100644 configs/reverse_text/debug_rl.toml create mode 100644 configs/reverse_text/debug_sft.toml diff --git a/configs/ci/integration/reverse_text_lora/resume.toml b/configs/ci/integration/reverse_text_lora/resume.toml index e2b7e66ca2..cef3d65e17 100644 --- a/configs/ci/integration/reverse_text_lora/resume.toml +++ b/configs/ci/integration/reverse_text_lora/resume.toml @@ -20,7 +20,7 @@ save_adapter_separately = true batch_size = 128 rollouts_per_example = 16 -[orchestrator.model.lora] +[orchestrator.student.model.lora] name = "r8-1e-4" [orchestrator.train.sampling] diff --git a/configs/ci/integration/reverse_text_lora/start.toml b/configs/ci/integration/reverse_text_lora/start.toml index 28e76d60f8..2460203c8a 100644 --- a/configs/ci/integration/reverse_text_lora/start.toml +++ b/configs/ci/integration/reverse_text_lora/start.toml @@ -19,7 +19,7 @@ save_adapter_separately = true batch_size = 128 rollouts_per_example = 16 -[orchestrator.model.lora] +[orchestrator.student.model.lora] name = "r8-1e-4" [orchestrator.train.sampling] diff --git a/configs/ci/integration/reverse_text_multi_run/orchestrator.toml b/configs/ci/integration/reverse_text_multi_run/orchestrator.toml index be34be66fb..92f403ff73 100644 --- a/configs/ci/integration/reverse_text_multi_run/orchestrator.toml +++ b/configs/ci/integration/reverse_text_multi_run/orchestrator.toml @@ -5,7 +5,7 @@ rollouts_per_example = 16 seq_len = 2048 max_steps = 20 -[model] +[model.model] name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT" [optim] diff --git a/configs/elastic/rl.toml b/configs/elastic/rl.toml index 387bfddbf2..4233b1e437 100644 --- a/configs/elastic/rl.toml +++ b/configs/elastic/rl.toml @@ -35,7 +35,7 @@ rollouts_per_example = 8 [orchestrator.train.sampling] max_completion_tokens = 768 -[orchestrator.client.elastic] +[orchestrator.model.client.elastic] hostname = "localhost" port = 8000 sync_interval = 5.0 diff --git a/configs/reverse_text/debug_opd.toml b/configs/reverse_text/debug_opd.toml new file mode 100644 index 0000000000..15aa8d75bb --- /dev/null +++ b/configs/reverse_text/debug_opd.toml @@ -0,0 +1,43 @@ +max_steps = 20 +seq_len = 2048 + +[model] +name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT" + +[wandb] +project = "reverse-text-debug" +name = "debug-opd" + +[orchestrator] +batch_size = 128 +rollouts_per_example = 16 +training_mode = "opd" + +[orchestrator.train.sampling] +max_completion_tokens = 128 + +[[orchestrator.train.env]] +id = "reverse-text" + +[orchestrator.eval] +interval = 5 + +[[orchestrator.eval.env]] +id = "reverse-text" + +[orchestrator.teacher.client] +base_url = ["https://api.pinference.ai/api/v1"] +api_key_var = "PRIME_API_KEY" + +[orchestrator.teacher.model] +name = "qwen/qwen3-4b-instruct" + +[trainer.optim] +lr = 3e-6 + +[trainer.loss] +teacher_tau = 0.5 + +[ckpt] + +[inference] diff --git a/configs/reverse_text/debug_rl.toml b/configs/reverse_text/debug_rl.toml new file mode 100644 index 0000000000..762d3fed35 --- /dev/null +++ b/configs/reverse_text/debug_rl.toml @@ -0,0 +1,32 @@ +max_steps = 20 +seq_len = 2048 + +[model] +name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT" + +[wandb] +project = "reverse-text-debug" +name = "debug-rl" + +[orchestrator] +batch_size = 128 +rollouts_per_example = 16 + +[orchestrator.train.sampling] +max_completion_tokens = 128 + +[[orchestrator.train.env]] +id = "reverse-text" + +[orchestrator.eval] +interval = 5 + +[[orchestrator.eval.env]] +id = "reverse-text" + +[trainer.optim] +lr = 3e-6 + +[ckpt] + +[inference] diff --git a/configs/reverse_text/debug_sft.toml b/configs/reverse_text/debug_sft.toml new file mode 100644 index 0000000000..d67f760c2a --- /dev/null +++ b/configs/reverse_text/debug_sft.toml @@ -0,0 +1,42 @@ +max_steps = 20 +seq_len = 2048 + +[model] +name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT" + +[wandb] +project = "reverse-text-debug" +name = "debug-sft" + +[orchestrator] +batch_size = 128 +rollouts_per_example = 4 +use_token_client = false +use_renderer = false +training_mode = "sft" + +[orchestrator.train.sampling] +max_completion_tokens = 128 + +[[orchestrator.train.env]] +id = "reverse-text" + +[orchestrator.eval] +interval = 5 + +[[orchestrator.eval.env]] +id = "reverse-text" + +[orchestrator.teacher.client] +base_url = ["https://api.pinference.ai/api/v1"] +api_key_var = "PRIME_API_KEY" + +[orchestrator.teacher.model] +name = "qwen/qwen3-4b-instruct" + +[trainer.optim] +lr = 3e-6 + +[ckpt] + +[inference] diff --git a/docs/on_policy_distillation.md b/docs/on_policy_distillation.md index 6f434bbfe1..01485bbf60 100644 --- a/docs/on_policy_distillation.md +++ b/docs/on_policy_distillation.md @@ -73,18 +73,18 @@ type = "sft" [orchestrator] use_token_client = false use_renderer = false -use_sft_loss = true +training_mode = "sft" -[orchestrator.teacher_rollout_model.client] +[orchestrator.teacher_model.client] base_url = ["https://your-openai-compatible-endpoint/v1"] skip_model_check = true -[orchestrator.teacher_rollout_model.model] +[orchestrator.teacher_model.model] name = "teacher-model-name" ``` In this mode: -- Rollouts are generated from `orchestrator.teacher_rollout_model` +- Rollouts are generated from `orchestrator.teacher_model` - The orchestrator uses text-level reconstruction with the student tokenizer - The RL trainer optimizes masked NLL (`trainer.loss.type = "sft"`) - Omit `[inference]` (no local inference server required) diff --git a/examples/alphabet_sort/sft_distill_hard.toml b/examples/alphabet_sort/sft_distill_hard.toml index 8eab8c9dc7..da3c77199c 100644 --- a/examples/alphabet_sort/sft_distill_hard.toml +++ b/examples/alphabet_sort/sft_distill_hard.toml @@ -31,17 +31,17 @@ batch_size = 256 rollouts_per_example = 4 use_token_client = false use_renderer = false -use_sft_loss = true +training_mode = "sft" [orchestrator.train.sampling] max_completion_tokens = 512 temperature = 0.7 -[orchestrator.teacher_rollout_model.client] +[orchestrator.teacher.client] base_url = ["https://api.pinference.ai/api/v1"] api_key_var = "PRIME_API_KEY" -[orchestrator.teacher_rollout_model.model] +[orchestrator.teacher.model] name = "qwen/qwen3-235b-a22b-instruct-2507" [[orchestrator.train.env]] diff --git a/examples/wiki_search/rl.toml b/examples/wiki_search/rl.toml index 6abbb3d815..81d78c8726 100644 --- a/examples/wiki_search/rl.toml +++ b/examples/wiki_search/rl.toml @@ -34,7 +34,7 @@ batch_size = 512 rollouts_per_example = 16 oversampling_factor = 2.0 -[orchestrator.model.lora] +[orchestrator.student.model.lora] name = "qwen3-4b-wiki-search" [orchestrator.train.sampling] 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 fb6c8d40ae..1364dfb067 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py @@ -880,33 +880,19 @@ class OrchestratorExperimentalConfig(BaseConfig): """Experimental features for the orchestrator.""" -class TeacherModelConfig(BaseConfig): - """Configures the teacher model for computing teacher logprobs (e.g. for distillation).""" - - client: Annotated[ - ClientConfig, - Field(description="The OAI client configuration for the teacher model."), - ] = ClientConfig() +class RolloutModelConfig(BaseConfig): + """Model + client pair for a rollout participant (student or teacher).""" model: Annotated[ ModelConfig, - Field(description="The model configuration for the teacher model."), + Field(description="The model configuration."), ] = ModelConfig() - -class TeacherRolloutModelConfig(BaseConfig): - """Configures an external teacher model used to generate rollout text.""" - client: Annotated[ ClientConfig, - Field(description="The OAI client configuration for rollout generation."), + Field(description="The OAI client configuration."), ] = ClientConfig() - model: Annotated[ - ModelConfig, - Field(description="The model configuration for rollout generation."), - ] = ModelConfig() - class OrchestratorConfig(BaseConfig): """Configures the orchestrator for RL training.""" @@ -914,11 +900,14 @@ class OrchestratorConfig(BaseConfig): # Training environments and sampling train: TrainConfig = TrainConfig() - # The OAI client configuration - client: ClientConfig = ClientConfig() - - # The model configuration - model: ModelConfig = ModelConfig() + # Student model + client (the model being trained) + student: Annotated[ + RolloutModelConfig, + Field( + validation_alias=AliasChoices("student", "model"), + description="Student model configuration (the model being trained).", + ), + ] = RolloutModelConfig() # The tokenizer configuration tokenizer: TokenizerConfig = TokenizerConfig() @@ -929,49 +918,30 @@ class OrchestratorConfig(BaseConfig): # The optimizer configuration (per-run LR for multi-run training) optim: OptimizerConfig = OptimizerConfig() - # The teacher model configuration (optional) - teacher_model: Annotated[ - TeacherModelConfig | None, - Field( - description="The teacher model configuration for computing teacher logprobs (e.g. for distillation). " - "If provided, teacher logprobs will be computed using the specified model. " - "If None, no teacher model will be used." - ), - ] = None - - # External teacher rollout model configuration (optional) - teacher_rollout_model: Annotated[ - TeacherRolloutModelConfig | None, + # Teacher model + client (optional; role determined by training_mode) + teacher: Annotated[ + RolloutModelConfig | None, Field( + validation_alias=AliasChoices("teacher", "teacher_model"), description=( - "Optional external teacher model used for rollout generation. " - "When set, rollouts are generated from this endpoint/model instead of the student inference server." + "Teacher model configuration. Role depends on training_mode: " + "opd — teacher computes logprobs; sft — teacher generates rollouts." ), ), ] = None - # When True, trainer uses SFT loss instead of RL loss (per-run override for hosted multi-tenant training) - use_sft_loss: Annotated[ - bool, + # Training mode: drives validation and runtime wiring + training_mode: Annotated[ + Literal["rl", "opd", "sft"], Field( description=( - "When True, use SFT masked NLL loss instead of the trainer's configured RL loss. " - "Requires a teacher_rollout_model to be configured." + "Training mode. " + "rl: student generates rollouts, no teacher. " + "opd: student generates rollouts, teacher computes logprobs (teacher_tau > 0). " + "sft: teacher generates rollouts, student inference pool used for evals and weight sync." ), ), - ] = False - - update_student_inference_weights: Annotated[ - bool | None, - Field( - description=( - "Whether teacher_rollout_model runs should push trained student weights to the student inference " - "server. Leave unset for the usual [inference] path; RLConfig enables it automatically. Set True " - "when teacher_rollout_model is configured and orchestrator.client points to an externally started " - "student inference server. Defaults to False for teacher_rollout_model runs and True otherwise." - ), - ), - ] = None + ] = "rl" # The evaluation configuration eval: EvalConfig | None = None @@ -1194,9 +1164,9 @@ def _env_to_train(cls, data: Any) -> Any: @model_validator(mode="after") def auto_setup_tokenizer(self): if self.tokenizer.name is None: - self.tokenizer.name = self.model.name + self.tokenizer.name = self.student.model.name if self.tokenizer.trust_remote_code is None: - self.tokenizer.trust_remote_code = self.model.trust_remote_code + self.tokenizer.trust_remote_code = self.student.model.trust_remote_code return self @model_validator(mode="after") @@ -1207,27 +1177,21 @@ def validate_unique_filter_types(self): return self @model_validator(mode="after") - def validate_sft_distill_mode(self): - """Enforce the SFT hard distill invariants that involve only orchestrator fields. - - Runs at ``OrchestratorConfig`` level so hosted deployments (which load this - config standalone via the ``orchestrator`` entrypoint) get the same guarantees - as the combined ``rl`` entrypoint. - """ - has_teacher = self.teacher_rollout_model is not None - if self.use_sft_loss and not has_teacher: + def validate_training_mode(self): + """Enforce training mode invariants that involve only orchestrator fields.""" + has_teacher = self.teacher is not None + if self.training_mode == "rl" and has_teacher: + raise ValueError("orchestrator.teacher must not be set when training_mode = 'rl'.") + if self.training_mode in ("opd", "sft") and not has_teacher: + raise ValueError(f"orchestrator.teacher must be configured when training_mode = '{self.training_mode}'.") + if self.training_mode == "sft" and self.use_token_client: raise ValueError( - "orchestrator.use_sft_loss = true requires orchestrator.teacher_rollout_model to be configured." - ) - if has_teacher and not self.use_sft_loss: - raise ValueError("orchestrator.teacher_rollout_model requires orchestrator.use_sft_loss = true.") - if has_teacher and self.use_token_client: - raise ValueError( - "orchestrator.use_token_client must be false when orchestrator.teacher_rollout_model is configured." + "orchestrator.use_token_client must be false when training_mode = 'sft' " + "(teacher rollout uses the plain OpenAI chat-completions client)." ) - if has_teacher and self.use_renderer: + if self.training_mode == "sft" and self.use_renderer: raise ValueError( - "orchestrator.use_renderer must be false when orchestrator.teacher_rollout_model is configured " + "orchestrator.use_renderer must be false when training_mode = 'sft' " "(teacher rollout uses the plain OpenAI chat-completions client)." ) return self @@ -1257,7 +1221,7 @@ def validate_renderer_vs_vlm(self): them client-side. VLMs need server-side image preprocessing and chat templating, so they must use MITO — fail loudly when both are set.""" - if self.use_renderer and self.model.vlm is not None: + if self.use_renderer and self.student.model.vlm is not None: raise ValueError( "orchestrator.use_renderer is not supported for VLMs. Use MITO " "(``use_token_client=false`` and ``use_renderer=false``) so image preprocessing and chat " @@ -1364,7 +1328,7 @@ def auto_setup_bench(self): @model_validator(mode="after") def resolve_env_config(self): """Populate extra_env_kwargs and vLLM sampling defaults from top-level fields.""" - is_vllm = self.teacher_rollout_model is None + is_vllm = self.training_mode != "sft" for env in self.train.env: env.extra_env_kwargs.update(max_seq_len=self.seq_len) if is_vllm: 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 14730c89c4..d8a57bd396 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/rl.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/rl.py @@ -426,10 +426,10 @@ def validate_quantize_in_weight_transfer(self): def validate_teacher_model(self): if ( self.trainer.loss.type == "default" and self.trainer.loss.teacher_tau > 0 - ) and not self.orchestrator.teacher_model: + ) and not self.orchestrator.teacher: raise ValueError( - "teacher_model must be configured when teacher_tau > 0. " - "Either set teacher_tau = 0, set deployment.num_teacher_gpus, or configure teacher_model manually." + "orchestrator.teacher must be configured when teacher_tau > 0. " + "Either set teacher_tau = 0, set deployment.num_teacher_gpus, or configure orchestrator.teacher manually." ) return self @@ -549,13 +549,13 @@ def auto_setup_model(self): inference_model_explicitly_set = "name" in self.inference.model.model_fields_set if not inference_model_explicitly_set: self.inference.model.name = self.model.name - self.orchestrator.model.name = self.inference.model.name + self.orchestrator.student.model.name = self.inference.model.name else: - self.orchestrator.model.name = self.model.name + self.orchestrator.student.model.name = self.model.name if self.model.vlm is not None: self.trainer.model.vlm = self.model.vlm - self.orchestrator.model.vlm = self.model.vlm + self.orchestrator.student.model.vlm = self.model.vlm if self.inference is not None: self.inference.model.vlm = self.model.vlm @@ -571,17 +571,21 @@ def auto_setup_tokenizer(self): # in name/trust_remote_code from model config where still unset. self.trainer.tokenizer = self.tokenizer.model_copy() self.orchestrator.tokenizer = self.tokenizer.model_copy() - for component in (self.trainer, self.orchestrator): - if component.tokenizer.name is None: - component.tokenizer.name = component.model.name - if component.tokenizer.trust_remote_code is None: - component.tokenizer.trust_remote_code = component.model.trust_remote_code + if self.trainer.tokenizer.name is None: + self.trainer.tokenizer.name = self.trainer.model.name + if self.trainer.tokenizer.trust_remote_code is None: + self.trainer.tokenizer.trust_remote_code = self.trainer.model.trust_remote_code + if self.orchestrator.tokenizer.name is None: + self.orchestrator.tokenizer.name = self.orchestrator.student.model.name + if self.orchestrator.tokenizer.trust_remote_code is None: + self.orchestrator.tokenizer.trust_remote_code = self.orchestrator.student.model.trust_remote_code else: # No shared tokenizer: re-derive from (now-correct) model names, # since auto_setup_tokenizer on sub-configs already ran with defaults. - for component in (self.trainer, self.orchestrator): - component.tokenizer.name = component.model.name - component.tokenizer.trust_remote_code = component.model.trust_remote_code + self.trainer.tokenizer.name = self.trainer.model.name + self.trainer.tokenizer.trust_remote_code = self.trainer.model.trust_remote_code + self.orchestrator.tokenizer.name = self.orchestrator.student.model.name + self.orchestrator.tokenizer.trust_remote_code = self.orchestrator.student.model.trust_remote_code # Propagate chat_template to inference (vLLM --chat-template) if self.inference is not None: @@ -705,40 +709,40 @@ def auto_setup_lora(self): if self.trainer.weight_broadcast.type == "nccl": raise ValueError("NCCL weight broadcast does not support LoRA yet.") - if self.orchestrator.model.lora is None: + if self.orchestrator.student.model.lora is None: from prime_rl.configs.orchestrator import LoRAConfig - self.orchestrator.model.lora = LoRAConfig() + self.orchestrator.student.model.lora = LoRAConfig() if ( - self.orchestrator.model.lora.rank is not None - and self.orchestrator.model.lora.rank != self.trainer.model.lora.rank + self.orchestrator.student.model.lora.rank is not None + and self.orchestrator.student.model.lora.rank != self.trainer.model.lora.rank ): raise ValueError( - f"orchestrator.model.lora.rank ({self.orchestrator.model.lora.rank}) conflicts with " + f"orchestrator.student.model.lora.rank ({self.orchestrator.student.model.lora.rank}) conflicts with " f"trainer.model.lora.rank ({self.trainer.model.lora.rank}). " - f"Remove orchestrator.model.lora.rank to inherit from trainer, or update trainer.model.lora.rank to match." + f"Remove orchestrator.student.model.lora.rank to inherit from trainer, or update trainer.model.lora.rank to match." ) if ( - self.orchestrator.model.lora.alpha is not None - and self.orchestrator.model.lora.alpha != self.trainer.model.lora.alpha + self.orchestrator.student.model.lora.alpha is not None + and self.orchestrator.student.model.lora.alpha != self.trainer.model.lora.alpha ): raise ValueError( - f"orchestrator.model.lora.alpha ({self.orchestrator.model.lora.alpha}) conflicts with " + f"orchestrator.student.model.lora.alpha ({self.orchestrator.student.model.lora.alpha}) conflicts with " f"trainer.model.lora.alpha ({self.trainer.model.lora.alpha}). " - f"Remove orchestrator.model.lora.alpha to inherit from trainer, or update trainer.model.lora.alpha to match." + f"Remove orchestrator.student.model.lora.alpha to inherit from trainer, or update trainer.model.lora.alpha to match." ) - if self.orchestrator.model.lora.rank is None: - self.orchestrator.model.lora.rank = self.trainer.model.lora.rank + if self.orchestrator.student.model.lora.rank is None: + self.orchestrator.student.model.lora.rank = self.trainer.model.lora.rank - if self.orchestrator.model.lora.alpha is None: - self.orchestrator.model.lora.alpha = self.trainer.model.lora.alpha + if self.orchestrator.student.model.lora.alpha is None: + self.orchestrator.student.model.lora.alpha = self.trainer.model.lora.alpha - if self.orchestrator.model.lora.name is None: - self.orchestrator.model.lora.name = ( - f"r{self.orchestrator.model.lora.rank}-a{self.orchestrator.model.lora.alpha}" + if self.orchestrator.student.model.lora.name is None: + self.orchestrator.student.model.lora.name = ( + f"r{self.orchestrator.student.model.lora.rank}-a{self.orchestrator.student.model.lora.alpha}" ) if self.inference is not None: @@ -756,7 +760,7 @@ def auto_setup_lora(self): @model_validator(mode="after") def auto_setup_session_headers(self): """Ensure X-Session-ID header is always set for sticky DP-aware routing at the inference router.""" - self.orchestrator.client.extra_headers_from_state.setdefault("X-Session-ID", "example_id") + self.orchestrator.student.client.extra_headers_from_state.setdefault("X-Session-ID", "example_id") return self @model_validator(mode="after") @@ -902,27 +906,22 @@ def auto_setup_disaggregated_inference(self): return self @model_validator(mode="after") - def auto_setup_dp_rank_count(self): - """Auto-set orchestrator client dp_rank_count from inference DP size. + def auto_setup_inference_client(self): + """Auto-configure orchestrator student client from the inference server config. - Uses data_parallel_size_local (per-node DP) when set, since each base URL - points to a single node whose API server only knows about its local ranks. - Falls back to the global parallel.dp for single-node setups. + For all modes, sets dp_rank_count from inference DP size. For SFT mode, + also sets base_url so setup_external_rollout_model can detect via + model_fields_set whether the student inference server is actually configured. """ - if self.inference is not None and "dp_rank_count" not in self.orchestrator.client.model_fields_set: - self.orchestrator.client.dp_rank_count = ( - self.inference.data_parallel_size_local or self.inference.parallel.dp - ) - return self - - @model_validator(mode="after") - def auto_setup_teacher_rollout_weight_updates(self): - if ( - self.orchestrator.teacher_rollout_model is not None - and self.inference is not None - and "update_student_inference_weights" not in self.orchestrator.model_fields_set - ): - self.orchestrator.update_student_inference_weights = True + if self.inference is None: + return self + client = self.orchestrator.student.client + if "dp_rank_count" not in client.model_fields_set: + client.dp_rank_count = self.inference.data_parallel_size_local or self.inference.parallel.dp + if self.orchestrator.training_mode == "sft" and "base_url" not in client.model_fields_set: + host = self.inference.server.host or "localhost" + port = self.inference.server.port + client.base_url = [f"http://{host}:{port}/v1"] return self @model_validator(mode="after") @@ -935,7 +934,7 @@ def auto_setup_teacher_inference(self): import copy - from prime_rl.configs.orchestrator import TeacherModelConfig + from prime_rl.configs.orchestrator import RolloutModelConfig if self.teacher_inference is None: if self.inference is None: @@ -957,12 +956,12 @@ def auto_setup_teacher_inference(self): assert num_teacher_gpus > 0, "num_teacher_gpus cannot be zero" self.teacher_inference.parallel.dp = num_teacher_gpus // tp - if self.orchestrator.teacher_model is None: - self.orchestrator.teacher_model = TeacherModelConfig() + if self.orchestrator.teacher is None: + self.orchestrator.teacher = RolloutModelConfig() host = self.teacher_inference.server.host or "localhost" port = self.teacher_inference.server.port - self.orchestrator.teacher_model.client.base_url = [f"http://{host}:{port}/v1"] - self.orchestrator.teacher_model.model.name = self.teacher_inference.model.name + self.orchestrator.teacher.client.base_url = [f"http://{host}:{port}/v1"] + self.orchestrator.teacher.model.name = self.teacher_inference.model.name return self diff --git a/packages/prime-rl-configs/src/prime_rl/utils/validation.py b/packages/prime-rl-configs/src/prime_rl/utils/validation.py index 9912cb47d7..6ebd0d1267 100644 --- a/packages/prime-rl-configs/src/prime_rl/utils/validation.py +++ b/packages/prime-rl-configs/src/prime_rl/utils/validation.py @@ -36,18 +36,18 @@ def validate_shared_model_name( ) -> None: # Orchestrator must match inference (it queries the inference server) if inference is not None: - if inference.model.name != orchestrator.model.name: + if inference.model.name != orchestrator.student.model.name: raise ValueError( - f"Inference model name ({inference.model.name}) and orchestrator model name ({orchestrator.model.name}) are not the same. " + f"Inference model name ({inference.model.name}) and orchestrator model name ({orchestrator.student.model.name}) are not the same. " "The orchestrator queries the inference server and must use the same model name." ) return if trainer.model.name.startswith("Jackmin108/"): # The TT MoE models will have a different name on the orchestrator return - if trainer.model.name != orchestrator.model.name: + if trainer.model.name != orchestrator.student.model.name: raise ValueError( - f"Trainer model name ({trainer.model.name}) and orchestrator model name ({orchestrator.model.name}) are not the same. Please specify the same model name for both." + f"Trainer model name ({trainer.model.name}) and orchestrator model name ({orchestrator.student.model.name}) are not the same. Please specify the same model name for both." ) diff --git a/skills/config/SKILL.md b/skills/config/SKILL.md index 3b0b0cc2db..9bafc8b91a 100644 --- a/skills/config/SKILL.md +++ b/skills/config/SKILL.md @@ -155,9 +155,9 @@ If you wish to configure values of the default variant, you don't need to set th ### SFT hard distill override -For hosted multi-tenant runs where the trainer image's `trainer.loss.type` is fixed, the orchestrator exposes a per-run override that forces SFT loss on every micro-batch without rebuilding the trainer. Set `orchestrator.use_sft_loss = true` alongside `orchestrator.teacher_rollout_model`; both must be configured together (the orchestrator validator enforces this). The orchestrator stamps each `TrainingSample.sft_loss = True`, which the trainer's `compute_loss` honors by dispatching to `sft_loss_fn` per batch, independent of the trainer's configured default loss. +Set `orchestrator.training_mode = "sft"` and configure `orchestrator.teacher_model` with the teacher endpoint. The orchestrator stamps each `TrainingSample.sft_loss = True`, which the trainer's `compute_loss` honors by dispatching to `sft_loss_fn` per batch, independent of the trainer's configured default loss. -When hard distill also needs online evals or policy weight sync against the student model, point `orchestrator.client` at the student inference server and `orchestrator.teacher_rollout_model.client` at the teacher. In the RL entrypoint this is usually done by configuring `[inference]`, which starts the student inference server, enables `orchestrator.update_student_inference_weights`, and leaves `teacher_rollout_model` scoped to training rollout generation only. For externally started student inference, set `orchestrator.update_student_inference_weights = true` explicitly. If weight updates are not enabled, hard distill keeps the teacher-only rollout behavior and skips policy updates. +When SFT hard distill also needs online evals or policy weight sync against the student model, configure `[inference]` in the RL entrypoint — this starts the student inference server and auto-configures `orchestrator.model.client`, enabling student weight sync. For externally started student inference, set `orchestrator.model.client.base_url` explicitly. If the student client is not configured, SFT keeps teacher-only rollout behavior and skips student policy updates. ### RL rollout client defaults diff --git a/src/prime_rl/entrypoints/rl.py b/src/prime_rl/entrypoints/rl.py index e1d6aceeaf..4dec4ababe 100644 --- a/src/prime_rl/entrypoints/rl.py +++ b/src/prime_rl/entrypoints/rl.py @@ -153,16 +153,16 @@ def rl_local(config: RLConfig): check_gpus_available(all_gpu_ids) # Validate client port matches inference server port - if config.inference is not None and not config.orchestrator.client.is_elastic: + if config.inference is not None and not config.orchestrator.student.client.is_elastic: from urllib.parse import urlparse - base_url = config.orchestrator.client.base_url[0] + base_url = config.orchestrator.student.client.base_url[0] parsed = urlparse(base_url) client_port = parsed.port expected_port = config.inference.server.port if client_port != expected_port: raise ValueError( - f"orchestrator.client.base_url port ({client_port}) does not match " + f"orchestrator.student.client.base_url port ({client_port}) does not match " f"inference.server.port ({expected_port}). " f"Update the base_url to use port {expected_port} to match the inference server." ) @@ -215,14 +215,12 @@ def sigterm_handler(signum, frame): monitor_thread.start() monitor_threads.append(monitor_thread) else: - if config.orchestrator.teacher_rollout_model is None: + if config.orchestrator.training_mode != "sft": logger.warning( "No inference config specified, skipping starting inference server. Make sure your inference server is running." ) else: - logger.info( - "No inference config specified, using orchestrator.teacher_rollout_model for rollout generation." - ) + logger.info("No inference config specified, using teacher model for rollout generation (sft mode).") # Optionally, start teacher inference process if config.teacher_inference: @@ -230,7 +228,7 @@ def sigterm_handler(signum, frame): raise ValueError( "teacher_inference is configured but deployment.num_teacher_gpus is not set. " "Either set deployment.num_teacher_gpus to start a teacher inference server, " - "or omit teacher_inference and configure orchestrator.teacher_model to use an existing server." + "or omit teacher_inference and configure orchestrator.teacher to use an existing server." ) teacher_inference_cmd = ["inference", "@", (config_dir / TEACHER_INFERENCE_TOML).as_posix()] @@ -260,10 +258,10 @@ def sigterm_handler(signum, frame): monitor_threads.append(monitor_thread) elif ( config.trainer.loss.type == "default" and config.trainer.loss.teacher_tau > 0 - ) or config.orchestrator.teacher_model: + ) or config.orchestrator.teacher: logger.warning( "No teacher_inference config specified, skipping starting teacher inference server. " - "Is your teacher inference server running? Make sure orchestrator.teacher_model is configured." + "Is your teacher inference server running? Make sure orchestrator.teacher is configured." ) # Start orchestrator process diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index 2b5c539e91..826a814ffa 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -122,22 +122,23 @@ async def orchestrate(config: OrchestratorConfig): # Setup rollout inference pool (handles both static and elastic modes) rollout_client_config, rollout_model_name, enable_policy_updates = setup_external_rollout_model(config, logger) - # Setup teacher inference pool if configured - if config.teacher_model: + # Setup teacher inference pool (opd: logprob distillation) + if config.training_mode == "opd": + assert config.teacher is not None logger.info( - f"Initializing teacher inference pool (base_url={', '.join(config.teacher_model.client.base_url)}, " - f"model={config.teacher_model.model.name})" + f"Initializing teacher inference pool (base_url={', '.join(config.teacher.client.base_url)}, " + f"model={config.teacher.model.name})" ) teacher_inference_pool = await setup_inference_pool( - config.teacher_model.client, - model_name=config.teacher_model.model.name, + config.teacher.client, + model_name=config.teacher.model.name, train_client_type="openai_chat_completions", ) else: teacher_inference_pool = None # Check if this is a vision-language model (used throughout for VLM-specific paths) - is_vlm = config.model.vlm is not None + is_vlm = config.student.model.vlm is not None # Load tokenizer and processor (processor only for VLM models) logger.info(f"Initializing tokenizer ({config.tokenizer})") @@ -145,25 +146,25 @@ async def orchestrate(config: OrchestratorConfig): processor = None if is_vlm: - logger.info(f"Loading VLM processor for {config.model.name}") + logger.info(f"Loading VLM processor for {config.student.model.name}") processor = AutoProcessor.from_pretrained( - config.model.name, trust_remote_code=config.model.trust_remote_code, use_fast=True + config.student.model.name, trust_remote_code=config.student.model.trust_remote_code, use_fast=True ) - teacher_rollout_clients = None - teacher_rollout_model_name = None - use_teacher_rollout_override = config.teacher_rollout_model is not None and enable_policy_updates - if use_teacher_rollout_override: + teacher_clients = None + teacher_model_name = None + use_sft_override = config.training_mode == "sft" and enable_policy_updates + if use_sft_override: logger.info(f"Using teacher rollout override (MITO, model={rollout_model_name})") - teacher_rollout_clients = setup_clients( + teacher_clients = setup_clients( rollout_client_config, client_type="openai_chat_completions", ) - teacher_rollout_model_name = rollout_model_name + teacher_model_name = rollout_model_name renderer = None inference_pool = await setup_inference_pool( - config.client, - model_name=config.model.name, + config.student.client, + model_name=config.student.model.name, train_client_type="openai_chat_completions", eval_client_type="openai_chat_completions", ) @@ -256,22 +257,22 @@ async def orchestrate(config: OrchestratorConfig): train_envs=train_envs, buffer=buffer, inference_pool=inference_pool, - teacher_rollout_clients=teacher_rollout_clients, - teacher_rollout_model_name=teacher_rollout_model_name, + teacher_clients=teacher_clients, + teacher_model_name=teacher_model_name, max_inflight_rollouts=config.max_inflight_rollouts, max_async_level=config.max_async_level, max_off_policy_steps=config.max_off_policy_steps, strict_async_level=config.strict_async_level, tasks_per_minute=config.tasks_per_minute, enable_policy_updates=enable_policy_updates, - lora_name=config.model.lora.name if config.model.lora else None, + lora_name=config.student.model.lora.name if config.student.model.lora else None, config=config, ) - scheduler.model_name = config.model.name if use_teacher_rollout_override else rollout_model_name + scheduler.model_name = config.student.model.name if use_sft_override else rollout_model_name # Check health of the inference pool logger.info("Waiting for inference pool to be ready") - inference_model_name = config.model.name if use_teacher_rollout_override else rollout_model_name + inference_model_name = config.student.model.name if use_sft_override else rollout_model_name await inference_pool.wait_for_ready(inference_model_name) logger.success("Inference pool ready") @@ -281,10 +282,11 @@ async def orchestrate(config: OrchestratorConfig): inference_metrics_collector = InferenceMetricsCollector(inference_pool.admin_clients) await inference_metrics_collector.start() - # Check health of teacher inference server if configured - if config.teacher_model and teacher_inference_pool: + # Check health of teacher inference server if configured (opd mode) + if config.training_mode == "opd" and teacher_inference_pool: + assert config.teacher is not None logger.info("Waiting for teacher inference pool to be ready") - await teacher_inference_pool.wait_for_ready(config.teacher_model.model.name) + await teacher_inference_pool.wait_for_ready(config.teacher.model.name) logger.success("Teacher inference pool ready") # Set up weight broadcast backend @@ -333,7 +335,7 @@ async def orchestrate(config: OrchestratorConfig): weights_path = get_weight_dir( config.output_dir, scheduler.ckpt_step, check_exists=check_exists, wait_timeout=wait_timeout ) - lora_name = config.model.lora.name if config.model.lora else None + lora_name = config.student.model.lora.name if config.student.model.lora else None await inference_pool.update_weights(weights_path, lora_name=lora_name, step=scheduler.ckpt_step) if lora_name is not None: inference_pool.update_model_name(lora_name) @@ -573,7 +575,7 @@ def process_rollout(rollout: vf.RolloutOutput, rollout_idx: int) -> list[Trainin for sample in samples: sample.advantage = rollout["advantage"] sample.reward = rollout["reward"] - if config.use_sft_loss: + if config.training_mode == "sft": sample.sft_loss = True sample_decode_tokens = sum(sample.completion_mask) sample_prefill_tokens = len(sample.prompt_ids) + len(sample.completion_mask) - sample_decode_tokens @@ -594,12 +596,12 @@ def process_rollout(rollout: vf.RolloutOutput, rollout_idx: int) -> list[Trainin # Compute teacher logprobs if teacher model is configured teacher_logprobs_time = 0 - if config.teacher_model and teacher_inference_pool: + if config.teacher and teacher_inference_pool: logger.info(f"Computing teacher logprobs for {len(train_examples)} training examples") teacher_logprobs_start_time = time.perf_counter() teacher_logprobs_list = await compute_teacher_logprobs( clients=teacher_inference_pool.train_clients, - model_name=config.teacher_model.model.name, + model_name=config.teacher.model.name, samples=train_examples, ) for train_example, teacher_logprobs in zip(train_examples, teacher_logprobs_list): @@ -939,7 +941,7 @@ async def setup_rollout_inference_pool( - both False → MITO (``openai_chat_completions``). VLMs land here too. """ - if config.teacher_rollout_model is not None: + if config.training_mode == "sft": logger.info("Using external rollout model (MITO) without renderer client") inference_pool = await setup_inference_pool( rollout_client_config, @@ -958,7 +960,7 @@ async def setup_rollout_inference_pool( preserve_all_thinking=config.renderer.preserve_all_thinking, preserve_thinking_between_tool_calls=config.renderer.preserve_thinking_between_tool_calls, ) - logger.info(f"Initialized {type(renderer).__name__} for {config.model.name}") + logger.info(f"Initialized {type(renderer).__name__} for {config.student.model.name}") inference_pool = await setup_inference_pool( rollout_client_config, model_name=rollout_model_name, diff --git a/src/prime_rl/orchestrator/scheduler.py b/src/prime_rl/orchestrator/scheduler.py index 321e300a4a..edc4cfe9e0 100644 --- a/src/prime_rl/orchestrator/scheduler.py +++ b/src/prime_rl/orchestrator/scheduler.py @@ -83,8 +83,8 @@ def __init__( tasks_per_minute: int | None, enable_policy_updates: bool = True, lora_name: str | None = None, - teacher_rollout_clients: list[vf.ClientConfig] | None = None, - teacher_rollout_model_name: str | None = None, + teacher_clients: list[vf.ClientConfig] | None = None, + teacher_model_name: str | None = None, ): self.logger = get_logger() if tasks_per_minute is not None: @@ -103,12 +103,12 @@ def __init__( self.strict_async_level = strict_async_level self.enable_policy_updates = enable_policy_updates self.lora_name = lora_name - self.model_name = self.config.model.name + self.model_name = self.config.student.model.name self.json_logging = config.log.json_logging self.inference_pool = inference_pool - self.teacher_rollout_clients = teacher_rollout_clients - self.teacher_rollout_model_name = teacher_rollout_model_name + self.teacher_clients = teacher_clients + self.teacher_model_name = teacher_model_name group_scoring_envs = [env.name for env in train_envs if env.requires_group_scoring] if group_scoring_envs: @@ -186,15 +186,15 @@ async def _select_least_loaded_client(self) -> vf.ClientConfig: return min(clients, key=lambda c: inflight[self._client_identity(c)]) def _resolve_rollout_request_target(self, client_config: vf.ClientConfig) -> tuple[vf.ClientConfig, str]: - if self.teacher_rollout_clients is None: + if self.teacher_clients is None: return client_config, self.model_name # The scheduler pins/load-balances against the student inference pool. # Map that selected logical client onto the teacher client set, which may # have a different size due to different base_url or dp_rank_count settings. - teacher_client = self.teacher_rollout_clients[client_config.client_idx % len(self.teacher_rollout_clients)] - assert self.teacher_rollout_model_name is not None - return teacher_client, self.teacher_rollout_model_name + teacher_client = self.teacher_clients[client_config.client_idx % len(self.teacher_clients)] + assert self.teacher_model_name is not None + return teacher_client, self.teacher_model_name async def drop_group(self, group_id: int) -> int: """Drop a group and cancel any remaining in-flight rollouts for it. Returns the number of cancelled rollouts.""" diff --git a/src/prime_rl/orchestrator/utils.py b/src/prime_rl/orchestrator/utils.py index 219b9c102b..4096637117 100644 --- a/src/prime_rl/orchestrator/utils.py +++ b/src/prime_rl/orchestrator/utils.py @@ -180,21 +180,20 @@ def find_stable_dir() -> Path | None: def setup_external_rollout_model(config: OrchestratorConfig, logger) -> tuple[Any, str, bool]: - """Resolve rollout client/model and whether policy updates should be enabled.""" - rollout_client_config = config.client - rollout_model_name = config.model.name - enable_policy_updates = ( - config.update_student_inference_weights if config.update_student_inference_weights is not None else True - ) - - if config.teacher_rollout_model is not None: - rollout_client_config = config.teacher_rollout_model.client - rollout_model_name = config.teacher_rollout_model.model.name - if config.update_student_inference_weights is None: - enable_policy_updates = False - logger.info( - f"Using external teacher rollout model (base_url={', '.join(rollout_client_config.base_url)}, " - f"model={rollout_model_name})" - ) + """Resolve rollout client/model and whether student policy updates are enabled. + - rl/opd: student generates rollouts, policy updates always enabled. + - sft: teacher generates rollouts; policy updates enabled iff the student + inference client is configured (non-empty base_url). + """ + if config.training_mode in ("rl", "opd"): + return config.student.client, config.student.model.name, True + + assert config.teacher is not None # validated by validate_training_mode + rollout_client_config = config.teacher.client + rollout_model_name = config.teacher.model.name + enable_policy_updates = "base_url" in config.student.client.model_fields_set + logger.info( + f"Using teacher rollout model (base_url={', '.join(rollout_client_config.base_url)}, model={rollout_model_name})" + ) return rollout_client_config, rollout_model_name, enable_policy_updates diff --git a/src/prime_rl/trainer/rl/broadcast/filesystem.py b/src/prime_rl/trainer/rl/broadcast/filesystem.py index 55a92c832d..e4ce958c92 100644 --- a/src/prime_rl/trainer/rl/broadcast/filesystem.py +++ b/src/prime_rl/trainer/rl/broadcast/filesystem.py @@ -78,7 +78,7 @@ def broadcast_weights(self, model: nn.Module, step: int) -> None: self.logger.debug(f"Saving weights for run {idx} to {save_dir}") save_state_dict(state_dict, save_dir, self.save_format, self.save_sharded, adapter=adapter_only) if adapter_only: - orch_lora = self.multi_run_manager.config[idx].model.lora + orch_lora = self.multi_run_manager.config[idx].student.model.lora save_lora_config( model, save_dir, diff --git a/src/prime_rl/trainer/runs.py b/src/prime_rl/trainer/runs.py index 9372386764..b01e19f044 100644 --- a/src/prime_rl/trainer/runs.py +++ b/src/prime_rl/trainer/runs.py @@ -511,19 +511,21 @@ def setup_multi_run_manager( def validate_lora_rank(orch_config: "OrchestratorConfig") -> tuple[bool, str]: # Default to trainer's rank/alpha if not specified - if orch_config.model.lora.rank is None: - orch_config.model.lora.rank = trainer_lora.rank - if orch_config.model.lora.alpha is None: - orch_config.model.lora.alpha = trainer_lora.alpha - if orch_config.model.lora.rank > trainer_lora.rank: + if orch_config.student.model.lora.rank is None: + orch_config.student.model.lora.rank = trainer_lora.rank + if orch_config.student.model.lora.alpha is None: + orch_config.student.model.lora.alpha = trainer_lora.alpha + if orch_config.student.model.lora.rank > trainer_lora.rank: return ( False, - f"model.lora.rank ({orch_config.model.lora.rank}) exceeds trainer max rank ({trainer_lora.rank})", + f"student.model.lora.rank ({orch_config.student.model.lora.rank}) exceeds trainer max rank ({trainer_lora.rank})", ) return True, "" def on_run_discovered(idx: int, run_id: str, orch_config: "OrchestratorConfig") -> None: - _MULTI_RUN_MANAGER.scaling_factors[idx] = orch_config.model.lora.alpha / orch_config.model.lora.rank + _MULTI_RUN_MANAGER.scaling_factors[idx] = ( + orch_config.student.model.lora.alpha / orch_config.student.model.lora.rank + ) _MULTI_RUN_MANAGER.register_config_validation_hook(validate_lora_rank) _MULTI_RUN_MANAGER.register_discovered_hook(on_run_discovered) diff --git a/tests/unit/orchestrator/test_orchestrator_setup.py b/tests/unit/orchestrator/test_orchestrator_setup.py index 2bfe050743..d6c3f59d3e 100644 --- a/tests/unit/orchestrator/test_orchestrator_setup.py +++ b/tests/unit/orchestrator/test_orchestrator_setup.py @@ -6,12 +6,12 @@ from prime_rl.orchestrator.utils import setup_external_rollout_model -def test_setup_rollout_inference_pool_uses_plain_client_for_external_teacher_rollout(): +def test_setup_rollout_inference_pool_uses_plain_client_for_sft_mode(): async def run() -> None: tokenizer = object() config = SimpleNamespace( - teacher_rollout_model=SimpleNamespace(), - model=SimpleNamespace(renderer="auto", name="student-model"), + training_mode="sft", + student=SimpleNamespace(renderer="auto", model=SimpleNamespace(name="student-model")), ) rollout_client_config = SimpleNamespace(base_url=["https://api.pinference.ai/api/v1"]) logger = MagicMock() @@ -45,26 +45,26 @@ async def run() -> None: asyncio.run(run()) -def test_setup_external_rollout_model_uses_explicit_weight_update_opt_in(): +def test_setup_external_rollout_model_sft_uses_teacher_and_checks_student_client(): + from prime_rl.configs.orchestrator import ClientConfig, RolloutModelConfig + teacher_client = SimpleNamespace(base_url=["https://teacher.example/v1"]) - student_client = SimpleNamespace(base_url=["http://localhost:8000/v1"]) - config = SimpleNamespace( - client=student_client, - model=SimpleNamespace(name="student-model"), - update_student_inference_weights=None, - teacher_rollout_model=SimpleNamespace( - client=teacher_client, - model=SimpleNamespace(name="teacher-model"), - ), - ) logger = MagicMock() + # SFT mode, student client has default base_url (not in model_fields_set) → policy updates disabled + config = SimpleNamespace( + training_mode="sft", + student=RolloutModelConfig(), # default client — base_url not explicitly set + teacher=SimpleNamespace(client=teacher_client, model=SimpleNamespace(name="teacher-model")), + ) rollout_client, rollout_model, enable_policy_updates = setup_external_rollout_model(config, logger) assert rollout_client is teacher_client assert rollout_model == "teacher-model" assert not enable_policy_updates - config.update_student_inference_weights = True + # SFT mode, student client base_url explicitly set → policy updates enabled + student_model = RolloutModelConfig(client=ClientConfig(base_url=["http://localhost:8000/v1"])) + config.student = student_model rollout_client, rollout_model, enable_policy_updates = setup_external_rollout_model(config, logger) assert rollout_client is teacher_client assert rollout_model == "teacher-model" @@ -75,10 +75,10 @@ def test_setup_rollout_inference_pool_uses_direct_renderer_client_for_local_vllm async def run() -> None: tokenizer = object() config = SimpleNamespace( - teacher_rollout_model=None, + training_mode="rl", use_renderer=True, use_token_client=False, - model=SimpleNamespace(name="student-model"), + student=SimpleNamespace(model=SimpleNamespace(name="student-model")), renderer=SimpleNamespace( name="qwen3_vl", tool_parser=None, diff --git a/tests/unit/orchestrator/test_scheduler.py b/tests/unit/orchestrator/test_scheduler.py index 3af71616b7..3832a99c4d 100644 --- a/tests/unit/orchestrator/test_scheduler.py +++ b/tests/unit/orchestrator/test_scheduler.py @@ -32,8 +32,8 @@ def make_scheduler() -> Scheduler: scheduler.update_policy_task = None scheduler.enable_policy_updates = True scheduler.rate_limiter = None - scheduler.teacher_rollout_clients = None - scheduler.teacher_rollout_model_name = None + scheduler.teacher_clients = None + scheduler.teacher_model_name = None return scheduler @@ -184,7 +184,7 @@ async def run() -> None: scheduler = make_scheduler() scheduler.model_name = "student-model" scheduler.lora_name = "student-lora" - scheduler.teacher_rollout_model_name = "teacher-model" + scheduler.teacher_model_name = "teacher-model" scheduler.inference_pool = SimpleNamespace( update_weights=AsyncMock(), @@ -201,7 +201,7 @@ async def run() -> None: scheduler.inference_pool.update_weights.assert_awaited_once() scheduler.inference_pool.update_model_name.assert_called_once_with("student-lora") assert scheduler.model_name == "student-lora" - assert scheduler.teacher_rollout_model_name == "teacher-model" + assert scheduler.teacher_model_name == "teacher-model" asyncio.run(run()) @@ -216,8 +216,8 @@ async def run() -> None: run_rollout=AsyncMock(return_value=[]), ) scheduler.inference_pool = SimpleNamespace(train_clients=[student_client]) - scheduler.teacher_rollout_clients = [teacher_client] - scheduler.teacher_rollout_model_name = "teacher-model" + scheduler.teacher_clients = [teacher_client] + scheduler.teacher_model_name = "teacher-model" scheduler.train_envs = SimpleNamespace(get=MagicMock(return_value=env)) scheduler.groups = { 0: GroupState( diff --git a/tests/unit/test_configs.py b/tests/unit/test_configs.py index 863622b645..ae3d69ea1c 100644 --- a/tests/unit/test_configs.py +++ b/tests/unit/test_configs.py @@ -161,10 +161,12 @@ def test_orchestrator_vlm_configs_must_disable_renderer(): with pytest.raises(ValidationError, match="orchestrator.use_renderer is not supported for VLMs"): OrchestratorConfig.model_validate( { - "model": { - "vlm": { - "vision_encoder_attr": "model.visual", - "language_model_attr": "model.language_model", + "student": { + "model": { + "vlm": { + "vision_encoder_attr": "model.visual", + "language_model_attr": "model.language_model", + } } } } @@ -172,10 +174,12 @@ def test_orchestrator_vlm_configs_must_disable_renderer(): config = OrchestratorConfig.model_validate( { - "model": { - "vlm": { - "vision_encoder_attr": "model.visual", - "language_model_attr": "model.language_model", + "student": { + "model": { + "vlm": { + "vision_encoder_attr": "model.visual", + "language_model_attr": "model.language_model", + } } }, "use_token_client": False, @@ -192,32 +196,24 @@ def test_selective_activation_checkpointing_requires_custom_impl(): TrainerModelConfig.model_validate({"impl": "hf", "ac": {"mode": "selective"}}) -def test_teacher_rollout_weight_updates_require_inference_config_or_explicit_opt_in(): +def test_sft_training_mode_enables_student_pool_when_inference_configured(): base_config = { "trainer": {}, "orchestrator": { - "use_sft_loss": True, + "training_mode": "sft", "use_token_client": False, "use_renderer": False, - "teacher_rollout_model": { + "teacher": { "client": {"base_url": ["http://teacher.example/v1"]}, "model": {"name": "teacher-model"}, }, }, } + # Without inference, student client base_url not explicitly set → policy updates disabled config = RLConfig.model_validate(base_config) - assert config.orchestrator.update_student_inference_weights is None + assert "base_url" not in config.orchestrator.student.client.model_fields_set + # With inference, student client base_url is auto-set → policy updates enabled config = RLConfig.model_validate({**base_config, "inference": {}}) - assert config.orchestrator.update_student_inference_weights - - explicit_config = { - **base_config, - "orchestrator": { - **base_config["orchestrator"], - "update_student_inference_weights": True, - }, - } - config = RLConfig.model_validate(explicit_config) - assert config.orchestrator.update_student_inference_weights + assert "base_url" in config.orchestrator.student.client.model_fields_set From 0f029c6611a95261e2557567895209346d0639a9 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Mon, 18 May 2026 21:37:10 +0000 Subject: [PATCH 08/47] chore(configs): cap reverse-text debug eval at 64 examples Co-Authored-By: Claude Sonnet 4.6 --- configs/reverse_text/debug_opd.toml | 1 + configs/reverse_text/debug_rl.toml | 1 + configs/reverse_text/debug_sft.toml | 1 + 3 files changed, 3 insertions(+) diff --git a/configs/reverse_text/debug_opd.toml b/configs/reverse_text/debug_opd.toml index 15aa8d75bb..afb68f6a5f 100644 --- a/configs/reverse_text/debug_opd.toml +++ b/configs/reverse_text/debug_opd.toml @@ -21,6 +21,7 @@ id = "reverse-text" [orchestrator.eval] interval = 5 +num_examples = 64 [[orchestrator.eval.env]] id = "reverse-text" diff --git a/configs/reverse_text/debug_rl.toml b/configs/reverse_text/debug_rl.toml index 762d3fed35..00f7a7a258 100644 --- a/configs/reverse_text/debug_rl.toml +++ b/configs/reverse_text/debug_rl.toml @@ -20,6 +20,7 @@ id = "reverse-text" [orchestrator.eval] interval = 5 +num_examples = 64 [[orchestrator.eval.env]] id = "reverse-text" diff --git a/configs/reverse_text/debug_sft.toml b/configs/reverse_text/debug_sft.toml index d67f760c2a..c630164adf 100644 --- a/configs/reverse_text/debug_sft.toml +++ b/configs/reverse_text/debug_sft.toml @@ -23,6 +23,7 @@ id = "reverse-text" [orchestrator.eval] interval = 5 +num_examples = 64 [[orchestrator.eval.env]] id = "reverse-text" From 6eb239d0d0c04e3099c641bba8c38196cb9a3c03 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Mon, 18 May 2026 21:38:38 +0000 Subject: [PATCH 09/47] chore(configs): cap reverse-text debug eval completion tokens at 128 Co-Authored-By: Claude Sonnet 4.6 --- configs/reverse_text/debug_opd.toml | 3 +++ configs/reverse_text/debug_rl.toml | 3 +++ configs/reverse_text/debug_sft.toml | 3 +++ 3 files changed, 9 insertions(+) diff --git a/configs/reverse_text/debug_opd.toml b/configs/reverse_text/debug_opd.toml index afb68f6a5f..7b9f3c37aa 100644 --- a/configs/reverse_text/debug_opd.toml +++ b/configs/reverse_text/debug_opd.toml @@ -23,6 +23,9 @@ id = "reverse-text" interval = 5 num_examples = 64 +[orchestrator.eval.sampling] +max_completion_tokens = 128 + [[orchestrator.eval.env]] id = "reverse-text" diff --git a/configs/reverse_text/debug_rl.toml b/configs/reverse_text/debug_rl.toml index 00f7a7a258..35aa2bafcb 100644 --- a/configs/reverse_text/debug_rl.toml +++ b/configs/reverse_text/debug_rl.toml @@ -22,6 +22,9 @@ id = "reverse-text" interval = 5 num_examples = 64 +[orchestrator.eval.sampling] +max_completion_tokens = 128 + [[orchestrator.eval.env]] id = "reverse-text" diff --git a/configs/reverse_text/debug_sft.toml b/configs/reverse_text/debug_sft.toml index c630164adf..9843797a09 100644 --- a/configs/reverse_text/debug_sft.toml +++ b/configs/reverse_text/debug_sft.toml @@ -25,6 +25,9 @@ id = "reverse-text" interval = 5 num_examples = 64 +[orchestrator.eval.sampling] +max_completion_tokens = 128 + [[orchestrator.eval.env]] id = "reverse-text" From c30c9155b3ef182b4a37c26ffe623f1bf0a30916 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Mon, 18 May 2026 21:46:16 +0000 Subject: [PATCH 10/47] fix(configs): use correct PI inference teacher model name qwen/qwen3-4b-instruct doesn't exist on PI inference; use Qwen/Qwen3-4B-Instruct-2507 instead. Co-Authored-By: Claude Sonnet 4.6 --- configs/reverse_text/debug_opd.toml | 2 +- configs/reverse_text/debug_sft.toml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/configs/reverse_text/debug_opd.toml b/configs/reverse_text/debug_opd.toml index 7b9f3c37aa..c191c0d76e 100644 --- a/configs/reverse_text/debug_opd.toml +++ b/configs/reverse_text/debug_opd.toml @@ -34,7 +34,7 @@ base_url = ["https://api.pinference.ai/api/v1"] api_key_var = "PRIME_API_KEY" [orchestrator.teacher.model] -name = "qwen/qwen3-4b-instruct" +name = "Qwen/Qwen3-4B-Instruct-2507" [trainer.optim] lr = 3e-6 diff --git a/configs/reverse_text/debug_sft.toml b/configs/reverse_text/debug_sft.toml index 9843797a09..5901ca5963 100644 --- a/configs/reverse_text/debug_sft.toml +++ b/configs/reverse_text/debug_sft.toml @@ -36,7 +36,7 @@ base_url = ["https://api.pinference.ai/api/v1"] api_key_var = "PRIME_API_KEY" [orchestrator.teacher.model] -name = "qwen/qwen3-4b-instruct" +name = "Qwen/Qwen3-4B-Instruct-2507" [trainer.optim] lr = 3e-6 From 3ff74a1cec7cf0ff7637e5c7888288fe6e3ce59e Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Mon, 18 May 2026 21:46:43 +0000 Subject: [PATCH 11/47] fix(configs): use qwen3-30b-a3b-instruct as debug teacher model Co-Authored-By: Claude Sonnet 4.6 --- configs/reverse_text/debug_opd.toml | 2 +- configs/reverse_text/debug_sft.toml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/configs/reverse_text/debug_opd.toml b/configs/reverse_text/debug_opd.toml index c191c0d76e..0096014340 100644 --- a/configs/reverse_text/debug_opd.toml +++ b/configs/reverse_text/debug_opd.toml @@ -34,7 +34,7 @@ base_url = ["https://api.pinference.ai/api/v1"] api_key_var = "PRIME_API_KEY" [orchestrator.teacher.model] -name = "Qwen/Qwen3-4B-Instruct-2507" +name = "qwen/qwen3-30b-a3b-instruct" [trainer.optim] lr = 3e-6 diff --git a/configs/reverse_text/debug_sft.toml b/configs/reverse_text/debug_sft.toml index 5901ca5963..86a33ede7f 100644 --- a/configs/reverse_text/debug_sft.toml +++ b/configs/reverse_text/debug_sft.toml @@ -36,7 +36,7 @@ base_url = ["https://api.pinference.ai/api/v1"] api_key_var = "PRIME_API_KEY" [orchestrator.teacher.model] -name = "Qwen/Qwen3-4B-Instruct-2507" +name = "qwen/qwen3-30b-a3b-instruct" [trainer.optim] lr = 3e-6 From 99183387ed9728f54935981d784f32c6ac25c99a Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Mon, 18 May 2026 21:47:53 +0000 Subject: [PATCH 12/47] fix(configs): correct teacher model name to qwen3-30b-a3b-instruct-2507 Co-Authored-By: Claude Sonnet 4.6 --- configs/reverse_text/debug_opd.toml | 2 +- configs/reverse_text/debug_sft.toml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/configs/reverse_text/debug_opd.toml b/configs/reverse_text/debug_opd.toml index 0096014340..22f45b8574 100644 --- a/configs/reverse_text/debug_opd.toml +++ b/configs/reverse_text/debug_opd.toml @@ -34,7 +34,7 @@ base_url = ["https://api.pinference.ai/api/v1"] api_key_var = "PRIME_API_KEY" [orchestrator.teacher.model] -name = "qwen/qwen3-30b-a3b-instruct" +name = "qwen/qwen3-30b-a3b-instruct-2507" [trainer.optim] lr = 3e-6 diff --git a/configs/reverse_text/debug_sft.toml b/configs/reverse_text/debug_sft.toml index 86a33ede7f..8481fcd347 100644 --- a/configs/reverse_text/debug_sft.toml +++ b/configs/reverse_text/debug_sft.toml @@ -36,7 +36,7 @@ base_url = ["https://api.pinference.ai/api/v1"] api_key_var = "PRIME_API_KEY" [orchestrator.teacher.model] -name = "qwen/qwen3-30b-a3b-instruct" +name = "qwen/qwen3-30b-a3b-instruct-2507" [trainer.optim] lr = 3e-6 From 18cad301e19eef46164edee6333594e1a2e923fe Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Mon, 18 May 2026 21:50:25 +0000 Subject: [PATCH 13/47] feat(client): add headers_from_env to ClientConfig Maps header names to env var names, resolved at client setup time. Analogous to api_key_var but for arbitrary headers. Useful for e.g. X-Prime-Team-ID when hitting the PI inference API. Co-Authored-By: Claude Sonnet 4.6 --- configs/reverse_text/debug_opd.toml | 3 +++ configs/reverse_text/debug_sft.toml | 3 +++ .../prime-rl-configs/src/prime_rl/configs/shared.py | 7 +++++++ src/prime_rl/utils/client.py | 10 ++++++++-- src/prime_rl/utils/elastic.py | 2 ++ 5 files changed, 23 insertions(+), 2 deletions(-) diff --git a/configs/reverse_text/debug_opd.toml b/configs/reverse_text/debug_opd.toml index 22f45b8574..fd425eefb5 100644 --- a/configs/reverse_text/debug_opd.toml +++ b/configs/reverse_text/debug_opd.toml @@ -33,6 +33,9 @@ id = "reverse-text" base_url = ["https://api.pinference.ai/api/v1"] api_key_var = "PRIME_API_KEY" +[orchestrator.teacher.client.headers_from_env] +X-Prime-Team-ID = "PRIME_TEAM_ID" + [orchestrator.teacher.model] name = "qwen/qwen3-30b-a3b-instruct-2507" diff --git a/configs/reverse_text/debug_sft.toml b/configs/reverse_text/debug_sft.toml index 8481fcd347..8f39d65ba6 100644 --- a/configs/reverse_text/debug_sft.toml +++ b/configs/reverse_text/debug_sft.toml @@ -35,6 +35,9 @@ id = "reverse-text" base_url = ["https://api.pinference.ai/api/v1"] api_key_var = "PRIME_API_KEY" +[orchestrator.teacher.client.headers_from_env] +X-Prime-Team-ID = "PRIME_TEAM_ID" + [orchestrator.teacher.model] name = "qwen/qwen3-30b-a3b-instruct-2507" diff --git a/packages/prime-rl-configs/src/prime_rl/configs/shared.py b/packages/prime-rl-configs/src/prime_rl/configs/shared.py index 962b7abbff..5f871fcd53 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/shared.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/shared.py @@ -299,6 +299,13 @@ class ClientConfig(BaseConfig): ), ] = {} + headers_from_env: Annotated[ + dict[str, str], + Field( + description='Maps HTTP header names to environment variable names. At runtime each entry is resolved via os.getenv and merged into the request headers. e.g. {"X-Prime-Team-ID": "PRIME_TEAM_ID"}.', + ), + ] = {} + extra_headers_from_state: Annotated[ dict[str, str], Field( diff --git a/src/prime_rl/utils/client.py b/src/prime_rl/utils/client.py index 4b7942348d..511c32fff3 100644 --- a/src/prime_rl/utils/client.py +++ b/src/prime_rl/utils/client.py @@ -208,9 +208,12 @@ def setup_clients( "preserve_all_thinking": preserve_all_thinking, "preserve_thinking_between_tool_calls": preserve_thinking_between_tool_calls, } + env_headers = { + k: v for k, v in ((k, os.getenv(v)) for k, v in client_config.headers_from_env.items()) if v is not None + } for base_url in client_config.base_url: for dp_rank in range(client_config.dp_rank_count): - headers = client_config.headers.copy() + headers = {**client_config.headers, **env_headers} if client_config.dp_rank_count > 1: headers["X-data-parallel-rank"] = str(dp_rank) clients.append( @@ -248,7 +251,10 @@ def setup_admin_clients(client_config: ClientConfig) -> list[AsyncClient]: urls = client_config.admin_base_url if client_config.admin_base_url else client_config.base_url def _setup_admin_client(base_url: str) -> httpx.AsyncClient: - headers = client_config.headers.copy() # avoid mutating config + env_headers = { + k: v for k, v in ((k, os.getenv(v)) for k, v in client_config.headers_from_env.items()) if v is not None + } + headers = {**client_config.headers, **env_headers} api_key = os.getenv(client_config.api_key_var, "EMPTY") if api_key and api_key != "EMPTY": headers["Authorization"] = f"Bearer {api_key}" diff --git a/src/prime_rl/utils/elastic.py b/src/prime_rl/utils/elastic.py index c59f81e27f..cef2bcb012 100644 --- a/src/prime_rl/utils/elastic.py +++ b/src/prime_rl/utils/elastic.py @@ -210,6 +210,7 @@ def _rebuild_clients(self) -> None: base_url=urls, api_key_var=self.client_config.api_key_var, headers=self.client_config.headers, + headers_from_env=self.client_config.headers_from_env, dp_rank_count=self.client_config.dp_rank_count, extra_headers_from_state=self.client_config.extra_headers_from_state, ) @@ -267,6 +268,7 @@ async def _create_admin_client(self, ip: str) -> AsyncClient: base_url=[f"{url}/v1"], api_key_var=self.client_config.api_key_var, headers=self.client_config.headers, + headers_from_env=self.client_config.headers_from_env, ) return setup_admin_clients(config)[0] From f11b8b1ee6a27265d561696695a6fc979758f90d Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Mon, 18 May 2026 21:53:02 +0000 Subject: [PATCH 14/47] feat(client): auto-inject X-Prime-Team-ID header for PI inference When ClientConfig.base_url targets pinference.ai, auto-add a headers_from_env mapping for X-Prime-Team-ID -> PRIME_TEAM_ID so team billing works without each config repeating the boilerplate. Co-Authored-By: Claude Sonnet 4.6 --- configs/reverse_text/debug_opd.toml | 3 --- configs/reverse_text/debug_sft.toml | 3 --- packages/prime-rl-configs/src/prime_rl/configs/shared.py | 6 ++++++ 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/configs/reverse_text/debug_opd.toml b/configs/reverse_text/debug_opd.toml index fd425eefb5..22f45b8574 100644 --- a/configs/reverse_text/debug_opd.toml +++ b/configs/reverse_text/debug_opd.toml @@ -33,9 +33,6 @@ id = "reverse-text" base_url = ["https://api.pinference.ai/api/v1"] api_key_var = "PRIME_API_KEY" -[orchestrator.teacher.client.headers_from_env] -X-Prime-Team-ID = "PRIME_TEAM_ID" - [orchestrator.teacher.model] name = "qwen/qwen3-30b-a3b-instruct-2507" diff --git a/configs/reverse_text/debug_sft.toml b/configs/reverse_text/debug_sft.toml index 8f39d65ba6..8481fcd347 100644 --- a/configs/reverse_text/debug_sft.toml +++ b/configs/reverse_text/debug_sft.toml @@ -35,9 +35,6 @@ id = "reverse-text" base_url = ["https://api.pinference.ai/api/v1"] api_key_var = "PRIME_API_KEY" -[orchestrator.teacher.client.headers_from_env] -X-Prime-Team-ID = "PRIME_TEAM_ID" - [orchestrator.teacher.model] name = "qwen/qwen3-30b-a3b-instruct-2507" diff --git a/packages/prime-rl-configs/src/prime_rl/configs/shared.py b/packages/prime-rl-configs/src/prime_rl/configs/shared.py index 5f871fcd53..82cc784408 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/shared.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/shared.py @@ -362,6 +362,12 @@ class ClientConfig(BaseConfig): ), ] = None + @model_validator(mode="after") + def auto_setup_pinference_team_header(self): + if any("pinference.ai" in url for url in self.base_url): + self.headers_from_env.setdefault("X-Prime-Team-ID", "PRIME_TEAM_ID") + return self + @property def is_elastic(self) -> bool: """Check if elastic mode is enabled.""" From 90f50d916ac29e42df19292f2c7aa9f64b0024fc Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Mon, 18 May 2026 21:54:21 +0000 Subject: [PATCH 15/47] chore(configs): set debug eval to 256 rollouts (64 examples x 4) Co-Authored-By: Claude Sonnet 4.6 --- configs/reverse_text/debug_opd.toml | 1 + configs/reverse_text/debug_rl.toml | 1 + configs/reverse_text/debug_sft.toml | 1 + 3 files changed, 3 insertions(+) diff --git a/configs/reverse_text/debug_opd.toml b/configs/reverse_text/debug_opd.toml index 22f45b8574..109fbab97a 100644 --- a/configs/reverse_text/debug_opd.toml +++ b/configs/reverse_text/debug_opd.toml @@ -22,6 +22,7 @@ id = "reverse-text" [orchestrator.eval] interval = 5 num_examples = 64 +rollouts_per_example = 4 [orchestrator.eval.sampling] max_completion_tokens = 128 diff --git a/configs/reverse_text/debug_rl.toml b/configs/reverse_text/debug_rl.toml index 35aa2bafcb..0a391f29eb 100644 --- a/configs/reverse_text/debug_rl.toml +++ b/configs/reverse_text/debug_rl.toml @@ -21,6 +21,7 @@ id = "reverse-text" [orchestrator.eval] interval = 5 num_examples = 64 +rollouts_per_example = 4 [orchestrator.eval.sampling] max_completion_tokens = 128 diff --git a/configs/reverse_text/debug_sft.toml b/configs/reverse_text/debug_sft.toml index 8481fcd347..391378d72d 100644 --- a/configs/reverse_text/debug_sft.toml +++ b/configs/reverse_text/debug_sft.toml @@ -24,6 +24,7 @@ id = "reverse-text" [orchestrator.eval] interval = 5 num_examples = 64 +rollouts_per_example = 4 [orchestrator.eval.sampling] max_completion_tokens = 128 From 680f7e2c7dfd2599b1186178a41b58760d21df46 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Mon, 18 May 2026 22:25:31 +0000 Subject: [PATCH 16/47] fix(configs): use local vLLM teacher for debug OPD instead of PI inference PI inference is OpenAI-compatible but doesn't expose vLLM's /inference/v1/generate prefill endpoint that compute_teacher_logprobs needs - so OPD against PI inference 404s. SFT mode kept working because that path only hits /chat/completions. Switch debug_opd.toml to use the existing teacher_inference auto-setup: set num_teacher_gpus = 1, and an empty [orchestrator.teacher] block to satisfy the orchestrator validator (which runs before auto_setup_teacher_inference can populate it). The validator chain spins up a local vLLM teacher on port 8001 using the same reverse-text student model. Co-Authored-By: Claude Sonnet 4.6 --- configs/reverse_text/debug_opd.toml | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/configs/reverse_text/debug_opd.toml b/configs/reverse_text/debug_opd.toml index 109fbab97a..83d7b78707 100644 --- a/configs/reverse_text/debug_opd.toml +++ b/configs/reverse_text/debug_opd.toml @@ -8,6 +8,9 @@ name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT" project = "reverse-text-debug" name = "debug-opd" +[deployment] +num_teacher_gpus = 1 + [orchestrator] batch_size = 128 rollouts_per_example = 16 @@ -30,12 +33,7 @@ max_completion_tokens = 128 [[orchestrator.eval.env]] id = "reverse-text" -[orchestrator.teacher.client] -base_url = ["https://api.pinference.ai/api/v1"] -api_key_var = "PRIME_API_KEY" - -[orchestrator.teacher.model] -name = "qwen/qwen3-30b-a3b-instruct-2507" +[orchestrator.teacher] [trainer.optim] lr = 3e-6 From 5b94e1358221ee447c46da6badd83d4d947e5188 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Mon, 18 May 2026 22:43:18 +0000 Subject: [PATCH 17/47] docs+validate: training modes grid + cross-config constraints - Add docs/training_modes.md with the rl/opd/sft role grid, implications (e.g. OPD needs local vLLM teacher, SFT teacher is any OAI-compatible endpoint), and minimal config per mode. Linked from mint.json nav. - Add validate_training_mode_loss_consistency on RLConfig enforcing: sft mode <-> sft loss type, opd mode requires teacher_tau > 0, non-sft modes forbid sft loss type. These cross orchestrator and trainer fields so they live on RLConfig (per-component validators in OrchestratorConfig already cover the within-orchestrator constraints). - Fix debug_sft.toml to set trainer.loss.type = "sft" (was relying on default loss which is now explicitly rejected). - Make training_mode explicit in debug_rl.toml and debug_opd.toml. Co-Authored-By: Claude Sonnet 4.6 --- configs/reverse_text/debug_opd.toml | 7 +- configs/reverse_text/debug_rl.toml | 6 +- configs/reverse_text/debug_sft.toml | 15 ++-- docs/mint.json | 1 + docs/training_modes.md | 68 +++++++++++++++++++ .../src/prime_rl/configs/rl.py | 23 +++++++ 6 files changed, 106 insertions(+), 14 deletions(-) create mode 100644 docs/training_modes.md diff --git a/configs/reverse_text/debug_opd.toml b/configs/reverse_text/debug_opd.toml index 83d7b78707..85b09cbad9 100644 --- a/configs/reverse_text/debug_opd.toml +++ b/configs/reverse_text/debug_opd.toml @@ -12,9 +12,9 @@ name = "debug-opd" num_teacher_gpus = 1 [orchestrator] +training_mode = "opd" batch_size = 128 rollouts_per_example = 16 -training_mode = "opd" [orchestrator.train.sampling] max_completion_tokens = 128 @@ -23,9 +23,8 @@ max_completion_tokens = 128 id = "reverse-text" [orchestrator.eval] -interval = 5 -num_examples = 64 -rollouts_per_example = 4 +interval = 1 +num_examples = 128 [orchestrator.eval.sampling] max_completion_tokens = 128 diff --git a/configs/reverse_text/debug_rl.toml b/configs/reverse_text/debug_rl.toml index 0a391f29eb..10f6a9e906 100644 --- a/configs/reverse_text/debug_rl.toml +++ b/configs/reverse_text/debug_rl.toml @@ -9,6 +9,7 @@ project = "reverse-text-debug" name = "debug-rl" [orchestrator] +training_mode = "rl" batch_size = 128 rollouts_per_example = 16 @@ -19,9 +20,8 @@ max_completion_tokens = 128 id = "reverse-text" [orchestrator.eval] -interval = 5 -num_examples = 64 -rollouts_per_example = 4 +interval = 1 +num_examples = 128 [orchestrator.eval.sampling] max_completion_tokens = 128 diff --git a/configs/reverse_text/debug_sft.toml b/configs/reverse_text/debug_sft.toml index 391378d72d..cdda571428 100644 --- a/configs/reverse_text/debug_sft.toml +++ b/configs/reverse_text/debug_sft.toml @@ -9,11 +9,10 @@ project = "reverse-text-debug" name = "debug-sft" [orchestrator] +training_mode = "sft" batch_size = 128 rollouts_per_example = 4 -use_token_client = false use_renderer = false -training_mode = "sft" [orchestrator.train.sampling] max_completion_tokens = 128 @@ -22,9 +21,8 @@ max_completion_tokens = 128 id = "reverse-text" [orchestrator.eval] -interval = 5 -num_examples = 64 -rollouts_per_example = 4 +interval = 1 +num_examples = 128 [orchestrator.eval.sampling] max_completion_tokens = 128 @@ -32,12 +30,15 @@ max_completion_tokens = 128 [[orchestrator.eval.env]] id = "reverse-text" +[orchestrator.teacher.model] +name = "qwen/qwen3-30b-a3b-instruct-2507" + [orchestrator.teacher.client] base_url = ["https://api.pinference.ai/api/v1"] api_key_var = "PRIME_API_KEY" -[orchestrator.teacher.model] -name = "qwen/qwen3-30b-a3b-instruct-2507" +[trainer.loss] +type = "sft" [trainer.optim] lr = 3e-6 diff --git a/docs/mint.json b/docs/mint.json index 216fbe4fa4..8bb9daee2b 100644 --- a/docs/mint.json +++ b/docs/mint.json @@ -7,6 +7,7 @@ "index", "entrypoints", "configs", + "training_modes", "environments", "async", "logging", diff --git a/docs/training_modes.md b/docs/training_modes.md new file mode 100644 index 0000000000..d9df7c6143 --- /dev/null +++ b/docs/training_modes.md @@ -0,0 +1,68 @@ +# Training Modes + +prime-rl supports three training modes, selected via `orchestrator.training_mode`: + +- **`rl`** — standard reinforcement learning from rewards +- **`opd`** — on-policy distillation: RL with an extra KL term toward a teacher's logprobs +- **`sft`** — hard distillation: supervised fine-tuning on teacher-generated rollouts + +The mode determines who generates rollouts, what role the teacher plays, and what must be configured. + +## Mode comparison + +| | **rl** | **opd** | **sft** | +|---|---|---|---| +| **Student does** | generate rollouts → get trained on them | generate rollouts → get trained on them | get trained on teacher's rollouts; optionally serve inference for evals | +| **Teacher does** | nothing (must be unset) | score student rollouts (token-level logprobs) | generate rollouts | +| **Loss** | reward-based (advantage) | reward + KL to teacher logprobs (`teacher_tau > 0`) | pure NLL on teacher tokens (hard distill) | +| **Student inference** (`[inference]`) | **required** | **required** | **optional** — only if you want evals or weight-sync the student | +| **Teacher inference** (`[teacher_inference]`) | forbidden | **required, must be vLLM** | not used (teacher is external) | +| **`[orchestrator.teacher]`** | must be `None` | auto-wired from `[teacher_inference]` | **required** — `client.base_url` + `model.name` of external endpoint | +| **`num_teacher_gpus`** | unset | **required** (`> 0`) | unset (teacher is external) | +| **Teacher endpoint type** | n/a | **local vLLM only** | **any OpenAI-compatible** (PI inference, OpenAI, Anthropic, local vLLM…) | +| **Weight sync (trainer → ?)** | → student inference | → student inference (teacher frozen) | → student inference if configured; teacher never touched | +| **Evals** | student | student | only if `[inference]` is set (then student evals) | + +## Key implications + +**OPD's teacher cannot be an external API.** `compute_teacher_logprobs` (`src/prime_rl/orchestrator/utils.py`) calls vLLM's `/inference/v1/generate` with `prompt_logprobs=1`. That endpoint is vLLM-specific; PI inference, OpenAI, etc. return 404. For OPD, set `num_teacher_gpus` and let `[teacher_inference]` spin up a local vLLM. + +**SFT's teacher is just chat completions.** It only needs `/v1/chat/completions`. Point `[orchestrator.teacher.client]` at anything OpenAI-compatible. No local GPU needed for the teacher. + +**SFT student inference is optional but enabling it changes behavior.** If you set `[inference]`, you get (a) student-side evals during training, and (b) weight sync from trainer to student inference (so the student inference pool reflects training progress). Without `[inference]`, the run is teacher-rollout-only with no online evals. + +**RL forbids any teacher.** Even a stray `[orchestrator.teacher]` block fails validation. + +**Student model name is always the model being trained.** In SFT this is *not* the rollout-generating model — that's the teacher. The student model field still determines tokenizer, trainer init weights, and what gets saved as checkpoints. + +## Minimal config per mode + +```toml +# rl +training_mode = "rl" +[inference] +``` + +```toml +# opd +training_mode = "opd" +[deployment] +num_teacher_gpus = 1 +[orchestrator.teacher] # empty block; auto-wired from teacher_inference +[trainer.loss] +teacher_tau = 0.5 +[inference] +# Override [teacher_inference.model] to use a different teacher model. +``` + +```toml +# sft +training_mode = "sft" +[orchestrator.teacher.client] +base_url = ["https://api.pinference.ai/api/v1"] +[orchestrator.teacher.model] +name = "qwen/qwen3-30b-a3b-instruct-2507" +[inference] # optional — drop if you don't want student-side evals +``` + +See [On-Policy Distillation](on_policy_distillation) for OPD/SFT-specific details (pure distillation, VLM support, parameters). 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 d8a57bd396..57f1d3ae7f 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/rl.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/rl.py @@ -433,6 +433,29 @@ def validate_teacher_model(self): ) return self + @model_validator(mode="after") + def validate_training_mode_loss_consistency(self): + """Cross-config invariants between orchestrator.training_mode and trainer.loss.""" + mode = self.orchestrator.training_mode + loss_type = self.trainer.loss.type + + if mode == "sft" and loss_type != "sft": + raise ValueError( + f"training_mode = 'sft' requires trainer.loss.type = 'sft' (got '{loss_type}'). " + "Either set trainer.loss.type = 'sft' or change training_mode." + ) + if mode in ("rl", "opd") and loss_type == "sft": + raise ValueError( + f"trainer.loss.type = 'sft' requires training_mode = 'sft' (got '{mode}'). " + "The sft loss path expects teacher-generated rollouts." + ) + if mode == "opd" and loss_type == "default" and self.trainer.loss.teacher_tau <= 0: + raise ValueError( + "training_mode = 'opd' requires trainer.loss.teacher_tau > 0. " + "Either set teacher_tau > 0 or change training_mode to 'rl'." + ) + return self + ### Auto-setup and validate shared configs @model_validator(mode="after") From 20027686ec7459d4e95f7cc1814b28cd71d31195 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Mon, 18 May 2026 22:47:23 +0000 Subject: [PATCH 18/47] feat(configs): lift training_mode to RLConfig + point debug teacher at local vLLM - Add shared training_mode field to RLConfig (rl | opd | sft). A mode="before" validator propagates it into orchestrator.training_mode and (for sft) trainer.loss.type before nested validation runs. Must be a before-validator because OrchestratorConfig.validate_training_mode would otherwise reject configs like training_mode = "opd" + [orchestrator.teacher] - the orchestrator-level default of "rl" would conflict with the teacher block before the after-validator could propagate the shared value. - Point debug_opd and debug_sft at a manually-deployed teacher vLLM (Qwen3-0.6B-Reverse-Text-RL on localhost:8001). Add a comment with the exact teacher start command at the top of each config. Drop the num_teacher_gpus auto-launch from debug_opd since we now use the external teacher. - Set student inference gpu_memory_utilization = 0.5 in all three debug configs (matches the teacher's mem budget, leaves room when both are co-located). - Add configs/reverse_text/README.md with the start commands. Co-Authored-By: Claude Sonnet 4.6 --- configs/reverse_text/README.md | 35 ++++++++++++++++++ configs/reverse_text/debug_opd.toml | 18 ++++++--- configs/reverse_text/debug_rl.toml | 1 + configs/reverse_text/debug_sft.toml | 17 +++++---- .../src/prime_rl/configs/rl.py | 37 +++++++++++++++++++ 5 files changed, 96 insertions(+), 12 deletions(-) create mode 100644 configs/reverse_text/README.md diff --git a/configs/reverse_text/README.md b/configs/reverse_text/README.md new file mode 100644 index 0000000000..19aef9f59f --- /dev/null +++ b/configs/reverse_text/README.md @@ -0,0 +1,35 @@ +# Reverse Text — Debug Configs + +Minimal end-to-end configs for the three training modes against the `reverse-text` env using `PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT` as the student. + +| Config | Mode | Teacher | +|---|---|---| +| `debug_rl.toml` | `rl` | none | +| `debug_opd.toml` | `opd` | local vLLM (`Qwen3-0.6B-Reverse-Text-RL`) | +| `debug_sft.toml` | `sft` | local vLLM (`Qwen3-0.6B-Reverse-Text-RL`) | + +The student inference server is auto-launched on GPU 0 at `http://localhost:8000/v1` with `gpu_memory_utilization=0.5`. The teacher (used by `debug_opd.toml` and `debug_sft.toml`) is **not** auto-launched — start it manually on GPU 1. + +## Start the teacher (only needed for opd/sft) + +```bash +CUDA_VISIBLE_DEVICES=1 uv run vllm serve PrimeIntellect/Qwen3-0.6B-Reverse-Text-RL \ + --port 8001 \ + --gpu-memory-utilization 0.5 \ + --enforce-eager +``` + +## Run the debug configs + +```bash +# RL (no teacher) +uv run rl @ configs/reverse_text/debug_rl.toml + +# OPD (needs teacher on port 8001) +uv run rl @ configs/reverse_text/debug_opd.toml + +# SFT hard distill (needs teacher on port 8001) +uv run rl @ configs/reverse_text/debug_sft.toml +``` + +See [docs/training_modes.md](../../docs/training_modes.md) for what each mode does. diff --git a/configs/reverse_text/debug_opd.toml b/configs/reverse_text/debug_opd.toml index 85b09cbad9..393ebbac7b 100644 --- a/configs/reverse_text/debug_opd.toml +++ b/configs/reverse_text/debug_opd.toml @@ -1,5 +1,12 @@ +# Start the teacher inference server first (on a separate GPU): +# CUDA_VISIBLE_DEVICES=1 uv run vllm serve PrimeIntellect/Qwen3-0.6B-Reverse-Text-RL \ +# --port 8001 --gpu-memory-utilization 0.5 --enforce-eager +# Then: +# uv run rl @ configs/reverse_text/debug_opd.toml + max_steps = 20 seq_len = 2048 +training_mode = "opd" [model] name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT" @@ -8,11 +15,7 @@ name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT" project = "reverse-text-debug" name = "debug-opd" -[deployment] -num_teacher_gpus = 1 - [orchestrator] -training_mode = "opd" batch_size = 128 rollouts_per_example = 16 @@ -32,7 +35,11 @@ max_completion_tokens = 128 [[orchestrator.eval.env]] id = "reverse-text" -[orchestrator.teacher] +[orchestrator.teacher.model] +name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-RL" + +[orchestrator.teacher.client] +base_url = ["http://localhost:8001/v1"] [trainer.optim] lr = 3e-6 @@ -43,3 +50,4 @@ teacher_tau = 0.5 [ckpt] [inference] +gpu_memory_utilization = 0.5 diff --git a/configs/reverse_text/debug_rl.toml b/configs/reverse_text/debug_rl.toml index 10f6a9e906..beb6c3ad61 100644 --- a/configs/reverse_text/debug_rl.toml +++ b/configs/reverse_text/debug_rl.toml @@ -35,3 +35,4 @@ lr = 3e-6 [ckpt] [inference] +gpu_memory_utilization = 0.5 diff --git a/configs/reverse_text/debug_sft.toml b/configs/reverse_text/debug_sft.toml index cdda571428..c998e33575 100644 --- a/configs/reverse_text/debug_sft.toml +++ b/configs/reverse_text/debug_sft.toml @@ -1,5 +1,12 @@ +# Start the teacher inference server first (on a separate GPU): +# CUDA_VISIBLE_DEVICES=1 uv run vllm serve PrimeIntellect/Qwen3-0.6B-Reverse-Text-RL \ +# --port 8001 --gpu-memory-utilization 0.5 --enforce-eager +# Then: +# uv run rl @ configs/reverse_text/debug_sft.toml + max_steps = 20 seq_len = 2048 +training_mode = "sft" [model] name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT" @@ -9,7 +16,6 @@ project = "reverse-text-debug" name = "debug-sft" [orchestrator] -training_mode = "sft" batch_size = 128 rollouts_per_example = 4 use_renderer = false @@ -31,14 +37,10 @@ max_completion_tokens = 128 id = "reverse-text" [orchestrator.teacher.model] -name = "qwen/qwen3-30b-a3b-instruct-2507" +name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-RL" [orchestrator.teacher.client] -base_url = ["https://api.pinference.ai/api/v1"] -api_key_var = "PRIME_API_KEY" - -[trainer.loss] -type = "sft" +base_url = ["http://localhost:8001/v1"] [trainer.optim] lr = 3e-6 @@ -46,3 +48,4 @@ lr = 3e-6 [ckpt] [inference] +gpu_memory_utilization = 0.5 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 57f1d3ae7f..f6716dcc08 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/rl.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/rl.py @@ -313,6 +313,15 @@ class RLConfig(BaseConfig): ), ] = None + training_mode: Annotated[ + Literal["rl", "opd", "sft"] | None, + Field( + description="Shared training mode. Propagates to orchestrator.training_mode and, " + "for 'sft', switches trainer.loss to SFTLossConfig. " + "Explicitly set per-component values always take precedence." + ), + ] = None + max_steps: Annotated[ int | None, Field( @@ -422,6 +431,34 @@ def validate_quantize_in_weight_transfer(self): return self + @model_validator(mode="before") + @classmethod + def auto_setup_training_mode(cls, data): + """Propagate shared training_mode into orchestrator.training_mode and trainer.loss.type. + + Runs before nested validation so that OrchestratorConfig.validate_training_mode + sees the propagated value. Only propagates to components that don't already set + the field explicitly. For 'sft' mode, defaults trainer.loss.type to 'sft'. + """ + if not isinstance(data, dict): + return data + mode = data.get("training_mode") + if mode is None: + return data + + orch = data.setdefault("orchestrator", {}) + if isinstance(orch, dict) and "training_mode" not in orch: + orch["training_mode"] = mode + + if mode == "sft": + trainer = data.setdefault("trainer", {}) + if isinstance(trainer, dict): + loss = trainer.setdefault("loss", {}) + if isinstance(loss, dict) and "type" not in loss: + loss["type"] = "sft" + + return data + @model_validator(mode="after") def validate_teacher_model(self): if ( From daebcc9d3edce8325942199b253543496fc0c014 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Mon, 18 May 2026 22:49:30 +0000 Subject: [PATCH 19/47] chore: drop GPU-occupancy runtime check; teacher uses uv run inference - Remove check_gpus_available in rl entrypoint - it raised RuntimeError when other processes held GPUs, which blocked iterating with a long-lived teacher inference server on a sibling GPU. - Update debug_opd / debug_sft / README to start the teacher via uv run inference (the prime-rl variant) instead of raw vllm serve. Same port (8001), same model, same flags. Co-Authored-By: Claude Sonnet 4.6 --- configs/reverse_text/README.md | 7 ++++--- configs/reverse_text/debug_opd.toml | 5 +++-- configs/reverse_text/debug_sft.toml | 5 +++-- src/prime_rl/entrypoints/rl.py | 24 ------------------------ 4 files changed, 10 insertions(+), 31 deletions(-) diff --git a/configs/reverse_text/README.md b/configs/reverse_text/README.md index 19aef9f59f..74a6d27981 100644 --- a/configs/reverse_text/README.md +++ b/configs/reverse_text/README.md @@ -13,10 +13,11 @@ The student inference server is auto-launched on GPU 0 at `http://localhost:8000 ## Start the teacher (only needed for opd/sft) ```bash -CUDA_VISIBLE_DEVICES=1 uv run vllm serve PrimeIntellect/Qwen3-0.6B-Reverse-Text-RL \ - --port 8001 \ +CUDA_VISIBLE_DEVICES=1 uv run inference \ + --model.name PrimeIntellect/Qwen3-0.6B-Reverse-Text-RL \ + --server.port 8001 \ --gpu-memory-utilization 0.5 \ - --enforce-eager + --model.enforce-eager ``` ## Run the debug configs diff --git a/configs/reverse_text/debug_opd.toml b/configs/reverse_text/debug_opd.toml index 393ebbac7b..3835f0bc4a 100644 --- a/configs/reverse_text/debug_opd.toml +++ b/configs/reverse_text/debug_opd.toml @@ -1,6 +1,7 @@ # Start the teacher inference server first (on a separate GPU): -# CUDA_VISIBLE_DEVICES=1 uv run vllm serve PrimeIntellect/Qwen3-0.6B-Reverse-Text-RL \ -# --port 8001 --gpu-memory-utilization 0.5 --enforce-eager +# CUDA_VISIBLE_DEVICES=1 uv run inference \ +# --model.name PrimeIntellect/Qwen3-0.6B-Reverse-Text-RL \ +# --server.port 8001 --gpu-memory-utilization 0.5 --model.enforce-eager # Then: # uv run rl @ configs/reverse_text/debug_opd.toml diff --git a/configs/reverse_text/debug_sft.toml b/configs/reverse_text/debug_sft.toml index c998e33575..a9c6c96214 100644 --- a/configs/reverse_text/debug_sft.toml +++ b/configs/reverse_text/debug_sft.toml @@ -1,6 +1,7 @@ # Start the teacher inference server first (on a separate GPU): -# CUDA_VISIBLE_DEVICES=1 uv run vllm serve PrimeIntellect/Qwen3-0.6B-Reverse-Text-RL \ -# --port 8001 --gpu-memory-utilization 0.5 --enforce-eager +# CUDA_VISIBLE_DEVICES=1 uv run inference \ +# --model.name PrimeIntellect/Qwen3-0.6B-Reverse-Text-RL \ +# --server.port 8001 --gpu-memory-utilization 0.5 --model.enforce-eager # Then: # uv run rl @ configs/reverse_text/debug_sft.toml diff --git a/src/prime_rl/entrypoints/rl.py b/src/prime_rl/entrypoints/rl.py index 4dec4ababe..09eec505d1 100644 --- a/src/prime_rl/entrypoints/rl.py +++ b/src/prime_rl/entrypoints/rl.py @@ -78,26 +78,6 @@ def write_subconfigs(config: RLConfig, output_dir: Path) -> None: tomli_w.dump(teacher_inference.model_dump(exclude_none=True, mode="json"), f) -def check_gpus_available(gpu_ids: list[int]) -> None: - """Raise error if there are existing processes on the specified GPUs.""" - pynvml.nvmlInit() - - occupied = [] - for gpu_id in gpu_ids: - handle = pynvml.nvmlDeviceGetHandleByIndex(gpu_id) - processes = pynvml.nvmlDeviceGetComputeRunningProcesses(handle) - if processes: - pids = [p.pid for p in processes] - occupied.append((gpu_id, pids)) - - if occupied: - msg = "Existing processes found on GPUs:\n" - for gpu_id, pids in occupied: - msg += f" GPU {gpu_id}: PIDs {pids}\n" - msg += "Kill these processes or use different GPUs." - raise RuntimeError(msg) - - def rl_local(config: RLConfig): assert config.deployment.type == "single_node" @@ -148,10 +128,6 @@ def rl_local(config: RLConfig): wandb_shared_env["WANDB_SHARED_MODE"] = "1" wandb_shared_env["WANDB_SHARED_RUN_ID"] = os.environ.get("WANDB_SHARED_RUN_ID", uuid.uuid4().hex) - # Check for existing processes on GPUs - all_gpu_ids = list(set(infer_gpu_ids + trainer_gpu_ids + teacher_gpu_ids)) - check_gpus_available(all_gpu_ids) - # Validate client port matches inference server port if config.inference is not None and not config.orchestrator.student.client.is_elastic: from urllib.parse import urlparse From 97872d3e05993028e7507423f55938fb2013f598 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Mon, 18 May 2026 23:06:36 +0000 Subject: [PATCH 20/47] fix(configs): accept legacy flat [orchestrator.model.lora] layout Pre-refactor, orchestrator.model was a ModelConfig directly (name, lora, trust_remote_code, vlm). After the student/teacher rename it's a RolloutModelConfig wrapping ModelConfig + ClientConfig, so existing [orchestrator.model.lora] TOML sections stopped parsing because they landed on the wrong nesting level. Add a mode="before" validator on RolloutModelConfig that detects a flat ModelConfig dict (any of {name, trust_remote_code, vlm, lora} at the top level, no nested model/client keys) and re-nests it under "model". This lets the three existing LoRA configs revert to their original syntax. Also fix examples/alphabet_sort/sft_distill_hard.toml: hoist training_mode to the top-level shared field so trainer.loss.type = "sft" gets auto-set (previously the orchestrator-only flag triggered the new sft/loss consistency validator). Co-Authored-By: Claude Sonnet 4.6 --- .../integration/reverse_text_lora/resume.toml | 2 +- .../integration/reverse_text_lora/start.toml | 2 +- configs/reverse_text/debug_opd.toml | 2 +- examples/alphabet_sort/sft_distill_hard.toml | 2 +- examples/wiki_search/rl.toml | 2 +- .../src/prime_rl/configs/orchestrator.py | 19 +++++++++++++++++++ tests/unit/test_configs.py | 2 +- 7 files changed, 25 insertions(+), 6 deletions(-) diff --git a/configs/ci/integration/reverse_text_lora/resume.toml b/configs/ci/integration/reverse_text_lora/resume.toml index cef3d65e17..e2b7e66ca2 100644 --- a/configs/ci/integration/reverse_text_lora/resume.toml +++ b/configs/ci/integration/reverse_text_lora/resume.toml @@ -20,7 +20,7 @@ save_adapter_separately = true batch_size = 128 rollouts_per_example = 16 -[orchestrator.student.model.lora] +[orchestrator.model.lora] name = "r8-1e-4" [orchestrator.train.sampling] diff --git a/configs/ci/integration/reverse_text_lora/start.toml b/configs/ci/integration/reverse_text_lora/start.toml index 2460203c8a..28e76d60f8 100644 --- a/configs/ci/integration/reverse_text_lora/start.toml +++ b/configs/ci/integration/reverse_text_lora/start.toml @@ -19,7 +19,7 @@ save_adapter_separately = true batch_size = 128 rollouts_per_example = 16 -[orchestrator.student.model.lora] +[orchestrator.model.lora] name = "r8-1e-4" [orchestrator.train.sampling] diff --git a/configs/reverse_text/debug_opd.toml b/configs/reverse_text/debug_opd.toml index 3835f0bc4a..b3d3df41c3 100644 --- a/configs/reverse_text/debug_opd.toml +++ b/configs/reverse_text/debug_opd.toml @@ -7,7 +7,6 @@ max_steps = 20 seq_len = 2048 -training_mode = "opd" [model] name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT" @@ -17,6 +16,7 @@ project = "reverse-text-debug" name = "debug-opd" [orchestrator] +training_mode = "opd" batch_size = 128 rollouts_per_example = 16 diff --git a/examples/alphabet_sort/sft_distill_hard.toml b/examples/alphabet_sort/sft_distill_hard.toml index da3c77199c..13e9a06fb0 100644 --- a/examples/alphabet_sort/sft_distill_hard.toml +++ b/examples/alphabet_sort/sft_distill_hard.toml @@ -1,5 +1,6 @@ max_steps = 24 seq_len = 2048 +training_mode = "sft" [deployment] type = "single_node" @@ -31,7 +32,6 @@ batch_size = 256 rollouts_per_example = 4 use_token_client = false use_renderer = false -training_mode = "sft" [orchestrator.train.sampling] max_completion_tokens = 512 diff --git a/examples/wiki_search/rl.toml b/examples/wiki_search/rl.toml index 81d78c8726..6abbb3d815 100644 --- a/examples/wiki_search/rl.toml +++ b/examples/wiki_search/rl.toml @@ -34,7 +34,7 @@ batch_size = 512 rollouts_per_example = 16 oversampling_factor = 2.0 -[orchestrator.student.model.lora] +[orchestrator.model.lora] name = "qwen3-4b-wiki-search" [orchestrator.train.sampling] 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 1364dfb067..1891415ef2 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py @@ -893,6 +893,25 @@ class RolloutModelConfig(BaseConfig): Field(description="The OAI client configuration."), ] = ClientConfig() + @model_validator(mode="before") + @classmethod + def _accept_flat_model_layout(cls, data): + """Accept legacy flat ModelConfig layout (e.g. [orchestrator.model.lora]). + + Pre-refactor, orchestrator.model was a ModelConfig directly (name, lora, + trust_remote_code, vlm). Now it's a RolloutModelConfig wrapping ModelConfig + + ClientConfig. Detect dicts whose only keys are ModelConfig fields and + re-nest them under "model" so existing configs keep working. + """ + if not isinstance(data, dict): + return data + model_only_keys = {"name", "trust_remote_code", "vlm", "lora"} + if "model" in data or "client" in data: + return data + if any(k in model_only_keys for k in data.keys()): + return {"model": data} + return data + class OrchestratorConfig(BaseConfig): """Configures the orchestrator for RL training.""" diff --git a/tests/unit/test_configs.py b/tests/unit/test_configs.py index ae3d69ea1c..731d00c5d1 100644 --- a/tests/unit/test_configs.py +++ b/tests/unit/test_configs.py @@ -198,9 +198,9 @@ def test_selective_activation_checkpointing_requires_custom_impl(): def test_sft_training_mode_enables_student_pool_when_inference_configured(): base_config = { + "training_mode": "sft", "trainer": {}, "orchestrator": { - "training_mode": "sft", "use_token_client": False, "use_renderer": False, "teacher": { From d25184e065c89f6c2be6ae4b160eaedcf1315bdb Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Mon, 18 May 2026 23:27:30 +0000 Subject: [PATCH 21/47] refactor(orchestrator): split into student_inference + teacher_inference pools Replace the "one pool plus optional teacher-client overlay" model with two explicit pools: - student_inference: InferencePool | None - Required for rl/opd, optional for sft (set iff [inference] is configured) - Target for evals, weight broadcast, and inference metrics - When None, online evals and policy updates are skipped automatically - teacher_inference: InferencePool | None - Set whenever orchestrator.teacher is configured (opd or sft) - Source of teacher logprobs in opd - Source of train rollouts in sft - Always MITO (chat completions) for simplicity; external OAI-compatible teachers (PI inference, OpenAI) work as drop-in endpoints Scheduler now takes both pools and resolves the rollout target internally (rollout_inference = teacher_inference if sft else student_inference). The old _resolve_rollout_request_target overlay is gone - the rollout pool serves rollouts directly. Fixes a latent bug where SFT-mode LoRA updates would overwrite scheduler.model_name and route subsequent teacher requests with the student's LoRA name. Also: - Drop setup_external_rollout_model (logic is now inline + clearer) - Rename setup_rollout_inference_pool -> setup_student_inference_pool; always handles only the student pool, teacher is plain setup_inference_pool - Log the resolved training_mode + one-line description at orchestrator start - Update tests: cover both sft (rollout != student) and rl (rollout = student) LoRA paths; drop the now-obsolete teacher-overlay scheduler test Co-Authored-By: Claude Sonnet 4.6 --- src/prime_rl/orchestrator/orchestrator.py | 194 ++++++++---------- src/prime_rl/orchestrator/scheduler.py | 60 +++--- src/prime_rl/orchestrator/utils.py | 21 -- .../orchestrator/test_orchestrator_setup.py | 123 +++++------ tests/unit/orchestrator/test_scheduler.py | 69 +++++-- 5 files changed, 221 insertions(+), 246 deletions(-) diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index 826a814ffa..190911b224 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -47,7 +47,6 @@ get_weight_dir, print_benchmark, set_default_executor, - setup_external_rollout_model, ) from prime_rl.orchestrator.vf_utils import ( get_seq_len, @@ -57,7 +56,6 @@ from prime_rl.trainer.model import setup_tokenizer from prime_rl.utils.client import ( init_nccl_broadcast, - setup_clients, setup_inference_pool, ) from prime_rl.utils.config import cli @@ -95,6 +93,12 @@ async def orchestrate(config: OrchestratorConfig): ) intercept_vf_logging(logger="verifiers.serve", level="WARN") # show logs from env clients logger.info("Starting orchestrator") + _MODE_DESCRIPTIONS = { + "rl": "student generates rollouts, trained with reward-based advantage", + "opd": "student generates rollouts, trained on reward + KL to teacher logprobs (on-policy distillation)", + "sft": "teacher generates rollouts, student trained on teacher tokens (hard distillation)", + } + logger.info(f"Training mode: {config.training_mode} - {_MODE_DESCRIPTIONS[config.training_mode]}") set_default_executor() event_loop_lag_monitor = EventLoopLagMonitor() @@ -119,24 +123,6 @@ async def orchestrate(config: OrchestratorConfig): for env_id in env_ids_to_install: install_env(env_id, prerelease=config.env_install_prerelease) - # Setup rollout inference pool (handles both static and elastic modes) - rollout_client_config, rollout_model_name, enable_policy_updates = setup_external_rollout_model(config, logger) - - # Setup teacher inference pool (opd: logprob distillation) - if config.training_mode == "opd": - assert config.teacher is not None - logger.info( - f"Initializing teacher inference pool (base_url={', '.join(config.teacher.client.base_url)}, " - f"model={config.teacher.model.name})" - ) - teacher_inference_pool = await setup_inference_pool( - config.teacher.client, - model_name=config.teacher.model.name, - train_client_type="openai_chat_completions", - ) - else: - teacher_inference_pool = None - # Check if this is a vision-language model (used throughout for VLM-specific paths) is_vlm = config.student.model.vlm is not None @@ -151,32 +137,42 @@ async def orchestrate(config: OrchestratorConfig): config.student.model.name, trust_remote_code=config.student.model.trust_remote_code, use_fast=True ) - teacher_clients = None - teacher_model_name = None - use_sft_override = config.training_mode == "sft" and enable_policy_updates - if use_sft_override: - logger.info(f"Using teacher rollout override (MITO, model={rollout_model_name})") - teacher_clients = setup_clients( - rollout_client_config, - client_type="openai_chat_completions", - ) - teacher_model_name = rollout_model_name - renderer = None - inference_pool = await setup_inference_pool( - config.student.client, - model_name=config.student.model.name, - train_client_type="openai_chat_completions", - eval_client_type="openai_chat_completions", + # Set up student inference pool. Required for rl/opd; optional for sft (only + # configured when the user wrote [inference] - signal: student.client.base_url + # is in model_fields_set, set by auto_setup_inference_client). When absent, + # SFT runs in teacher-only mode: no online evals, no weight sync. + has_student_inference = config.training_mode != "sft" or "base_url" in config.student.client.model_fields_set + student_inference = None + renderer = None + if has_student_inference: + logger.info( + f"Initializing student inference pool (base_url={', '.join(config.student.client.base_url)}, " + f"model={config.student.model.name})" ) - else: - renderer, inference_pool = await setup_rollout_inference_pool( + renderer, student_inference = await setup_student_inference_pool( config=config, - rollout_client_config=rollout_client_config, - rollout_model_name=rollout_model_name, tokenizer=tokenizer, logger=logger, ) + # Set up teacher inference pool (configured for opd or sft). Always MITO for + # simplicity - this also keeps external OAI-compatible teachers (PI inference, + # OpenAI) working as drop-in endpoints. + teacher_inference = None + if config.teacher is not None: + logger.info( + f"Initializing teacher inference pool (base_url={', '.join(config.teacher.client.base_url)}, " + f"model={config.teacher.model.name})" + ) + teacher_inference = await setup_inference_pool( + config.teacher.client, + model_name=config.teacher.model.name, + train_client_type="openai_chat_completions", + ) + + # Weight sync is only possible when a student inference pool exists. + enable_policy_updates = student_inference is not None + # Setup monitor (may register the run and set RUN_ID in the environment) logger.info(f"Initializing monitor (wandb={config.wandb}, prime_monitor={config.prime_monitor})") monitor = setup_monitor( @@ -256,9 +252,8 @@ async def orchestrate(config: OrchestratorConfig): scheduler = Scheduler( train_envs=train_envs, buffer=buffer, - inference_pool=inference_pool, - teacher_clients=teacher_clients, - teacher_model_name=teacher_model_name, + student_inference=student_inference, + teacher_inference=teacher_inference, max_inflight_rollouts=config.max_inflight_rollouts, max_async_level=config.max_async_level, max_off_policy_steps=config.max_off_policy_steps, @@ -268,33 +263,31 @@ async def orchestrate(config: OrchestratorConfig): lora_name=config.student.model.lora.name if config.student.model.lora else None, config=config, ) - scheduler.model_name = config.student.model.name if use_sft_override else rollout_model_name - - # Check health of the inference pool - logger.info("Waiting for inference pool to be ready") - inference_model_name = config.student.model.name if use_sft_override else rollout_model_name - await inference_pool.wait_for_ready(inference_model_name) - logger.success("Inference pool ready") - - # Start inference metrics collector (requires W&B) - inference_metrics_collector = None - if config.wandb is not None and config.collect_inference_metrics: - inference_metrics_collector = InferenceMetricsCollector(inference_pool.admin_clients) - await inference_metrics_collector.start() - # Check health of teacher inference server if configured (opd mode) - if config.training_mode == "opd" and teacher_inference_pool: + # Wait for pools to be ready + if student_inference is not None: + logger.info("Waiting for student inference pool to be ready") + await student_inference.wait_for_ready(config.student.model.name) + logger.success("Student inference pool ready") + if teacher_inference is not None: assert config.teacher is not None logger.info("Waiting for teacher inference pool to be ready") - await teacher_inference_pool.wait_for_ready(config.teacher.model.name) + await teacher_inference.wait_for_ready(config.teacher.model.name) logger.success("Teacher inference pool ready") - # Set up weight broadcast backend + # Start inference metrics collector (requires W&B + student inference pool) + inference_metrics_collector = None + if config.wandb is not None and config.collect_inference_metrics and student_inference is not None: + inference_metrics_collector = InferenceMetricsCollector(student_inference.admin_clients) + await inference_metrics_collector.start() + + # Set up weight broadcast backend (targets student inference) if enable_policy_updates: + assert student_inference is not None logger.info(f"Initializing weight broadcast ({config.weight_broadcast})") if config.weight_broadcast.type == "nccl": await init_nccl_broadcast( - inference_pool.admin_clients, + student_inference.admin_clients, config.weight_broadcast.host, config.weight_broadcast.port, config.weight_broadcast.timeout, @@ -302,7 +295,7 @@ async def orchestrate(config: OrchestratorConfig): quantize_in_weight_transfer=config.weight_broadcast.quantize_in_weight_transfer, ) else: - logger.info("Skipping weight broadcast initialization (SFT distillation mode)") + logger.info("Skipping weight broadcast initialization (no student inference pool)") # Setup training batch sender for sending training examples to trainer logger.info(f"Initializing training batch sender ({config.rollout_transport})") @@ -329,6 +322,7 @@ async def orchestrate(config: OrchestratorConfig): prev_ckpt_step = scheduler.ckpt_step - 1 if enable_policy_updates: + assert student_inference is not None # In NCCL mode, skip existence check - weights are broadcasted, not stored on disk check_exists = config.weight_broadcast.type != "nccl" wait_timeout = config.ckpt.wait_for_weights_timeout if config.ckpt else None @@ -336,10 +330,11 @@ async def orchestrate(config: OrchestratorConfig): config.output_dir, scheduler.ckpt_step, check_exists=check_exists, wait_timeout=wait_timeout ) lora_name = config.student.model.lora.name if config.student.model.lora else None - await inference_pool.update_weights(weights_path, lora_name=lora_name, step=scheduler.ckpt_step) + await student_inference.update_weights(weights_path, lora_name=lora_name, step=scheduler.ckpt_step) if lora_name is not None: - inference_pool.update_model_name(lora_name) - scheduler.model_name = lora_name + student_inference.update_model_name(lora_name) + if scheduler.rollout_inference is student_inference: + scheduler.model_name = lora_name else: logger.info("Training from scratch") @@ -383,7 +378,7 @@ async def orchestrate(config: OrchestratorConfig): # scheduler.checkpoint_ready during eval to ensure consistent weights. # Each eval env has its own interval, so we check each independently. envs_to_eval: list[EvalEnv] = [] - if config.eval: + if config.eval and student_inference is not None: assert eval_envs is not None for eval_env in eval_envs: eval_ckpt_step = compute_eval_ckpt_step( @@ -398,6 +393,7 @@ async def orchestrate(config: OrchestratorConfig): envs_to_eval.append(eval_env) if envs_to_eval: + assert student_inference is not None env_names = ", ".join(e.name for e in envs_to_eval) logger.info(f"Running evals at {ckpt_step=} for {env_names}") @@ -413,8 +409,8 @@ async def orchestrate(config: OrchestratorConfig): eval_results = await asyncio.gather( *[ eval_env.evaluate( - model_name=inference_pool.model_name, - get_client=inference_pool.get_eval_client, + model_name=student_inference.model_name, + get_client=student_inference.get_eval_client, ckpt_step=ckpt_step, step=progress.step, cache_salt=str(ckpt_step), @@ -594,13 +590,14 @@ def process_rollout(rollout: vf.RolloutOutput, rollout_idx: int) -> list[Trainin f"to {len(train_examples)} training examples" ) - # Compute teacher logprobs if teacher model is configured + # Compute teacher logprobs (opd only - sft trains on teacher tokens directly) teacher_logprobs_time = 0 - if config.teacher and teacher_inference_pool: + if config.training_mode == "opd" and teacher_inference is not None: + assert config.teacher is not None logger.info(f"Computing teacher logprobs for {len(train_examples)} training examples") teacher_logprobs_start_time = time.perf_counter() teacher_logprobs_list = await compute_teacher_logprobs( - clients=teacher_inference_pool.train_clients, + clients=teacher_inference.train_clients, model_name=config.teacher.model.name, samples=train_examples, ) @@ -834,13 +831,13 @@ def compute_solve_rates(df): if heart is not None: heart.beat() - if config.eval and eval_envs is not None: + if config.eval and eval_envs is not None and student_inference is not None: logger.info("Running final evals") eval_results = await asyncio.gather( *[ eval_env.evaluate( - model_name=inference_pool.model_name, - get_client=inference_pool.get_eval_client, + model_name=student_inference.model_name, + get_client=student_inference.get_eval_client, ckpt_step=ckpt_step, step=progress.step, cache_salt=str(ckpt_step), @@ -876,9 +873,10 @@ async def _graceful_shutdown() -> None: await scheduler.stop() if inference_metrics_collector is not None: await inference_metrics_collector.stop() - await inference_pool.stop() - if teacher_inference_pool is not None: - await teacher_inference_pool.stop() + if student_inference is not None: + await student_inference.stop() + if teacher_inference is not None: + await teacher_inference.stop() event_loop_lag_monitor_task.cancel() # Shutdown env processes (also registered as atexit handler for crash safety) train_envs.shutdown() @@ -916,40 +914,28 @@ def main(): asyncio.run(orchestrate(cli(OrchestratorConfig))) -async def setup_rollout_inference_pool( +async def setup_student_inference_pool( *, config: OrchestratorConfig, - rollout_client_config, - rollout_model_name: str, tokenizer, logger, ): - """Set up rollout inference. + """Set up the student inference pool (rollouts when rl/opd, evals + weight sync always). Routing policy is driven by ``config.use_token_client`` and ``config.use_renderer`` (mutually exclusive — config-level validators block both being True): - - external teacher rollout → MITO (``openai_chat_completions``), - selected independently of the toggles (config-level validator rejects - ``use_token_client`` / ``use_renderer`` in that case) - ``use_renderer=True`` → renderer-backed TITO client (``/v1/generate``). - Default for text-only rollouts. - Not allowed for VLMs (validated at config time). - - ``use_token_client=True`` → server-tokenized TITO - (``openai_chat_completions_token``, ``/v1/chat/completions/tokens``). - - both False → MITO (``openai_chat_completions``). - VLMs land here too. + Default for text-only rollouts. Not allowed for VLMs (validated at config time). + - ``use_token_client=True`` → server-tokenized TITO (``/v1/chat/completions/tokens``). + - both False → MITO (``openai_chat_completions``). VLMs land here too. + + Eval clients always use MITO. In sft mode the renderer/tito knobs are forced + off by config validators, so the student pool is plain MITO end-to-end. """ - if config.training_mode == "sft": - logger.info("Using external rollout model (MITO) without renderer client") - inference_pool = await setup_inference_pool( - rollout_client_config, - model_name=rollout_model_name, - train_client_type="openai_chat_completions", - eval_client_type="openai_chat_completions", - ) - return None, inference_pool + client_config = config.student.client + model_name = config.student.model.name if config.use_renderer: renderer = create_renderer( @@ -960,10 +946,10 @@ async def setup_rollout_inference_pool( preserve_all_thinking=config.renderer.preserve_all_thinking, preserve_thinking_between_tool_calls=config.renderer.preserve_thinking_between_tool_calls, ) - logger.info(f"Initialized {type(renderer).__name__} for {config.student.model.name}") + logger.info(f"Initialized {type(renderer).__name__} for {model_name}") inference_pool = await setup_inference_pool( - rollout_client_config, - model_name=rollout_model_name, + client_config, + model_name=model_name, train_client_type="renderer", eval_client_type="openai_chat_completions", renderer_name=config.renderer.name, @@ -982,8 +968,8 @@ async def setup_rollout_inference_pool( else: logger.info("Using MITO (openai_chat_completions) for rollouts") inference_pool = await setup_inference_pool( - rollout_client_config, - model_name=rollout_model_name, + client_config, + model_name=model_name, train_client_type=train_client_type, eval_client_type="openai_chat_completions", ) diff --git a/src/prime_rl/orchestrator/scheduler.py b/src/prime_rl/orchestrator/scheduler.py index edc4cfe9e0..fbc4b79efc 100644 --- a/src/prime_rl/orchestrator/scheduler.py +++ b/src/prime_rl/orchestrator/scheduler.py @@ -73,7 +73,8 @@ class Scheduler: def __init__( self, train_envs: TrainEnvs, - inference_pool: InferencePool, + student_inference: InferencePool | None, + teacher_inference: InferencePool | None, buffer: Buffer, config: OrchestratorConfig, max_inflight_rollouts: int, @@ -83,8 +84,6 @@ def __init__( tasks_per_minute: int | None, enable_policy_updates: bool = True, lora_name: str | None = None, - teacher_clients: list[vf.ClientConfig] | None = None, - teacher_model_name: str | None = None, ): self.logger = get_logger() if tasks_per_minute is not None: @@ -103,12 +102,18 @@ def __init__( self.strict_async_level = strict_async_level self.enable_policy_updates = enable_policy_updates self.lora_name = lora_name - self.model_name = self.config.student.model.name self.json_logging = config.log.json_logging - self.inference_pool = inference_pool - self.teacher_clients = teacher_clients - self.teacher_model_name = teacher_model_name + # student_inference is the weight-sync target (None = no policy updates). + # teacher_inference is set in opd (for logprobs) and sft (for rollouts). + # rollout_inference is whichever pool serves train rollouts for this mode. + self.student_inference = student_inference + self.teacher_inference = teacher_inference + rollout = teacher_inference if config.training_mode == "sft" else student_inference + assert rollout is not None, "rollout_inference resolved to None - config validation should prevent this" + self.rollout_inference: InferencePool = rollout + # model_name is the name to send on rollout requests - matches the rollout pool + self.model_name = self.rollout_inference.model_name group_scoring_envs = [env.name for env in train_envs if env.requires_group_scoring] if group_scoring_envs: @@ -178,24 +183,13 @@ async def _select_least_loaded_client(self) -> vf.ClientConfig: Uses (api_base_url, dp_rank) as identity rather than client_idx so that load tracking survives elastic pool refreshes (which reassign indices). """ - clients = self.inference_pool.train_clients + clients = self.rollout_inference.train_clients while not clients: await asyncio.sleep(1) - clients = self.inference_pool.train_clients + clients = self.rollout_inference.train_clients inflight = Counter(self._client_identity(info.client_config) for info in self.inflight_requests.values()) return min(clients, key=lambda c: inflight[self._client_identity(c)]) - def _resolve_rollout_request_target(self, client_config: vf.ClientConfig) -> tuple[vf.ClientConfig, str]: - if self.teacher_clients is None: - return client_config, self.model_name - - # The scheduler pins/load-balances against the student inference pool. - # Map that selected logical client onto the teacher client set, which may - # have a different size due to different base_url or dp_rank_count settings. - teacher_client = self.teacher_clients[client_config.client_idx % len(self.teacher_clients)] - assert self.teacher_model_name is not None - return teacher_client, self.teacher_model_name - async def drop_group(self, group_id: int) -> int: """Drop a group and cancel any remaining in-flight rollouts for it. Returns the number of cancelled rollouts.""" tasks_to_cancel = [] @@ -229,16 +223,15 @@ async def schedule_rollout(self, group_id: int): env_name = group.example["env_name"] env = self.train_envs.get(env_name) - request_client_config, request_model_name = self._resolve_rollout_request_target(client_config) cache_salt = str(self.ckpt_step) if env.requires_group_scoring: rollout_count = group.rollouts_to_schedule group.rollouts_to_schedule = 0 task = asyncio.create_task( env.run_group( - client=request_client_config, + client=client_config, example=group.example, - model_name=request_model_name, + model_name=self.model_name, rollouts_per_example=rollout_count, cache_salt=cache_salt, ) @@ -248,9 +241,9 @@ async def schedule_rollout(self, group_id: int): group.rollouts_to_schedule -= 1 task = asyncio.create_task( env.run_rollout( - client=request_client_config, + client=client_config, example=group.example, - model_name=request_model_name, + model_name=self.model_name, cache_salt=cache_salt, ) ) @@ -335,14 +328,21 @@ async def _apply_policy_update(self, next_ckpt_step: int) -> None: update_weights_start_time = time.perf_counter() weights_path = get_step_path(get_broadcast_dir(self.config.output_dir), next_ckpt_step) - await self.inference_pool.update_weights(weights_path, lora_name=self.lora_name, step=next_ckpt_step) + assert self.student_inference is not None, ( + "weight sync requires student_inference - guard with enable_policy_updates" + ) + await self.student_inference.update_weights(weights_path, lora_name=self.lora_name, step=next_ckpt_step) self.update_weights_time = time.perf_counter() - update_weights_start_time self.logger.debug(f"Updated weights to step {next_ckpt_step} in {self.update_weights_time:.2f}s") self.ckpt_step = next_ckpt_step if self.lora_name is not None: - self.inference_pool.update_model_name(self.lora_name) - self.model_name = self.lora_name + self.student_inference.update_model_name(self.lora_name) + # Only redirect rollout requests to the new LoRA when rollouts come from + # student inference (rl/opd). In sft, rollouts go to the teacher and + # the student's LoRA name is irrelevant to them. + if self.rollout_inference is self.student_inference: + self.model_name = self.lora_name self.checkpoint_ready.set() await self._update_off_policy() @@ -613,7 +613,7 @@ def get_metrics(self) -> dict[str, float]: self.total_rollouts_by_env.clear() self.dropped_groups_by_env.clear() - # Add inference pool metrics (e.g. elastic pool server counts) - metrics.update(self.inference_pool.get_metrics()) + # Add train pool metrics (e.g. elastic pool server counts) + metrics.update(self.rollout_inference.get_metrics()) return metrics diff --git a/src/prime_rl/orchestrator/utils.py b/src/prime_rl/orchestrator/utils.py index 4096637117..121cc57c54 100644 --- a/src/prime_rl/orchestrator/utils.py +++ b/src/prime_rl/orchestrator/utils.py @@ -11,7 +11,6 @@ from rich.table import Table from verifiers.utils.client_utils import setup_openai_client -from prime_rl.configs.orchestrator import OrchestratorConfig from prime_rl.transport import TrainingSample from prime_rl.utils.logger import get_logger from prime_rl.utils.utils import ( @@ -177,23 +176,3 @@ def find_stable_dir() -> Path | None: return broadcast_weight_dir raise FileNotFoundError(f"No weight directory found for checkpoint step {step}") - - -def setup_external_rollout_model(config: OrchestratorConfig, logger) -> tuple[Any, str, bool]: - """Resolve rollout client/model and whether student policy updates are enabled. - - - rl/opd: student generates rollouts, policy updates always enabled. - - sft: teacher generates rollouts; policy updates enabled iff the student - inference client is configured (non-empty base_url). - """ - if config.training_mode in ("rl", "opd"): - return config.student.client, config.student.model.name, True - - assert config.teacher is not None # validated by validate_training_mode - rollout_client_config = config.teacher.client - rollout_model_name = config.teacher.model.name - enable_policy_updates = "base_url" in config.student.client.model_fields_set - logger.info( - f"Using teacher rollout model (base_url={', '.join(rollout_client_config.base_url)}, model={rollout_model_name})" - ) - return rollout_client_config, rollout_model_name, enable_policy_updates diff --git a/tests/unit/orchestrator/test_orchestrator_setup.py b/tests/unit/orchestrator/test_orchestrator_setup.py index d6c3f59d3e..6e5021ec9f 100644 --- a/tests/unit/orchestrator/test_orchestrator_setup.py +++ b/tests/unit/orchestrator/test_orchestrator_setup.py @@ -2,133 +2,110 @@ from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch -from prime_rl.orchestrator.orchestrator import setup_rollout_inference_pool -from prime_rl.orchestrator.utils import setup_external_rollout_model +from prime_rl.orchestrator.orchestrator import setup_student_inference_pool -def test_setup_rollout_inference_pool_uses_plain_client_for_sft_mode(): +def test_setup_student_inference_pool_uses_renderer_when_enabled(): async def run() -> None: tokenizer = object() config = SimpleNamespace( - training_mode="sft", - student=SimpleNamespace(renderer="auto", model=SimpleNamespace(name="student-model")), + training_mode="rl", + use_renderer=True, + use_token_client=False, + student=SimpleNamespace( + client=SimpleNamespace(base_url=["http://localhost:8000/v1"]), + model=SimpleNamespace(name="student-model"), + ), + renderer=SimpleNamespace( + name="qwen3_vl", + tool_parser=None, + reasoning_parser=None, + pool_size=None, + preserve_all_thinking=False, + preserve_thinking_between_tool_calls=False, + ), ) - rollout_client_config = SimpleNamespace(base_url=["https://api.pinference.ai/api/v1"]) logger = MagicMock() + renderer = object() inference_pool = object() with ( + patch("prime_rl.orchestrator.orchestrator.create_renderer", return_value=renderer) as create_renderer_mock, patch( "prime_rl.orchestrator.orchestrator.setup_inference_pool", new=AsyncMock(return_value=inference_pool), ) as setup_pool_mock, - patch("prime_rl.orchestrator.orchestrator.create_renderer") as create_renderer_mock, ): - renderer, returned_pool = await setup_rollout_inference_pool( + returned_renderer, returned_pool = await setup_student_inference_pool( config=config, - rollout_client_config=rollout_client_config, - rollout_model_name="teacher-model", tokenizer=tokenizer, logger=logger, ) - assert renderer is None + assert returned_renderer is renderer assert returned_pool is inference_pool - create_renderer_mock.assert_not_called() + create_renderer_mock.assert_called_once_with( + tokenizer, + renderer="qwen3_vl", + tool_parser=None, + reasoning_parser=None, + preserve_all_thinking=False, + preserve_thinking_between_tool_calls=False, + ) setup_pool_mock.assert_awaited_once_with( - rollout_client_config, - model_name="teacher-model", - train_client_type="openai_chat_completions", + config.student.client, + model_name="student-model", + train_client_type="renderer", eval_client_type="openai_chat_completions", + renderer_name="qwen3_vl", + tool_parser=None, + reasoning_parser=None, + renderer_pool_size=None, + preserve_all_thinking=False, + preserve_thinking_between_tool_calls=False, ) asyncio.run(run()) -def test_setup_external_rollout_model_sft_uses_teacher_and_checks_student_client(): - from prime_rl.configs.orchestrator import ClientConfig, RolloutModelConfig +def test_setup_student_inference_pool_defaults_to_mito(): + """No renderer, no token client -> plain MITO chat completions.""" - teacher_client = SimpleNamespace(base_url=["https://teacher.example/v1"]) - logger = MagicMock() - - # SFT mode, student client has default base_url (not in model_fields_set) → policy updates disabled - config = SimpleNamespace( - training_mode="sft", - student=RolloutModelConfig(), # default client — base_url not explicitly set - teacher=SimpleNamespace(client=teacher_client, model=SimpleNamespace(name="teacher-model")), - ) - rollout_client, rollout_model, enable_policy_updates = setup_external_rollout_model(config, logger) - assert rollout_client is teacher_client - assert rollout_model == "teacher-model" - assert not enable_policy_updates - - # SFT mode, student client base_url explicitly set → policy updates enabled - student_model = RolloutModelConfig(client=ClientConfig(base_url=["http://localhost:8000/v1"])) - config.student = student_model - rollout_client, rollout_model, enable_policy_updates = setup_external_rollout_model(config, logger) - assert rollout_client is teacher_client - assert rollout_model == "teacher-model" - assert enable_policy_updates - - -def test_setup_rollout_inference_pool_uses_direct_renderer_client_for_local_vllm(): async def run() -> None: tokenizer = object() config = SimpleNamespace( training_mode="rl", - use_renderer=True, + use_renderer=False, use_token_client=False, - student=SimpleNamespace(model=SimpleNamespace(name="student-model")), - renderer=SimpleNamespace( - name="qwen3_vl", - tool_parser=None, - reasoning_parser=None, - pool_size=None, - preserve_all_thinking=False, - preserve_thinking_between_tool_calls=False, + student=SimpleNamespace( + client=SimpleNamespace(base_url=["http://localhost:8000/v1"]), + model=SimpleNamespace(name="student-model"), ), ) - rollout_client_config = SimpleNamespace(base_url=["http://localhost:8000/v1"]) logger = MagicMock() - renderer = object() inference_pool = object() with ( - patch("prime_rl.orchestrator.orchestrator.create_renderer", return_value=renderer) as create_renderer_mock, + patch("prime_rl.orchestrator.orchestrator.create_renderer") as create_renderer_mock, patch( "prime_rl.orchestrator.orchestrator.setup_inference_pool", new=AsyncMock(return_value=inference_pool), ) as setup_pool_mock, ): - returned_renderer, returned_pool = await setup_rollout_inference_pool( + renderer, returned_pool = await setup_student_inference_pool( config=config, - rollout_client_config=rollout_client_config, - rollout_model_name="student-model", tokenizer=tokenizer, logger=logger, ) - assert returned_renderer is renderer + assert renderer is None assert returned_pool is inference_pool - create_renderer_mock.assert_called_once_with( - tokenizer, - renderer="qwen3_vl", - tool_parser=None, - reasoning_parser=None, - preserve_all_thinking=False, - preserve_thinking_between_tool_calls=False, - ) + create_renderer_mock.assert_not_called() setup_pool_mock.assert_awaited_once_with( - rollout_client_config, + config.student.client, model_name="student-model", - train_client_type="renderer", + train_client_type="openai_chat_completions", eval_client_type="openai_chat_completions", - renderer_name="qwen3_vl", - tool_parser=None, - reasoning_parser=None, - renderer_pool_size=None, - preserve_all_thinking=False, - preserve_thinking_between_tool_calls=False, ) asyncio.run(run()) diff --git a/tests/unit/orchestrator/test_scheduler.py b/tests/unit/orchestrator/test_scheduler.py index 3832a99c4d..1d49482362 100644 --- a/tests/unit/orchestrator/test_scheduler.py +++ b/tests/unit/orchestrator/test_scheduler.py @@ -32,8 +32,6 @@ def make_scheduler() -> Scheduler: scheduler.update_policy_task = None scheduler.enable_policy_updates = True scheduler.rate_limiter = None - scheduler.teacher_clients = None - scheduler.teacher_model_name = None return scheduler @@ -105,10 +103,11 @@ async def update_weights(weight_dir, lora_name=None, step=0) -> None: started.set() await release.wait() - scheduler.inference_pool = SimpleNamespace( + scheduler.student_inference = SimpleNamespace( update_weights=update_weights, update_model_name=MagicMock(), ) + scheduler.rollout_inference = scheduler.student_inference scheduler._update_off_policy = AsyncMock() with ( @@ -145,10 +144,11 @@ async def update_weights(weight_dir, lora_name=None, step=0) -> None: finally: cancelled.set() - scheduler.inference_pool = SimpleNamespace( + scheduler.student_inference = SimpleNamespace( update_weights=update_weights, update_model_name=MagicMock(), ) + scheduler.rollout_inference = scheduler.student_inference scheduler._update_off_policy = AsyncMock() with ( @@ -179,17 +179,53 @@ def test_client_identity_distinguishes_base_url_and_dp_rank(): assert Scheduler._client_identity(client_a) != Scheduler._client_identity(client_b) -def test_lora_policy_update_keeps_student_model_name_with_teacher_rollout_override(): +def test_lora_policy_update_in_sft_keeps_teacher_model_name(): + """In sft mode, train_pool is the teacher. LoRA updates the student inference + pool but must not change scheduler.model_name (which is what gets sent to the + teacher endpoint on each rollout request).""" + + async def run() -> None: + scheduler = make_scheduler() + scheduler.model_name = "teacher-model" + scheduler.lora_name = "student-lora" + + student_inference = SimpleNamespace( + update_weights=AsyncMock(), + update_model_name=MagicMock(), + ) + teacher_inference = SimpleNamespace() + scheduler.student_inference = student_inference + scheduler.rollout_inference = teacher_inference # sft: train_pool != student_inference + scheduler._update_off_policy = AsyncMock() + + with ( + patch("prime_rl.orchestrator.scheduler.get_latest_ckpt_step", return_value=8), + patch("prime_rl.orchestrator.scheduler.wait_for_path", new=AsyncMock()), + ): + await scheduler.maybe_update_policy() + + student_inference.update_weights.assert_awaited_once() + student_inference.update_model_name.assert_called_once_with("student-lora") + assert scheduler.model_name == "teacher-model" + + asyncio.run(run()) + + +def test_lora_policy_update_in_rl_updates_model_name(): + """In rl/opd mode, train_pool is the student. LoRA updates redirect rollout + requests to the new LoRA name.""" + async def run() -> None: scheduler = make_scheduler() scheduler.model_name = "student-model" scheduler.lora_name = "student-lora" - scheduler.teacher_model_name = "teacher-model" - scheduler.inference_pool = SimpleNamespace( + student_inference = SimpleNamespace( update_weights=AsyncMock(), update_model_name=MagicMock(), ) + scheduler.student_inference = student_inference + scheduler.rollout_inference = student_inference # rl/opd: same pool scheduler._update_off_policy = AsyncMock() with ( @@ -198,26 +234,25 @@ async def run() -> None: ): await scheduler.maybe_update_policy() - scheduler.inference_pool.update_weights.assert_awaited_once() - scheduler.inference_pool.update_model_name.assert_called_once_with("student-lora") + student_inference.update_weights.assert_awaited_once() + student_inference.update_model_name.assert_called_once_with("student-lora") assert scheduler.model_name == "student-lora" - assert scheduler.teacher_model_name == "teacher-model" asyncio.run(run()) -def test_schedule_rollout_applies_teacher_override_at_request_submission(): +def test_schedule_rollout_uses_train_pool(): + """schedule_rollout dispatches to train_pool's clients with train_pool's model name.""" + async def run() -> None: scheduler = make_scheduler() - student_client = vf.ClientConfig(api_base_url="http://student.example/v1") + scheduler.model_name = "teacher-model" teacher_client = vf.ClientConfig(api_base_url="http://teacher.example/v1") env = SimpleNamespace( requires_group_scoring=False, run_rollout=AsyncMock(return_value=[]), ) - scheduler.inference_pool = SimpleNamespace(train_clients=[student_client]) - scheduler.teacher_clients = [teacher_client] - scheduler.teacher_model_name = "teacher-model" + scheduler.rollout_inference = SimpleNamespace(train_clients=[teacher_client]) scheduler.train_envs = SimpleNamespace(get=MagicMock(return_value=env)) scheduler.groups = { 0: GroupState( @@ -235,8 +270,6 @@ async def run() -> None: model_name="teacher-model", cache_salt="7", ) - assert scheduler.groups[0].pinned_client is student_client - [info] = scheduler.inflight_requests.values() - assert info.client_config is student_client + assert scheduler.groups[0].pinned_client is teacher_client asyncio.run(run()) From 6597e16fb0c9c029fd4a8ae24355759d9e8494cf Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Mon, 18 May 2026 23:52:44 +0000 Subject: [PATCH 22/47] docs+configs: legacy [orchestrator.client] shim + merge OPD doc into training_modes - Add _accept_top_level_client validator on OrchestratorConfig: legacy [orchestrator.client...] re-nests under [orchestrator.student.client...] matching the existing _accept_flat_model_layout shim pattern. Revert configs/elastic/rl.toml back to the original [orchestrator.client.elastic] syntax. - Revert configs/ci/integration/reverse_text_multi_run/orchestrator.toml [model.model] -> [model]; the flat-layout shim on RolloutModelConfig already handles it (verified). - Merge docs/on_policy_distillation.md into docs/training_modes.md. The training_modes page is now the single entry point; OPD-specific details (external teacher, pure distillation, monitoring, VLM, params) live as subsections. Delete the standalone OPD page. - Drop the redundant "Initializing static inference pool" / WandbMonitor / PrimeMonitor init logs - they duplicate the immediately-preceding per-pool init line. Co-Authored-By: Claude Sonnet 4.6 --- .../reverse_text_multi_run/orchestrator.toml | 2 +- configs/elastic/rl.toml | 2 +- docs/on_policy_distillation.md | 119 ------------------ docs/training_modes.md | 86 +++++++++++-- .../src/prime_rl/configs/orchestrator.py | 63 ++++++---- src/prime_rl/orchestrator/orchestrator.py | 15 +-- src/prime_rl/utils/client.py | 5 - src/prime_rl/utils/monitor/prime.py | 1 - src/prime_rl/utils/monitor/wandb.py | 1 - 9 files changed, 127 insertions(+), 167 deletions(-) delete mode 100644 docs/on_policy_distillation.md diff --git a/configs/ci/integration/reverse_text_multi_run/orchestrator.toml b/configs/ci/integration/reverse_text_multi_run/orchestrator.toml index 92f403ff73..be34be66fb 100644 --- a/configs/ci/integration/reverse_text_multi_run/orchestrator.toml +++ b/configs/ci/integration/reverse_text_multi_run/orchestrator.toml @@ -5,7 +5,7 @@ rollouts_per_example = 16 seq_len = 2048 max_steps = 20 -[model.model] +[model] name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT" [optim] diff --git a/configs/elastic/rl.toml b/configs/elastic/rl.toml index 4233b1e437..387bfddbf2 100644 --- a/configs/elastic/rl.toml +++ b/configs/elastic/rl.toml @@ -35,7 +35,7 @@ rollouts_per_example = 8 [orchestrator.train.sampling] max_completion_tokens = 768 -[orchestrator.model.client.elastic] +[orchestrator.client.elastic] hostname = "localhost" port = 8000 sync_interval = 5.0 diff --git a/docs/on_policy_distillation.md b/docs/on_policy_distillation.md deleted file mode 100644 index 01485bbf60..0000000000 --- a/docs/on_policy_distillation.md +++ /dev/null @@ -1,119 +0,0 @@ -# On-Policy Distillation - -On-policy distillation uses a teacher model to provide dense token-level feedback during RL training. The student generates rollouts, and the teacher's logprobs guide the student to stay close to stronger behavior while still learning from rewards. - -For more details, see [On-Policy Distillation](https://thinkingmachines.ai/blog/on-policy-distillation/) by Thinking Machines. - -## Quick Start - -Add `num_teacher_gpus` to `[deployment]` and set `teacher_tau > 0`: - -```toml -[deployment] -num_teacher_gpus = 2 - -[trainer.loss] -teacher_tau = 0.5 -``` - -This automatically starts a teacher inference server using the same model as inference. To use a different teacher model: - -```toml -[deployment] -num_teacher_gpus = 2 - -[teacher_inference.model] -name = "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B" - -[trainer.loss] -teacher_tau = 0.5 -``` - -## Using an External Teacher Server - -If the teacher is already running elsewhere: - -```toml -[trainer.loss] -teacher_tau = 0.5 - -[orchestrator.teacher_model.client] -base_url = ["http://teacher-server:8000/v1"] - -[orchestrator.teacher_model.model] -name = "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B" -``` - -## Pure Distillation (No Verification) - -For agentic environments where verification is expensive (code execution, tool use, multi-turn interactions), you can skip verification entirely and use only the teacher signal: - -```toml -[deployment] -num_teacher_gpus = 2 - -[trainer.loss] -teacher_tau = 1.0 -adv_tau = 0.0 # Disable reward-based learning - -[orchestrator.verification] -enabled = false # Skip expensive verification -``` - -This runs pure on-policy distillation: the student learns to match the teacher without needing any reward signal. - -## SFT Distillation ("Hard Distillation") From Teacher Rollouts - -Use this mode when you want to train from teacher-generated completions directly (hard distillation), without teacher token-level logprobs. - -```toml -[trainer.loss] -type = "sft" - -[orchestrator] -use_token_client = false -use_renderer = false -training_mode = "sft" - -[orchestrator.teacher_model.client] -base_url = ["https://your-openai-compatible-endpoint/v1"] -skip_model_check = true - -[orchestrator.teacher_model.model] -name = "teacher-model-name" -``` - -In this mode: -- Rollouts are generated from `orchestrator.teacher_model` -- The orchestrator uses text-level reconstruction with the student tokenizer -- The RL trainer optimizes masked NLL (`trainer.loss.type = "sft"`) -- Omit `[inference]` (no local inference server required) - -### Image Input (VLM) Support - -Yes, image input is supported in SFT/hard-distillation mode when the student model is multimodal (VLM). - -- Prompts can include OpenAI-style image items in `message.content`, e.g. `{"type": "image_url", "image_url": {"url": "data:image/..."}}` -- The orchestrator extracts and preprocesses images from trajectory prompts and attaches `pixel_values`/`image_grid_thw` to training samples -- No teacher token IDs/logprobs are required; reconstruction still happens from messages - -Notes: -- This path currently expects `data:image/...` payloads in message content -- The teacher rollout endpoint still needs to be able to handle the same multimodal prompts during generation - -Reference configs: -- `configs/alphabet_sort/sft_distill_hard_qwen4b_lora_prime_teacher.toml` -- `examples/alphabet_sort/sft_distill_hard.toml` - -## Parameters - -| Parameter | Default | Description | -|-----------|---------|-------------| -| `deployment.num_teacher_gpus` | `None` | Number of GPUs for teacher server. Auto-starts server when set. | -| `trainer.loss.teacher_tau` | `0.0` | Distillation strength. Set `> 0` to enable. | -| `trainer.loss.adv_tau` | `1.0` | Weight for RL advantage signal. Set `0` for pure distillation. | -| `orchestrator.verification.enabled` | `true` | Enable/disable verification. Set to `false` for pure distillation with `adv_tau = 0`. | - -## Monitoring - -The `teacher_kl` metric shows the KL divergence from teacher to student. Lower values mean the student is closer to the teacher. diff --git a/docs/training_modes.md b/docs/training_modes.md index d9df7c6143..d6c49285f1 100644 --- a/docs/training_modes.md +++ b/docs/training_modes.md @@ -1,9 +1,9 @@ # Training Modes -prime-rl supports three training modes, selected via `orchestrator.training_mode`: +prime-rl supports three training modes, selected via `training_mode`: - **`rl`** — standard reinforcement learning from rewards -- **`opd`** — on-policy distillation: RL with an extra KL term toward a teacher's logprobs +- **`opd`** — on-policy distillation: RL with an extra KL term toward a teacher's logprobs ([Thinking Machines blog post](https://thinkingmachines.ai/blog/on-policy-distillation/)) - **`sft`** — hard distillation: supervised fine-tuning on teacher-generated rollouts The mode determines who generates rollouts, what role the teacher plays, and what must be configured. @@ -44,25 +44,93 @@ training_mode = "rl" ``` ```toml -# opd +# opd — auto-launched local teacher training_mode = "opd" [deployment] -num_teacher_gpus = 1 -[orchestrator.teacher] # empty block; auto-wired from teacher_inference +num_teacher_gpus = 1 # spin up a teacher vLLM +[orchestrator.teacher] # empty block; client + model auto-wired [trainer.loss] teacher_tau = 0.5 [inference] -# Override [teacher_inference.model] to use a different teacher model. +# Override [teacher_inference.model] to use a different teacher model than the student. ``` ```toml -# sft +# sft — external teacher (PI inference) training_mode = "sft" [orchestrator.teacher.client] base_url = ["https://api.pinference.ai/api/v1"] [orchestrator.teacher.model] name = "qwen/qwen3-30b-a3b-instruct-2507" -[inference] # optional — drop if you don't want student-side evals +[inference] # optional — drop if you don't want student-side evals ``` -See [On-Policy Distillation](on_policy_distillation) for OPD/SFT-specific details (pure distillation, VLM support, parameters). +## OPD details + +### Using an external (already-running) teacher + +Skip `num_teacher_gpus` and point at the existing endpoint. The teacher **must** be a vLLM server (for the `/inference/v1/generate` + `prompt_logprobs` endpoint): + +```toml +training_mode = "opd" +[trainer.loss] +teacher_tau = 0.5 + +[orchestrator.teacher.client] +base_url = ["http://teacher-server:8000/v1"] + +[orchestrator.teacher.model] +name = "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B" +``` + +### Pure distillation (no verification) + +For agentic environments where verification is expensive (code execution, tool use, multi-turn interactions), skip verification and use only the teacher signal: + +```toml +training_mode = "opd" +[deployment] +num_teacher_gpus = 2 + +[trainer.loss] +teacher_tau = 1.0 +adv_tau = 0.0 # disable reward-based learning + +[orchestrator.verification] +enabled = false # skip expensive verification +``` + +The student learns to match the teacher without needing any reward signal. + +### Monitoring + +The `teacher_kl` metric shows the KL divergence from teacher to student. Lower means the student is closer to the teacher. + +## SFT details + +### VLM support + +Image input is supported in SFT mode when the student is a VLM: + +- Prompts can include OpenAI-style image items in `message.content`, e.g. `{"type": "image_url", "image_url": {"url": "data:image/..."}}` +- The orchestrator extracts and preprocesses images from trajectory prompts and attaches `pixel_values` / `image_grid_thw` to training samples +- No teacher token IDs / logprobs are required; reconstruction still happens from messages + +Notes: +- This path currently expects `data:image/...` payloads in message content +- The teacher rollout endpoint must also handle the same multimodal prompts during generation + +### Reference configs + +- `configs/alphabet_sort/sft_distill_hard_qwen4b_lora_prime_teacher.toml` +- `examples/alphabet_sort/sft_distill_hard.toml` + +## Parameter reference + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `training_mode` | `"rl"` | One of `rl`, `opd`, `sft`. Propagates to `orchestrator.training_mode` and (for sft) `trainer.loss.type`. | +| `deployment.num_teacher_gpus` | `None` | Number of GPUs for the teacher vLLM server. Auto-starts when set. OPD only. | +| `trainer.loss.teacher_tau` | `0.0` | Distillation strength. Must be `> 0` in OPD. | +| `trainer.loss.adv_tau` | `1.0` | Weight for the RL advantage signal. Set `0` for pure distillation. | +| `orchestrator.verification.enabled` | `true` | Enable/disable verification. Set to `false` for pure distillation with `adv_tau = 0`. | 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 1891415ef2..7fcb141314 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py @@ -916,8 +916,18 @@ def _accept_flat_model_layout(cls, data): class OrchestratorConfig(BaseConfig): """Configures the orchestrator for RL training.""" - # Training environments and sampling - train: TrainConfig = TrainConfig() + # Training mode: drives validation and runtime wiring + training_mode: Annotated[ + Literal["rl", "opd", "sft"], + Field( + description=( + "Training mode. " + "rl: student generates rollouts, no teacher. " + "opd: student generates rollouts, teacher computes logprobs (teacher_tau > 0). " + "sft: teacher generates rollouts, student inference pool used for evals and weight sync." + ), + ), + ] = "rl" # Student model + client (the model being trained) student: Annotated[ @@ -928,15 +938,6 @@ class OrchestratorConfig(BaseConfig): ), ] = RolloutModelConfig() - # The tokenizer configuration - tokenizer: TokenizerConfig = TokenizerConfig() - - # The renderer configuration (only used when use_renderer=True) - renderer: RendererConfig = RendererConfig() - - # The optimizer configuration (per-run LR for multi-run training) - optim: OptimizerConfig = OptimizerConfig() - # Teacher model + client (optional; role determined by training_mode) teacher: Annotated[ RolloutModelConfig | None, @@ -949,18 +950,17 @@ class OrchestratorConfig(BaseConfig): ), ] = None - # Training mode: drives validation and runtime wiring - training_mode: Annotated[ - Literal["rl", "opd", "sft"], - Field( - description=( - "Training mode. " - "rl: student generates rollouts, no teacher. " - "opd: student generates rollouts, teacher computes logprobs (teacher_tau > 0). " - "sft: teacher generates rollouts, student inference pool used for evals and weight sync." - ), - ), - ] = "rl" + # Training environments and sampling + train: TrainConfig = TrainConfig() + + # The tokenizer configuration + tokenizer: TokenizerConfig = TokenizerConfig() + + # The renderer configuration (only used when use_renderer=True) + renderer: RendererConfig = RendererConfig() + + # The optimizer configuration (per-run LR for multi-run training) + optim: OptimizerConfig = OptimizerConfig() # The evaluation configuration eval: EvalConfig | None = None @@ -1153,6 +1153,23 @@ class OrchestratorConfig(BaseConfig): Field(description="Experimental features for the orchestrator."), ] = OrchestratorExperimentalConfig() + @model_validator(mode="before") + @classmethod + def _accept_top_level_client(cls, data: Any) -> Any: + """Accept legacy [orchestrator.client] as shorthand for [orchestrator.student.client]. + + Pre-refactor, OrchestratorConfig had a top-level `client` field. After the + student/teacher rename it moved under `student.client`. Re-nest a top-level + `client` dict so existing configs keep working. + """ + if not isinstance(data, dict) or "client" not in data: + return data + student = data.setdefault("student", {}) + # If the user already nested model+client under student/model, leave it alone. + if isinstance(student, dict) and "client" not in student: + student["client"] = data.pop("client") + return data + @model_validator(mode="before") @classmethod def _env_to_train(cls, data: Any) -> Any: diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index 190911b224..7bad4015f7 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -92,15 +92,16 @@ async def orchestrate(config: OrchestratorConfig): json_logging=config.log.json_logging, ) intercept_vf_logging(logger="verifiers.serve", level="WARN") # show logs from env clients - logger.info("Starting orchestrator") - _MODE_DESCRIPTIONS = { - "rl": "student generates rollouts, trained with reward-based advantage", - "opd": "student generates rollouts, trained on reward + KL to teacher logprobs (on-policy distillation)", - "sft": "teacher generates rollouts, student trained on teacher tokens (hard distillation)", + + # Print start message + mode_descriptions = { + "rl": "student generates rollouts, no teacher", + "opd": "student generates rollouts, teacher judges", + "sft": "teacher generates rollouts, student trains on teacher tokens", } - logger.info(f"Training mode: {config.training_mode} - {_MODE_DESCRIPTIONS[config.training_mode]}") - set_default_executor() + logger.info(f"Starting orchestrator in {config.training_mode} mode ({mode_descriptions[config.training_mode]})") + set_default_executor() event_loop_lag_monitor = EventLoopLagMonitor() event_loop_lag_monitor_task = asyncio.create_task(event_loop_lag_monitor.run()) diff --git a/src/prime_rl/utils/client.py b/src/prime_rl/utils/client.py index 511c32fff3..a37b0c632f 100644 --- a/src/prime_rl/utils/client.py +++ b/src/prime_rl/utils/client.py @@ -167,11 +167,6 @@ async def setup_inference_pool( preserve_thinking_between_tool_calls=preserve_thinking_between_tool_calls, ) - logger.info( - f"Initializing static inference pool (base_url={', '.join(client_config.base_url)}, " - f"dp_rank_count={client_config.dp_rank_count}, " - f"api_key_var={client_config.api_key_var}, headers={client_config.headers})" - ) return StaticInferencePool( client_config, model_name=model_name, diff --git a/src/prime_rl/utils/monitor/prime.py b/src/prime_rl/utils/monitor/prime.py index 7e63b6c2d7..657037c6f7 100644 --- a/src/prime_rl/utils/monitor/prime.py +++ b/src/prime_rl/utils/monitor/prime.py @@ -134,7 +134,6 @@ def __init__( return assert config is not None - self.logger.info(f"Initializing {self.__class__.__name__} ({config})") api_key = os.getenv(config.api_key_var) if api_key is None: diff --git a/src/prime_rl/utils/monitor/wandb.py b/src/prime_rl/utils/monitor/wandb.py index 6474dc7a32..27a4192dec 100644 --- a/src/prime_rl/utils/monitor/wandb.py +++ b/src/prime_rl/utils/monitor/wandb.py @@ -44,7 +44,6 @@ def __init__( return assert config is not None - self.logger.info(f"Initializing {self.__class__.__name__} ({config})") self._maybe_overwrite_wandb_command() # WANDB_MODE=disabled/offline takes precedence over shared mode — shared mode From 5c25df0a3c68e92a4ccc9b19d7ba9d9990cefa92 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Mon, 18 May 2026 23:59:35 +0000 Subject: [PATCH 23/47] feat(orchestrator): probe teacher for prompt_logprobs before training In OPD mode, send a tiny probe request to the teacher's /inference/v1/generate endpoint after the pool is ready. If the endpoint 404s or doesn't return prompt_logprobs, raise an informative error pointing at docs/training_modes.md instead of crashing mid-training with a 404. Verified that PI inference fails this probe on all three plausible routes: - /inference/v1/generate -> 404 - /v1/chat/completions with logprobs=true -> logprobs:null in response - /v1/completions with echo=true,logprobs=1 -> 404 Also add configs/reverse_text/debug_sft_thinking.toml: SFT from qwen3-30b-a3b-thinking-2507 via PI inference (sft path needs only chat completions, so PI inference is fine here). Co-Authored-By: Claude Sonnet 4.6 --- configs/reverse_text/README.md | 5 ++ configs/reverse_text/debug_sft_thinking.toml | 52 +++++++++++++++++++ src/prime_rl/orchestrator/orchestrator.py | 5 ++ src/prime_rl/orchestrator/utils.py | 53 ++++++++++++++++++++ 4 files changed, 115 insertions(+) create mode 100644 configs/reverse_text/debug_sft_thinking.toml diff --git a/configs/reverse_text/README.md b/configs/reverse_text/README.md index 74a6d27981..9402778a19 100644 --- a/configs/reverse_text/README.md +++ b/configs/reverse_text/README.md @@ -7,6 +7,7 @@ Minimal end-to-end configs for the three training modes against the `reverse-tex | `debug_rl.toml` | `rl` | none | | `debug_opd.toml` | `opd` | local vLLM (`Qwen3-0.6B-Reverse-Text-RL`) | | `debug_sft.toml` | `sft` | local vLLM (`Qwen3-0.6B-Reverse-Text-RL`) | +| `debug_sft_thinking.toml` | `sft` | PI inference (`qwen/qwen3-30b-a3b-thinking-2507`) | The student inference server is auto-launched on GPU 0 at `http://localhost:8000/v1` with `gpu_memory_utilization=0.5`. The teacher (used by `debug_opd.toml` and `debug_sft.toml`) is **not** auto-launched — start it manually on GPU 1. @@ -31,6 +32,10 @@ uv run rl @ configs/reverse_text/debug_opd.toml # SFT hard distill (needs teacher on port 8001) uv run rl @ configs/reverse_text/debug_sft.toml + +# SFT hard distill from qwen3-30b-a3b-thinking via PI inference +# (requires PRIME_API_KEY + PRIME_TEAM_ID in env; no local teacher needed) +uv run rl @ configs/reverse_text/debug_sft_thinking.toml ``` See [docs/training_modes.md](../../docs/training_modes.md) for what each mode does. diff --git a/configs/reverse_text/debug_sft_thinking.toml b/configs/reverse_text/debug_sft_thinking.toml new file mode 100644 index 0000000000..514615fba4 --- /dev/null +++ b/configs/reverse_text/debug_sft_thinking.toml @@ -0,0 +1,52 @@ +# SFT from qwen3-30b-a3b-thinking-2507 via PI inference. +# X-Prime-Team-ID header is auto-injected from $PRIME_TEAM_ID for pinference.ai URLs. +# +# Run with: +# uv run rl @ configs/reverse_text/debug_sft_thinking.toml + +max_steps = 20 +seq_len = 2048 +training_mode = "sft" + +[model] +name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT" + +[wandb] +project = "reverse-text-debug" +name = "debug-sft-thinking" + +[orchestrator] +batch_size = 128 +rollouts_per_example = 4 +use_renderer = false + +[orchestrator.train.sampling] +max_completion_tokens = 128 + +[[orchestrator.train.env]] +id = "reverse-text" + +[orchestrator.eval] +interval = 1 +num_examples = 128 + +[orchestrator.eval.sampling] +max_completion_tokens = 128 + +[[orchestrator.eval.env]] +id = "reverse-text" + +[orchestrator.teacher.model] +name = "qwen/qwen3-30b-a3b-thinking-2507" + +[orchestrator.teacher.client] +base_url = ["https://api.pinference.ai/api/v1"] +api_key_var = "PRIME_API_KEY" + +[trainer.optim] +lr = 3e-6 + +[ckpt] + +[inference] +gpu_memory_utilization = 0.5 diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index 7bad4015f7..6bde5e6b66 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -46,6 +46,7 @@ compute_teacher_logprobs, get_weight_dir, print_benchmark, + probe_teacher_logprobs, set_default_executor, ) from prime_rl.orchestrator.vf_utils import ( @@ -275,6 +276,10 @@ async def orchestrate(config: OrchestratorConfig): logger.info("Waiting for teacher inference pool to be ready") await teacher_inference.wait_for_ready(config.teacher.model.name) logger.success("Teacher inference pool ready") + if config.training_mode == "opd": + logger.info("Probing teacher for prompt_logprobs support") + await probe_teacher_logprobs(teacher_inference.train_clients, config.teacher.model.name) + logger.success("Teacher supports prompt_logprobs") # Start inference metrics collector (requires W&B + student inference pool) inference_metrics_collector = None diff --git a/src/prime_rl/orchestrator/utils.py b/src/prime_rl/orchestrator/utils.py index 121cc57c54..9458ea64a5 100644 --- a/src/prime_rl/orchestrator/utils.py +++ b/src/prime_rl/orchestrator/utils.py @@ -133,6 +133,59 @@ async def _compute_single(client_config: vf.ClientConfig, sample: TrainingSample return await asyncio.gather(*[_compute_single(client, sample) for client, sample in zip(cycle(clients), samples)]) +async def probe_teacher_logprobs(clients: list[vf.ClientConfig], model_name: str) -> None: + """Probe the teacher endpoint for vLLM `prompt_logprobs` support. + + OPD requires the vLLM-specific ``/inference/v1/generate`` endpoint with + ``prompt_logprobs=1`` (see ``compute_teacher_logprobs``). External + OAI-compatible endpoints (PI inference, OpenAI, Anthropic) don't expose + this and would 404 mid-training. Probe once at startup with a tiny + request and raise a clear error if the endpoint can't serve it. + """ + import httpx + from vllm.entrypoints.serve.disagg.protocol import GenerateResponse + + client_config = clients[0] + client = setup_openai_client(client_config) + base = str(client.base_url).rstrip("/").removesuffix("/v1") + url = f"{base}/inference/v1/generate" + + hint = ( + "OPD requires a self-hosted vLLM teacher exposing /inference/v1/generate with " + "prompt_logprobs. External OAI-compatible endpoints (PI inference, OpenAI, Anthropic) " + "are not supported. See docs/training_modes.md." + ) + + try: + http_response = await client.post( + url, + cast_to=httpx.Response, + body={ + "model": model_name, + "token_ids": [1, 2, 3], + "sampling_params": { + "max_tokens": 1, + "temperature": 1.0, + "top_p": 1.0, + "prompt_logprobs": 1, + }, + }, + ) + except Exception as e: + raise RuntimeError(f"OPD teacher probe failed: {url} - {type(e).__name__}: {e}\n{hint}") from e + + try: + response = GenerateResponse.model_validate_json(http_response.content) + except Exception as e: + snippet = http_response.content[:200] if hasattr(http_response, "content") else "" + raise RuntimeError( + f"OPD teacher probe failed: {url} returned a non-vLLM response (got {snippet!r}).\n{hint}" + ) from e + + if not response.prompt_logprobs: + raise RuntimeError(f"OPD teacher probe failed: {url} returned no prompt_logprobs in the response.\n{hint}") + + def get_weight_dir(output_dir: Path, step: int, check_exists: bool = True, wait_timeout: int | None = None) -> Path: """Get the weight directory for a given checkpoint step. From 20c59a546987c7158ec6fc7535efd885498c65f6 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Tue, 19 May 2026 00:19:18 +0000 Subject: [PATCH 24/47] chore(orchestrator): deprecate use_token_client (server-tokenized TITO) Drop the server-tokenized TITO path (``openai_chat_completions_token`` / ``/v1/chat/completions/tokens``). The orchestrator now picks between renderer-backed TITO (``use_renderer = true``, default) and MITO (``use_renderer = false``, fallback for VLMs and external teacher rollouts). - Remove ``OrchestratorConfig.use_token_client`` field and its validators - Drop the TITO branch in ``setup_rollout_inference_pool`` - Drop the obsolete TITO warning in ``setup_inference_pool`` and flip the default ``train_client_type`` to ``openai_chat_completions`` in ``StaticInferencePool`` / ``ElasticInferencePool`` - Strip ``use_token_client = false`` from configs/examples/docs that carried the no-op fallback line - CHANGELOG entry documenting the breaking removal Co-Authored-By: Claude Opus 4.7 (1M context) --- CHANGELOG.md | 1 + .../ci/nightly/multimodal_color_codeword.toml | 1 - configs/multimodal/rl_color_codeword.toml | 1 - .../multimodal/rl_color_codeword_test.toml | 1 - docs/on_policy_distillation.md | 1 - examples/alphabet_sort/sft_distill_hard.toml | 1 - .../src/prime_rl/configs/orchestrator.py | 45 +++---------------- .../src/prime_rl/configs/rl.py | 2 +- skills/config/SKILL.md | 2 +- src/prime_rl/orchestrator/orchestrator.py | 26 ++++------- src/prime_rl/utils/client.py | 11 +---- src/prime_rl/utils/elastic.py | 4 +- .../orchestrator/test_orchestrator_setup.py | 1 - tests/unit/test_configs.py | 2 - 14 files changed, 22 insertions(+), 77 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9e0de85b98..0fb5069651 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,7 @@ Documenting **breaking** configuration changes — renamed, removed, or moved fields that require users to update existing configs. +- **`orchestrator.use_token_client` removed**: The server-tokenized TITO path (`openai_chat_completions_token` / `/v1/chat/completions/tokens`) has been deprecated. The orchestrator now picks between renderer-backed TITO (`use_renderer = true`, default) and MITO (`use_renderer = false`, fallback). Existing configs with `use_token_client = true` must migrate to `use_renderer = true` (or `use_renderer = false` for MITO); configs with `use_token_client = false` can simply drop the field. (2026-05-19) - **`orchestrator.advantage.length_penalty` → discriminated sub-config**: The scalar `length_penalty: Literal["tokens","turns"] | None` is replaced by a `LengthPenaltyConfig | None` discriminated on `type`. Token shaping now takes weighted completion + tool-response token costs. Migration: `length_penalty = "tokens"` becomes `[orchestrator.advantage.length_penalty]\ntype = "tokens"` (default weights `completion_weight = 1.0`, `tool_response_weight = 1.0` — total context). `length_penalty = "turns"` becomes `[orchestrator.advantage.length_penalty]\ntype = "turns"`. (2026-05-06) - **`orchestrator.advantage.length_shaping` → `orchestrator.advantage.length_penalty`**: The boolean `length_shaping` flag has been replaced by `length_penalty: Literal["tokens", "turns"] | None` (default: `None`). `length_shaping = true` becomes `length_penalty = "tokens"`; `length_shaping = false` becomes `length_penalty = None`. The new `"turns"` option applies the same correctness-gated efficiency shaping using trajectory turn count instead of completion-token count. (2026-05-01) - **`AdvantageInputs` API**: Replaced the `rewards`/`completion_lengths`/`num_turns` tensor fields with a single `rollouts: list[list[vf.RolloutOutput]]` (grouped by problem). Custom advantage functions can now access any rollout metadata. Existing custom advantages must update their signatures and extract per-rollout fields directly (e.g. `torch.tensor([[r["reward"] for r in g] for g in inputs.rollouts])`). (2026-05-01) diff --git a/configs/ci/nightly/multimodal_color_codeword.toml b/configs/ci/nightly/multimodal_color_codeword.toml index 4c7f44b557..173193c81c 100644 --- a/configs/ci/nightly/multimodal_color_codeword.toml +++ b/configs/ci/nightly/multimodal_color_codeword.toml @@ -16,7 +16,6 @@ language_model_attr = "model.language_model" [orchestrator] batch_size = 256 rollouts_per_example = 16 -use_token_client = false use_renderer = false [orchestrator.train.sampling] diff --git a/configs/multimodal/rl_color_codeword.toml b/configs/multimodal/rl_color_codeword.toml index 034a0ff662..a357bf0598 100644 --- a/configs/multimodal/rl_color_codeword.toml +++ b/configs/multimodal/rl_color_codeword.toml @@ -11,7 +11,6 @@ language_model_attr = "model.language_model" [orchestrator] batch_size = 256 rollouts_per_example = 16 -use_token_client = false use_renderer = false diff --git a/configs/multimodal/rl_color_codeword_test.toml b/configs/multimodal/rl_color_codeword_test.toml index 50e3170653..643cbd34f0 100644 --- a/configs/multimodal/rl_color_codeword_test.toml +++ b/configs/multimodal/rl_color_codeword_test.toml @@ -12,7 +12,6 @@ language_model_attr = "model.language_model" [orchestrator] batch_size = 16 rollouts_per_example = 2 -use_token_client = false use_renderer = false [orchestrator.train.sampling] diff --git a/docs/on_policy_distillation.md b/docs/on_policy_distillation.md index 6f434bbfe1..e57e7e3386 100644 --- a/docs/on_policy_distillation.md +++ b/docs/on_policy_distillation.md @@ -71,7 +71,6 @@ Use this mode when you want to train from teacher-generated completions directly type = "sft" [orchestrator] -use_token_client = false use_renderer = false use_sft_loss = true diff --git a/examples/alphabet_sort/sft_distill_hard.toml b/examples/alphabet_sort/sft_distill_hard.toml index 8eab8c9dc7..9e3ef0a082 100644 --- a/examples/alphabet_sort/sft_distill_hard.toml +++ b/examples/alphabet_sort/sft_distill_hard.toml @@ -29,7 +29,6 @@ save_adapter_separately = true [orchestrator] batch_size = 256 rollouts_per_example = 4 -use_token_client = false use_renderer = false use_sft_loss = true 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 d6921ad338..0ff529bcd7 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py @@ -1119,24 +1119,14 @@ class OrchestratorConfig(BaseConfig): HeartbeatConfig | None, Field(description="The heartbeat config for monitoring training progress.") ] = None - use_token_client: Annotated[ - bool, - Field( - description="Whether to use the server-tokenized token-in-token-out (TITO) client for training across all environments. " - "WARNING: Only use this if your environment has a linear history and the chat template has the extension " - "property (i.e. no tokens are ever removed or inserted by the chat template). Mutually exclusive with " - "``use_renderer``." - ), - ] = False - use_renderer: Annotated[ bool, Field( description="Whether to use the renderer-backed TITO client (client-side tokenization via the ``renderers`` package, " - "served by ``/v1/generate``). Mutually exclusive with ``use_token_client``. When True, the " - "``[orchestrator.renderer]`` block (name / tool_parser / reasoning_parser / pool_size) applies. " - "This is the default for text-only rollouts. Not supported for VLMs — VLMs must use MITO so " - "image preprocessing and chat templating stay server-side." + "served by ``/v1/generate``). When True, the ``[orchestrator.renderer]`` block " + "(name / tool_parser / reasoning_parser / pool_size) applies. This is the default for text-only " + "rollouts. Set to False to fall back to MITO (``openai_chat_completions``); VLMs and external " + "teacher rollouts require MITO." ), ] = True @@ -1209,10 +1199,6 @@ def validate_sft_distill_mode(self): ) if has_teacher and not self.use_sft_loss: raise ValueError("orchestrator.teacher_rollout_model requires orchestrator.use_sft_loss = true.") - if has_teacher and self.use_token_client: - raise ValueError( - "orchestrator.use_token_client must be false when orchestrator.teacher_rollout_model is configured." - ) if has_teacher and self.use_renderer: raise ValueError( "orchestrator.use_renderer must be false when orchestrator.teacher_rollout_model is configured " @@ -1220,25 +1206,6 @@ def validate_sft_distill_mode(self): ) return self - @model_validator(mode="after") - def validate_client_mode(self): - """The two client toggles select among three exclusive modes: - - - ``use_token_client=False`` + ``use_renderer=True`` → renderer-backed TITO (default) - - ``use_token_client=True`` + ``use_renderer=False`` → server-tokenized TITO - - ``use_token_client=False`` + ``use_renderer=False`` → MITO - - Both True is invalid: renderer-backed TITO and server-tokenized TITO are - different wire protocols (client-side vs server-side tokenization). - """ - if self.use_token_client and self.use_renderer: - raise ValueError( - "orchestrator.use_token_client and orchestrator.use_renderer are mutually exclusive. " - "Pick one TITO path: renderer client (client-side tokenization) or token client " - "(server-side tokenization)." - ) - return self - @model_validator(mode="after") def validate_renderer_vs_vlm(self): """The renderer client takes plain message dicts and tokenizes @@ -1248,8 +1215,8 @@ def validate_renderer_vs_vlm(self): if self.use_renderer and self.model.vlm is not None: raise ValueError( "orchestrator.use_renderer is not supported for VLMs. Use MITO " - "(``use_token_client=false`` and ``use_renderer=false``) so image preprocessing and chat " - "templating stay on the inference server." + "(``use_renderer=false``) so image preprocessing and chat templating stay on the " + "inference server." ) return self 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 a160af2c9f..b55d1314b2 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/rl.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/rl.py @@ -438,7 +438,7 @@ def validate_external_rollout_inference(self): """Forbid configuring a local inference server when rollouts come from an external teacher. Orchestrator-only invariants (``use_sft_loss`` paired with ``teacher_rollout_model``, - and ``use_token_client`` coupling) live on ``OrchestratorConfig`` so the hosted + and ``use_renderer`` coupling) live on ``OrchestratorConfig`` so the hosted orchestrator entrypoint also enforces them. """ if self.orchestrator.teacher_rollout_model is not None and self.inference is not None: diff --git a/skills/config/SKILL.md b/skills/config/SKILL.md index b8f8c48780..c4a33a592b 100644 --- a/skills/config/SKILL.md +++ b/skills/config/SKILL.md @@ -159,7 +159,7 @@ For hosted multi-tenant runs where the trainer image's `trainer.loss.type` is fi ### RL rollout client defaults -For text-only RL rollouts, the orchestrator defaults to renderer-backed TITO (`use_renderer = true`, `use_token_client = false`). VLM configs must explicitly use MITO (`use_token_client = false`, `use_renderer = false`) so image preprocessing and chat templating stay server-side. External teacher rollouts must also set `use_renderer = false`. +For text-only RL rollouts, the orchestrator defaults to renderer-backed TITO (`use_renderer = true`). VLM configs must explicitly fall back to MITO (`use_renderer = false`) so image preprocessing and chat templating stay server-side. External teacher rollouts must also set `use_renderer = false`. ### Model fields diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index d99208108c..e132c5980f 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -904,20 +904,16 @@ async def setup_rollout_inference_pool( ): """Set up rollout inference. - Routing policy is driven by ``config.use_token_client`` and - ``config.use_renderer`` (mutually exclusive — config-level validators - block both being True): + Routing policy is driven by ``config.use_renderer``: - external teacher rollout → MITO (``openai_chat_completions``), - forced regardless of the toggles (config-level validator - rejects ``use_token_client`` / ``use_renderer`` in that case) + forced regardless of the toggle (config-level validator rejects + ``use_renderer`` in that case) - ``use_renderer=True`` → renderer-backed TITO client (``/v1/generate``). - Default for text-only rollouts. - Not allowed for VLMs (validated at config time). - - ``use_token_client=True`` → server-tokenized TITO - (``openai_chat_completions_token``, ``/v1/chat/completions/tokens``). - - both False → MITO (``openai_chat_completions``). - VLMs land here too. + Default for text-only rollouts. Not allowed for VLMs (validated at + config time). + - ``use_renderer=False`` → MITO (``openai_chat_completions``). VLMs + land here too. """ if config.teacher_rollout_model is not None: logger.info("Using external rollout model (MITO) without renderer client") @@ -954,15 +950,11 @@ async def setup_rollout_inference_pool( logger.info("Using direct renderer rollout client") return renderer, inference_pool - train_client_type = "openai_chat_completions_token" if config.use_token_client else "openai_chat_completions" - if config.use_token_client: - logger.info("Using server-tokenized TITO for rollouts — /v1/chat/completions/tokens") - else: - logger.info("Using MITO (openai_chat_completions) for rollouts") + logger.info("Using MITO (openai_chat_completions) for rollouts") inference_pool = await setup_inference_pool( rollout_client_config, model_name=rollout_model_name, - train_client_type=train_client_type, + train_client_type="openai_chat_completions", eval_client_type="openai_chat_completions", ) return None, inference_pool diff --git a/src/prime_rl/utils/client.py b/src/prime_rl/utils/client.py index 9f59d1a2b1..bd053ad55c 100644 --- a/src/prime_rl/utils/client.py +++ b/src/prime_rl/utils/client.py @@ -62,7 +62,7 @@ def __init__( self, client_config: ClientConfig, model_name: str, - train_client_type: str = "openai_chat_completions_token", + train_client_type: str = "openai_chat_completions", eval_client_type: str = "openai_chat_completions", renderer_name: str = "auto", tool_parser: str | None = None, @@ -127,7 +127,7 @@ async def stop(self) -> None: async def setup_inference_pool( client_config: ClientConfig, model_name: str, - train_client_type: str = "openai_chat_completions_token", + train_client_type: str = "openai_chat_completions", eval_client_type: str = "openai_chat_completions", renderer_name: str = "auto", tool_parser: str | None = None, @@ -139,13 +139,6 @@ async def setup_inference_pool( """Create an inference pool from config (static or elastic).""" logger = get_logger() - if train_client_type == "openai_chat_completions_token": - logger.warning( - "Token-in-token-out (TITO) client is enabled for training. Only use " - "this if your environment has a linear history and the chat " - "template has the extension property." - ) - if client_config.is_elastic: from prime_rl.utils.elastic import ElasticInferencePool diff --git a/src/prime_rl/utils/elastic.py b/src/prime_rl/utils/elastic.py index c59f81e27f..19a2b2f2af 100644 --- a/src/prime_rl/utils/elastic.py +++ b/src/prime_rl/utils/elastic.py @@ -104,7 +104,7 @@ def __init__( self, client_config: ClientConfig, model_name: str, - train_client_type: str = "openai_chat_completions_token", + train_client_type: str = "openai_chat_completions", eval_client_type: str = "openai_chat_completions", renderer_name: str = "auto", tool_parser: str | None = None, @@ -150,7 +150,7 @@ async def from_config( cls, client_config: ClientConfig, model_name: str, - train_client_type: str = "openai_chat_completions_token", + train_client_type: str = "openai_chat_completions", eval_client_type: str = "openai_chat_completions", renderer_name: str = "auto", tool_parser: str | None = None, diff --git a/tests/unit/orchestrator/test_orchestrator_setup.py b/tests/unit/orchestrator/test_orchestrator_setup.py index d4567d9682..e981d051b5 100644 --- a/tests/unit/orchestrator/test_orchestrator_setup.py +++ b/tests/unit/orchestrator/test_orchestrator_setup.py @@ -43,7 +43,6 @@ async def run() -> None: config = SimpleNamespace( teacher_rollout_model=None, use_renderer=True, - use_token_client=False, model=SimpleNamespace(name="student-model"), renderer=SimpleNamespace( name="qwen3_vl", diff --git a/tests/unit/test_configs.py b/tests/unit/test_configs.py index 0af78319ff..f5b5dc130c 100644 --- a/tests/unit/test_configs.py +++ b/tests/unit/test_configs.py @@ -178,12 +178,10 @@ def test_orchestrator_vlm_configs_must_disable_renderer(): "language_model_attr": "model.language_model", } }, - "use_token_client": False, "use_renderer": False, } ) - assert config.use_token_client is False assert config.use_renderer is False From 7b04ecadb5333ffc8036b6d59d62aec5428f55e4 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Tue, 19 May 2026 00:31:28 +0000 Subject: [PATCH 25/47] chore(inference): remove server-side TITO route Drop the ``/v1/chat/completions/tokens`` endpoint and the ``OpenAIServingChatWithTokens`` / ``ChatCompletionRequestWithTokens`` wrappers now that no client routes to them. Also drop the obsolete ``base()`` helper and unused imports from ``server.py``. Co-Authored-By: Claude Opus 4.7 (1M context) --- CHANGELOG.md | 2 +- skills/entrypoints/SKILL.md | 1 - src/prime_rl/inference/vllm/server.py | 62 +---- .../vllm/serving_chat_with_tokens.py | 243 ------------------ src/prime_rl/inference/vllm/serving_tokens.py | 17 +- 5 files changed, 11 insertions(+), 314 deletions(-) delete mode 100644 src/prime_rl/inference/vllm/serving_chat_with_tokens.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 0fb5069651..59b0b126f6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,7 +2,7 @@ Documenting **breaking** configuration changes — renamed, removed, or moved fields that require users to update existing configs. -- **`orchestrator.use_token_client` removed**: The server-tokenized TITO path (`openai_chat_completions_token` / `/v1/chat/completions/tokens`) has been deprecated. The orchestrator now picks between renderer-backed TITO (`use_renderer = true`, default) and MITO (`use_renderer = false`, fallback). Existing configs with `use_token_client = true` must migrate to `use_renderer = true` (or `use_renderer = false` for MITO); configs with `use_token_client = false` can simply drop the field. (2026-05-19) +- **`orchestrator.use_token_client` removed**: The server-tokenized TITO path has been deprecated end-to-end. The orchestrator-side config flag (`use_token_client`), the verifiers client_type (`openai_chat_completions_token`), and the inference server's `/v1/chat/completions/tokens` route (along with the `OpenAIServingChatWithTokens` wrapper) are all gone. The orchestrator now picks between renderer-backed TITO (`use_renderer = true`, default) and MITO (`use_renderer = false`, fallback). Existing configs with `use_token_client = true` must migrate to `use_renderer = true` (or `use_renderer = false` for MITO); configs with `use_token_client = false` can simply drop the field. (2026-05-19) - **`orchestrator.advantage.length_penalty` → discriminated sub-config**: The scalar `length_penalty: Literal["tokens","turns"] | None` is replaced by a `LengthPenaltyConfig | None` discriminated on `type`. Token shaping now takes weighted completion + tool-response token costs. Migration: `length_penalty = "tokens"` becomes `[orchestrator.advantage.length_penalty]\ntype = "tokens"` (default weights `completion_weight = 1.0`, `tool_response_weight = 1.0` — total context). `length_penalty = "turns"` becomes `[orchestrator.advantage.length_penalty]\ntype = "turns"`. (2026-05-06) - **`orchestrator.advantage.length_shaping` → `orchestrator.advantage.length_penalty`**: The boolean `length_shaping` flag has been replaced by `length_penalty: Literal["tokens", "turns"] | None` (default: `None`). `length_shaping = true` becomes `length_penalty = "tokens"`; `length_shaping = false` becomes `length_penalty = None`. The new `"turns"` option applies the same correctness-gated efficiency shaping using trajectory turn count instead of completion-token count. (2026-05-01) - **`AdvantageInputs` API**: Replaced the `rewards`/`completion_lengths`/`num_turns` tensor fields with a single `rollouts: list[list[vf.RolloutOutput]]` (grouped by problem). Custom advantage functions can now access any rollout metadata. Existing custom advantages must update their signatures and extract per-rollout fields directly (e.g. `torch.tensor([[r["reward"] for r in g] for g in inputs.rollouts])`). (2026-05-01) diff --git a/skills/entrypoints/SKILL.md b/skills/entrypoints/SKILL.md index 0ec6c52f41..45ec77aee5 100644 --- a/skills/entrypoints/SKILL.md +++ b/skills/entrypoints/SKILL.md @@ -49,7 +49,6 @@ uv run inference --model.name Qwen/Qwen3-0.6B --model.enforce-eager Always use the `inference` entrypoint — never `vllm serve` directly. Custom endpoints beyond standard OpenAI API: -- `/v1/chat/completions/tokens` — accepts token IDs as prompt input - `/update_weights` — hot-reload model weights from the trainer - `/load_lora_adapter` — load LoRA adapters at runtime - `/init_broadcaster` — initialize weight broadcast for distributed training diff --git a/src/prime_rl/inference/vllm/server.py b/src/prime_rl/inference/vllm/server.py index 53ae22c104..b8aae37148 100644 --- a/src/prime_rl/inference/vllm/server.py +++ b/src/prime_rl/inference/vllm/server.py @@ -1,22 +1,17 @@ import asyncio from argparse import Namespace -from http import HTTPStatus from typing import Any import uvloop -from fastapi import APIRouter, Depends, Request -from fastapi.responses import JSONResponse, StreamingResponse +from fastapi import APIRouter, Request +from fastapi.responses import JSONResponse from starlette.datastructures import State from vllm.engine.protocol import EngineClient from vllm.entrypoints.openai.api_server import init_app_state -from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionResponse from vllm.entrypoints.openai.cli_args import make_arg_parser, validate_parsed_serve_args from vllm.entrypoints.openai.engine.protocol import ErrorResponse -from vllm.entrypoints.openai.engine.serving import OpenAIServing from vllm.entrypoints.openai.models.serving import OpenAIServingModels -from vllm.entrypoints.openai.utils import validate_json_request from vllm.entrypoints.serve.lora.protocol import LoadLoRAAdapterRequest -from vllm.entrypoints.utils import load_aware_call, with_cancellation from vllm.logger import init_logger from vllm.utils.argparse_utils import FlexibleArgumentParser @@ -140,10 +135,6 @@ def resolve_tool_call_parser(model_name: str, tool_call_parser: str | None) -> s monkey_patch_load_lora_adapter, monkey_patch_tokenize_params_validation, ) -from prime_rl.inference.vllm.serving_chat_with_tokens import ( - ChatCompletionRequestWithTokens, - OpenAIServingChatWithTokens, -) # NOTE: Fix harmony stop token propagation for GPT-OSS models # Upstream issue still open: https://github.com/vllm-project/vllm/issues/22519 @@ -165,10 +156,6 @@ def engine_client(request: Request) -> EngineClient: return request.app.state.engine_client -def base(request: Request) -> OpenAIServing: - return request.app.state.openai_serving_tokenization - - def models(request: Request) -> OpenAIServingModels: return request.app.state.openai_serving_models @@ -179,36 +166,6 @@ def models(request: Request) -> OpenAIServingModels: } -def chat_with_tokens(request: Request) -> OpenAIServingChatWithTokens | None: - return request.app.state.openai_serving_chat_with_tokens - - -@router.post( - "/v1/chat/completions/tokens", - dependencies=[Depends(validate_json_request)], - responses={ - HTTPStatus.OK.value: {"content": {"text/event-stream": {}}}, - HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse}, - HTTPStatus.NOT_FOUND.value: {"model": ErrorResponse}, - HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse}, - }, -) -@with_cancellation -@load_aware_call -async def _chat_with_tokens(request: ChatCompletionRequestWithTokens, raw_request: Request): - handler = chat_with_tokens(raw_request) - if handler is None: - return base(raw_request).create_error_response(message="The model does not support Chat Completions API") - generator = await handler.create_chat_completion_with_tokens(request, raw_request) - if isinstance(generator, ErrorResponse): - return JSONResponse(content=generator.model_dump(), status_code=generator.error.code) - - elif isinstance(generator, ChatCompletionResponse): - return JSONResponse(content=generator.model_dump()) - - return StreamingResponse(content=generator, media_type="text/event-stream") - - @router.post("/pause") async def pause(request: Request): await engine_client(request).pause_generation(mode="keep", clear_cache=False) @@ -277,10 +234,7 @@ async def custom_init_app_state( Modifies init_app_state: 1. Call the original init_app_state to set up standard state, including vLLM 0.20's ``serving_tokens`` for ``/inference/v1/generate``. - 2. Replace ``serving_chat`` with our ``OpenAIServingChatWithTokens`` wrapper - so the ``/v1/chat/completions/tokens`` (TITO) endpoint can stream - token IDs alongside the rendered chat completion. - 3. Replace ``serving_tokens`` with ``PrimeRlServingTokens`` so DP-rank + 2. Replace ``serving_tokens`` with ``PrimeRlServingTokens`` so DP-rank routing and ``routed_experts`` export survive the migration off the legacy ``/v1/generate`` endpoint. """ @@ -289,16 +243,6 @@ async def custom_init_app_state( state.reset_prefix_cache_after_update = getattr(args, "reset_prefix_cache_after_update", True) state.liveness_timeout_seconds = args.liveness_timeout_seconds - # TITO: server-side chat templating + token IDs. - if "generate" in supported_tasks and state.openai_serving_chat is not None: - original_chat = state.openai_serving_chat - serving_chat = object.__new__(OpenAIServingChatWithTokens) - serving_chat.__dict__.update(original_chat.__dict__) - state.openai_serving_chat = serving_chat - state.openai_serving_chat_with_tokens = serving_chat - else: - state.openai_serving_chat_with_tokens = None - # Swap in our ServingTokens subclass for /inference/v1/generate so the # X-data-parallel-rank header and routed_experts response field — both # used by prime-RL's renderer / router-replay paths — keep working. diff --git a/src/prime_rl/inference/vllm/serving_chat_with_tokens.py b/src/prime_rl/inference/vllm/serving_chat_with_tokens.py deleted file mode 100644 index fae9465fbe..0000000000 --- a/src/prime_rl/inference/vllm/serving_chat_with_tokens.py +++ /dev/null @@ -1,243 +0,0 @@ -from collections.abc import AsyncGenerator, AsyncIterator -from typing import ClassVar, Optional, Union - -from fastapi import Request -from pydantic import Field -from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest, ChatCompletionResponse -from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat -from vllm.entrypoints.openai.engine.protocol import ErrorResponse, RequestResponseMetadata -from vllm.entrypoints.openai.engine.serving import GenerationError -from vllm.entrypoints.utils import get_max_tokens -from vllm.exceptions import VLLMValidationError -from vllm.logger import init_logger -from vllm.outputs import RequestOutput -from vllm.reasoning import ReasoningParser -from vllm.sampling_params import BeamSearchParams, SamplingParams - -from prime_rl.inference.vllm.serving_tokens import _RoutedExpertsCaptureBase - -logger = init_logger(__name__) - - -class _RoutedExpertsCapture(_RoutedExpertsCaptureBase): - """Chat-endpoint variant: mutates choices in-place because - ``ChatCompletionResponseChoice`` is ``extra='allow'``, so an extra - ``routed_experts`` attribute survives serialization.""" - - def post_process(self, response: ChatCompletionResponse) -> None: - for choice in response.choices: - if choice.index in self.routed_experts: - choice.routed_experts = self.routed_experts[choice.index] - - -class ChatCompletionRequestWithTokens(ChatCompletionRequest): - field_names: ClassVar[Optional[set[str]]] = None - tokens: list[int] = Field(description=("Prompt tokens to use for the request.")) - - -class OpenAIServingChatWithTokens(OpenAIServingChat): - """OpenAI-compatible generate API that allows token-in and routed experts capture.""" - - async def chat_completion_full_generator( - self, - request: ChatCompletionRequest, - result_generator: AsyncIterator[RequestOutput], - request_id: str, - model_name: str, - conversation, - tokenizer, - request_metadata: RequestResponseMetadata, - reasoning_parser: ReasoningParser | None = None, - ) -> ErrorResponse | ChatCompletionResponse: - # We need to override the full_generator to be able to capture the routed experts - # By default, VLLM does not save the routed experts into ChatCompletionResponse.choices, so we need to capture them manually - # How this works: - # 1. We create a custom generator that encapsulates the original result_generator in self._generator - # 2. We override it's __aiter__ method to also capture the routed experts as an extra field in ChatCompletionResponse.choices - # 3. We override the full_generator method to use the custom generator instead of the original one if expert routing is enabled - if self.model_config.enable_return_routed_experts: - capture = _RoutedExpertsCapture(result_generator) - result_generator = capture - else: - capture = None - - response = await super().chat_completion_full_generator( - request, - result_generator, - request_id, - model_name, - conversation, - tokenizer, - request_metadata, - reasoning_parser, - ) - - if capture and isinstance(response, ChatCompletionResponse): - capture.post_process(response) - - return response - - async def create_chat_completion_with_tokens( - self, - request: ChatCompletionRequestWithTokens, - raw_request: Optional[Request] = None, - ) -> Union[AsyncGenerator[str, None], ChatCompletionResponse, ErrorResponse]: - """ - Chat Completion API similar to OpenAI's API. - - See https://platform.openai.com/docs/api-reference/chat/create - for the API specification. This API mimics the OpenAI - Chat Completion API. - """ - # Streaming response - tokenizer = self.renderer.tokenizer - assert tokenizer is not None - reasoning_parser: ReasoningParser | None = None - try: - if self.reasoning_parser_cls: - # Pass the same chat template kwargs as used in tokenization - chat_template_kwargs = self._prepare_extra_chat_template_kwargs( - request.chat_template_kwargs, - self.default_chat_template_kwargs, - ) - reasoning_parser = self.reasoning_parser_cls( - tokenizer, - chat_template_kwargs=chat_template_kwargs, # type: ignore[call-arg] - ) - except RuntimeError as e: - logger.exception("Error in reasoning parser creation.") - return self.create_error_response(str(e)) - result = await self.render_chat_request(request) - if isinstance(result, ErrorResponse): - return result - - conversation, engine_prompts = result - - # We override prompt tokens directly. - # VLM conversations use MITO (message-based) instead of TITO, so - # multi_modal_data is not expected here. - engine_prompts[0]["prompt_token_ids"] = request.tokens # type: ignore - - request_id = f"chatcmpl-{self._base_request_id(raw_request, request.request_id)}" - - request_metadata = RequestResponseMetadata(request_id=request_id) - if raw_request: - raw_request.state.request_metadata = request_metadata - - try: - lora_request = self._maybe_get_adapters(request, supports_default_mm_loras=True) - - model_name = self.models.model_name(lora_request) - except (ValueError, TypeError, RuntimeError) as e: - logger.exception("Error preparing request components") - return self.create_error_response(e) - - # Extract data_parallel_rank from header (router can inject it) - data_parallel_rank = self._get_data_parallel_rank(raw_request) - - # Schedule the request and get the result generator. - max_model_len = self.model_config.max_model_len - generators: list[AsyncGenerator[RequestOutput, None]] = [] - try: - for i, engine_prompt in enumerate(engine_prompts): - prompt_token_ids = self._extract_prompt_components(engine_prompt).token_ids - - # If we are creating sub requests for multiple prompts, ensure that they - # have unique request ids. - sub_request_id = request_id if len(engine_prompts) == 1 else f"{request_id}_{i}" - - prompt_len = self._extract_prompt_len(engine_prompt) - if prompt_len >= max_model_len: - raise VLLMValidationError( - f"This model's maximum context length is " - f"{max_model_len} tokens. However, your request has " - f"{prompt_len} input tokens. Please reduce the length of " - "the input messages.", - parameter="input_tokens", - value=prompt_len, - ) - - max_tokens = get_max_tokens( - max_model_len, - request.max_completion_tokens if request.max_completion_tokens is not None else request.max_tokens, - self._extract_prompt_len(engine_prompt), - self.default_sampling_params, - self.override_max_tokens, - ) - - sampling_params: SamplingParams | BeamSearchParams - if request.use_beam_search: - sampling_params = request.to_beam_search_params(max_tokens, self.default_sampling_params) - else: - sampling_params = request.to_sampling_params( - max_tokens, - self.default_sampling_params, - ) - - self._log_inputs( - sub_request_id, - engine_prompt, - params=sampling_params, - lora_request=lora_request, - ) - - trace_headers = None if raw_request is None else await self._get_trace_headers(raw_request.headers) - - if isinstance(sampling_params, BeamSearchParams): - generator = self.beam_search( - prompt=engine_prompt, - request_id=sub_request_id, - params=sampling_params, - lora_request=lora_request, - trace_headers=trace_headers, - ) - else: - reasoning_ended = ( - reasoning_parser.is_reasoning_end(prompt_token_ids or []) if reasoning_parser else None - ) - - generator = self.engine_client.generate( - engine_prompt, - sampling_params, - sub_request_id, - lora_request=lora_request, - trace_headers=trace_headers, - priority=request.priority, - data_parallel_rank=data_parallel_rank, - reasoning_ended=reasoning_ended, - ) - - generators.append(generator) - except ValueError as e: - return self.create_error_response(e) - - assert len(generators) == 1 - (result_generator,) = generators - - if request.stream: - return self.chat_completion_stream_generator( - request, - result_generator, - request_id, - model_name, - conversation, - tokenizer, - request_metadata, - reasoning_parser, - ) - - try: - return await self.chat_completion_full_generator( - request, - result_generator, - request_id, - model_name, - conversation, - tokenizer, - request_metadata, - reasoning_parser, - ) - except GenerationError: - raise # Let FastAPI's global generation_error_handler handle it - except ValueError as e: - return self.create_error_response(e) diff --git a/src/prime_rl/inference/vllm/serving_tokens.py b/src/prime_rl/inference/vllm/serving_tokens.py index 359df83d11..932ebeaa60 100644 --- a/src/prime_rl/inference/vllm/serving_tokens.py +++ b/src/prime_rl/inference/vllm/serving_tokens.py @@ -226,10 +226,8 @@ async def serve_tokens( # but never threads it into the engine, so PD disagg never fires on # ``/inference/v1/generate`` — decode receives an empty NIXL handshake # and ends up re-prefilling the prompt locally (~100× slower under - # concurrency). The chat path bridges this via - # ``ChatCompletionRequestWithTokens.to_sampling_params``; we mirror that - # bridge here so the engine's KV connector picks the params up out of - # ``sampling_params.extra_args``. + # concurrency). Bridge it through ``sampling_params.extra_args`` so the + # engine's KV connector picks the params up. # # Upstream fix: https://github.com/vllm-project/vllm/pull/42644 — drop # this block once we pin a vLLM version that includes it. @@ -300,12 +298,11 @@ async def serve_tokens_full_generator( # type: ignore[override] model_name: str, request_metadata: RequestResponseMetadata, ) -> ErrorResponse | GenerateResponse: - # Mirror serving_chat_with_tokens: wrap the result generator to capture - # routed_experts as it streams, defer the rest to upstream, then post- - # process the response into our PrimeRlGenerateResponse subclass so the - # encoded experts surface in the JSON. Skipping the wrapper when the - # engine isn't producing routed experts keeps us a no-op subclass on - # the common path. + # Wrap the result generator to capture routed_experts as it streams, + # defer the rest to upstream, then post-process the response into our + # PrimeRlGenerateResponse subclass so the encoded experts surface in + # the JSON. Skipping the wrapper when the engine isn't producing routed + # experts keeps us a no-op subclass on the common path. capture: _RoutedExpertsCapture | None = None if self.model_config.enable_return_routed_experts: capture = _RoutedExpertsCapture(result_generator) From 5587044ee97bd229cc2b7b3205d5588273d44f25 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Tue, 19 May 2026 00:45:48 +0000 Subject: [PATCH 26/47] refactor: require student_inference in all training modes Student inference was already required for rl/opd and optional for sft; making it required everywhere simplifies the orchestrator + scheduler: - Drop has_student_inference / enable_policy_updates gating - student_inference: InferencePool (no longer Optional) in Scheduler.__init__ - Weight broadcast init, weight sync, metrics collector, eval routing, and shutdown all unconditionally use student_inference - Drop the now-stale OPD probe (will revisit in a follow-up PR) - Add is_vlm property on BaseModelConfig; replace local is_vlm vars in the orchestrator with config.student.model.is_vlm - Merge the two legacy student-layout shims on OrchestratorConfig into one before-validator (_accept_legacy_student_layout) covering both [orchestrator.client.*] and flat [orchestrator.model.] paths - Drop redundant pool / WandbMonitor / PrimeMonitor init log lines - Reorder OrchestratorConfig fields: training_mode > student > teacher - Drop "OAI" from RolloutModelConfig.client description - Update docs/training_modes.md: SFT student inference is now required; evals and weight sync are unconditional Co-Authored-By: Claude Sonnet 4.6 --- docs/training_modes.md | 12 +- .../src/prime_rl/configs/orchestrator.py | 74 ++++++----- .../src/prime_rl/configs/shared.py | 4 + src/prime_rl/orchestrator/orchestrator.py | 116 +++++++----------- src/prime_rl/orchestrator/scheduler.py | 44 +++---- src/prime_rl/orchestrator/utils.py | 53 -------- tests/unit/orchestrator/test_scheduler.py | 1 - tests/unit/test_configs.py | 23 ---- 8 files changed, 113 insertions(+), 214 deletions(-) diff --git a/docs/training_modes.md b/docs/training_modes.md index d6c49285f1..596211125c 100644 --- a/docs/training_modes.md +++ b/docs/training_modes.md @@ -12,16 +12,16 @@ The mode determines who generates rollouts, what role the teacher plays, and wha | | **rl** | **opd** | **sft** | |---|---|---|---| -| **Student does** | generate rollouts → get trained on them | generate rollouts → get trained on them | get trained on teacher's rollouts; optionally serve inference for evals | +| **Student does** | generate rollouts → get trained on them | generate rollouts → get trained on them | serve inference (for evals + weight sync); get trained on teacher's rollouts | | **Teacher does** | nothing (must be unset) | score student rollouts (token-level logprobs) | generate rollouts | | **Loss** | reward-based (advantage) | reward + KL to teacher logprobs (`teacher_tau > 0`) | pure NLL on teacher tokens (hard distill) | -| **Student inference** (`[inference]`) | **required** | **required** | **optional** — only if you want evals or weight-sync the student | +| **Student inference** (`[inference]`) | **required** | **required** | **required** | | **Teacher inference** (`[teacher_inference]`) | forbidden | **required, must be vLLM** | not used (teacher is external) | | **`[orchestrator.teacher]`** | must be `None` | auto-wired from `[teacher_inference]` | **required** — `client.base_url` + `model.name` of external endpoint | | **`num_teacher_gpus`** | unset | **required** (`> 0`) | unset (teacher is external) | | **Teacher endpoint type** | n/a | **local vLLM only** | **any OpenAI-compatible** (PI inference, OpenAI, Anthropic, local vLLM…) | -| **Weight sync (trainer → ?)** | → student inference | → student inference (teacher frozen) | → student inference if configured; teacher never touched | -| **Evals** | student | student | only if `[inference]` is set (then student evals) | +| **Weight sync (trainer → ?)** | → student inference | → student inference (teacher frozen) | → student inference (teacher frozen) | +| **Evals** | student | student | student | ## Key implications @@ -29,8 +29,6 @@ The mode determines who generates rollouts, what role the teacher plays, and wha **SFT's teacher is just chat completions.** It only needs `/v1/chat/completions`. Point `[orchestrator.teacher.client]` at anything OpenAI-compatible. No local GPU needed for the teacher. -**SFT student inference is optional but enabling it changes behavior.** If you set `[inference]`, you get (a) student-side evals during training, and (b) weight sync from trainer to student inference (so the student inference pool reflects training progress). Without `[inference]`, the run is teacher-rollout-only with no online evals. - **RL forbids any teacher.** Even a stray `[orchestrator.teacher]` block fails validation. **Student model name is always the model being trained.** In SFT this is *not* the rollout-generating model — that's the teacher. The student model field still determines tokenizer, trainer init weights, and what gets saved as checkpoints. @@ -62,7 +60,7 @@ training_mode = "sft" base_url = ["https://api.pinference.ai/api/v1"] [orchestrator.teacher.model] name = "qwen/qwen3-30b-a3b-instruct-2507" -[inference] # optional — drop if you don't want student-side evals +[inference] ``` ## OPD details 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 7fcb141314..787eb66320 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py @@ -890,28 +890,9 @@ class RolloutModelConfig(BaseConfig): client: Annotated[ ClientConfig, - Field(description="The OAI client configuration."), + Field(description="The client configuration."), ] = ClientConfig() - @model_validator(mode="before") - @classmethod - def _accept_flat_model_layout(cls, data): - """Accept legacy flat ModelConfig layout (e.g. [orchestrator.model.lora]). - - Pre-refactor, orchestrator.model was a ModelConfig directly (name, lora, - trust_remote_code, vlm). Now it's a RolloutModelConfig wrapping ModelConfig - + ClientConfig. Detect dicts whose only keys are ModelConfig fields and - re-nest them under "model" so existing configs keep working. - """ - if not isinstance(data, dict): - return data - model_only_keys = {"name", "trust_remote_code", "vlm", "lora"} - if "model" in data or "client" in data: - return data - if any(k in model_only_keys for k in data.keys()): - return {"model": data} - return data - class OrchestratorConfig(BaseConfig): """Configures the orchestrator for RL training.""" @@ -1155,19 +1136,52 @@ class OrchestratorConfig(BaseConfig): @model_validator(mode="before") @classmethod - def _accept_top_level_client(cls, data: Any) -> Any: - """Accept legacy [orchestrator.client] as shorthand for [orchestrator.student.client]. + def _accept_legacy_student_layout(cls, data: Any) -> Any: + """Backward-compat shims for the pre-refactor student layout. + + Pre-refactor OrchestratorConfig had top-level `model: ModelConfig` and + `client: ClientConfig` fields. The student/teacher rename consolidated + both under `student: RolloutModelConfig` (with `model` as a legacy alias + for `student`). Re-nest legacy keys so old configs still parse: + + - [orchestrator.client.*] -> [orchestrator.student.client.*] + - [orchestrator.model.] -> [orchestrator.student.model.] + (where is a ModelConfig field: name, trust_remote_code, vlm, lora) - Pre-refactor, OrchestratorConfig had a top-level `client` field. After the - student/teacher rename it moved under `student.client`. Re-nest a top-level - `client` dict so existing configs keep working. + Teacher was always nested pre-refactor (teacher_model.model + + teacher_model.client), so we don't touch it. """ - if not isinstance(data, dict) or "client" not in data: + if not isinstance(data, dict): return data - student = data.setdefault("student", {}) - # If the user already nested model+client under student/model, leave it alone. - if isinstance(student, dict) and "client" not in student: - student["client"] = data.pop("client") + + # 1. Re-nest top-level [orchestrator.client] under student.client. + if "client" in data: + student = data.setdefault("student", {}) + if isinstance(student, dict) and "client" not in student: + student["client"] = data.pop("client") + + # 2. Consolidate the legacy `model` alias into `student` so the + # flat-layout fix-up below sees a single target. + legacy_model = data.pop("model", None) + if legacy_model is not None: + existing = data.get("student") + if existing is None: + data["student"] = legacy_model + elif isinstance(existing, dict) and isinstance(legacy_model, dict): + for k, v in legacy_model.items(): + existing.setdefault(k, v) + else: + # Mismatched types - put it back and let pydantic surface the error. + data["model"] = legacy_model + + # 3. Re-nest flat ModelConfig keys under student.model. + model_only_keys = {"name", "trust_remote_code", "vlm", "lora"} + student = data.get("student") + if isinstance(student, dict): + flat = {k: student.pop(k) for k in list(student) if k in model_only_keys} + if flat: + student.setdefault("model", {}).update(flat) + return data @model_validator(mode="before") diff --git a/packages/prime-rl-configs/src/prime_rl/configs/shared.py b/packages/prime-rl-configs/src/prime_rl/configs/shared.py index 82cc784408..7af041e540 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/shared.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/shared.py @@ -131,6 +131,10 @@ class BaseModelConfig(BaseConfig): ), ] = None + @property + def is_vlm(self) -> bool: + return self.vlm is not None + class RendererConfig(BaseConfig): """Configures the client-side renderer (chat-template + response parsing). diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index 6bde5e6b66..0710fd25a8 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -46,7 +46,6 @@ compute_teacher_logprobs, get_weight_dir, print_benchmark, - probe_teacher_logprobs, set_default_executor, ) from prime_rl.orchestrator.vf_utils import ( @@ -125,37 +124,27 @@ async def orchestrate(config: OrchestratorConfig): for env_id in env_ids_to_install: install_env(env_id, prerelease=config.env_install_prerelease) - # Check if this is a vision-language model (used throughout for VLM-specific paths) - is_vlm = config.student.model.vlm is not None - # Load tokenizer and processor (processor only for VLM models) logger.info(f"Initializing tokenizer ({config.tokenizer})") tokenizer = setup_tokenizer(config.tokenizer) processor = None - if is_vlm: + if config.student.model.is_vlm: logger.info(f"Loading VLM processor for {config.student.model.name}") processor = AutoProcessor.from_pretrained( config.student.model.name, trust_remote_code=config.student.model.trust_remote_code, use_fast=True ) - # Set up student inference pool. Required for rl/opd; optional for sft (only - # configured when the user wrote [inference] - signal: student.client.base_url - # is in model_fields_set, set by auto_setup_inference_client). When absent, - # SFT runs in teacher-only mode: no online evals, no weight sync. - has_student_inference = config.training_mode != "sft" or "base_url" in config.student.client.model_fields_set - student_inference = None - renderer = None - if has_student_inference: - logger.info( - f"Initializing student inference pool (base_url={', '.join(config.student.client.base_url)}, " - f"model={config.student.model.name})" - ) - renderer, student_inference = await setup_student_inference_pool( - config=config, - tokenizer=tokenizer, - logger=logger, - ) + # Set up student inference pool (required for all training modes). + logger.info( + f"Initializing student inference pool (base_url={', '.join(config.student.client.base_url)}, " + f"model={config.student.model.name})" + ) + renderer, student_inference = await setup_student_inference_pool( + config=config, + tokenizer=tokenizer, + logger=logger, + ) # Set up teacher inference pool (configured for opd or sft). Always MITO for # simplicity - this also keeps external OAI-compatible teachers (PI inference, @@ -172,9 +161,6 @@ async def orchestrate(config: OrchestratorConfig): train_client_type="openai_chat_completions", ) - # Weight sync is only possible when a student inference pool exists. - enable_policy_updates = student_inference is not None - # Setup monitor (may register the run and set RUN_ID in the environment) logger.info(f"Initializing monitor (wandb={config.wandb}, prime_monitor={config.prime_monitor})") monitor = setup_monitor( @@ -261,47 +247,37 @@ async def orchestrate(config: OrchestratorConfig): max_off_policy_steps=config.max_off_policy_steps, strict_async_level=config.strict_async_level, tasks_per_minute=config.tasks_per_minute, - enable_policy_updates=enable_policy_updates, lora_name=config.student.model.lora.name if config.student.model.lora else None, config=config, ) # Wait for pools to be ready - if student_inference is not None: - logger.info("Waiting for student inference pool to be ready") - await student_inference.wait_for_ready(config.student.model.name) - logger.success("Student inference pool ready") + logger.info("Waiting for student inference pool to be ready") + await student_inference.wait_for_ready(config.student.model.name) + logger.success("Student inference pool ready") if teacher_inference is not None: assert config.teacher is not None logger.info("Waiting for teacher inference pool to be ready") await teacher_inference.wait_for_ready(config.teacher.model.name) logger.success("Teacher inference pool ready") - if config.training_mode == "opd": - logger.info("Probing teacher for prompt_logprobs support") - await probe_teacher_logprobs(teacher_inference.train_clients, config.teacher.model.name) - logger.success("Teacher supports prompt_logprobs") - # Start inference metrics collector (requires W&B + student inference pool) + # Start inference metrics collector (requires W&B) inference_metrics_collector = None - if config.wandb is not None and config.collect_inference_metrics and student_inference is not None: + if config.wandb is not None and config.collect_inference_metrics: inference_metrics_collector = InferenceMetricsCollector(student_inference.admin_clients) await inference_metrics_collector.start() # Set up weight broadcast backend (targets student inference) - if enable_policy_updates: - assert student_inference is not None - logger.info(f"Initializing weight broadcast ({config.weight_broadcast})") - if config.weight_broadcast.type == "nccl": - await init_nccl_broadcast( - student_inference.admin_clients, - config.weight_broadcast.host, - config.weight_broadcast.port, - config.weight_broadcast.timeout, - inference_world_size=config.weight_broadcast.inference_world_size, - quantize_in_weight_transfer=config.weight_broadcast.quantize_in_weight_transfer, - ) - else: - logger.info("Skipping weight broadcast initialization (no student inference pool)") + logger.info(f"Initializing weight broadcast ({config.weight_broadcast})") + if config.weight_broadcast.type == "nccl": + await init_nccl_broadcast( + student_inference.admin_clients, + config.weight_broadcast.host, + config.weight_broadcast.port, + config.weight_broadcast.timeout, + inference_world_size=config.weight_broadcast.inference_world_size, + quantize_in_weight_transfer=config.weight_broadcast.quantize_in_weight_transfer, + ) # Setup training batch sender for sending training examples to trainer logger.info(f"Initializing training batch sender ({config.rollout_transport})") @@ -327,20 +303,18 @@ async def orchestrate(config: OrchestratorConfig): # Allow eval at resumed step by setting prev_ckpt_step one behind prev_ckpt_step = scheduler.ckpt_step - 1 - if enable_policy_updates: - assert student_inference is not None - # In NCCL mode, skip existence check - weights are broadcasted, not stored on disk - check_exists = config.weight_broadcast.type != "nccl" - wait_timeout = config.ckpt.wait_for_weights_timeout if config.ckpt else None - weights_path = get_weight_dir( - config.output_dir, scheduler.ckpt_step, check_exists=check_exists, wait_timeout=wait_timeout - ) - lora_name = config.student.model.lora.name if config.student.model.lora else None - await student_inference.update_weights(weights_path, lora_name=lora_name, step=scheduler.ckpt_step) - if lora_name is not None: - student_inference.update_model_name(lora_name) - if scheduler.rollout_inference is student_inference: - scheduler.model_name = lora_name + # In NCCL mode, skip existence check - weights are broadcasted, not stored on disk + check_exists = config.weight_broadcast.type != "nccl" + wait_timeout = config.ckpt.wait_for_weights_timeout if config.ckpt else None + weights_path = get_weight_dir( + config.output_dir, scheduler.ckpt_step, check_exists=check_exists, wait_timeout=wait_timeout + ) + lora_name = config.student.model.lora.name if config.student.model.lora else None + await student_inference.update_weights(weights_path, lora_name=lora_name, step=scheduler.ckpt_step) + if lora_name is not None: + student_inference.update_model_name(lora_name) + if scheduler.rollout_inference is student_inference: + scheduler.model_name = lora_name else: logger.info("Training from scratch") @@ -356,7 +330,7 @@ async def orchestrate(config: OrchestratorConfig): raise RuntimeError(f"Run evicted by trainer: {reason}") # Capture ckpt_step once for consistency (it's updated inside the scheduler) - ckpt_step = scheduler.ckpt_step if enable_policy_updates else progress.step + ckpt_step = scheduler.ckpt_step scheduler.ckpt_step = ckpt_step # Save checkpoint (if we are at an interval step and not at the first or last step) @@ -384,7 +358,7 @@ async def orchestrate(config: OrchestratorConfig): # scheduler.checkpoint_ready during eval to ensure consistent weights. # Each eval env has its own interval, so we check each independently. envs_to_eval: list[EvalEnv] = [] - if config.eval and student_inference is not None: + if config.eval: assert eval_envs is not None for eval_env in eval_envs: eval_ckpt_step = compute_eval_ckpt_step( @@ -399,7 +373,6 @@ async def orchestrate(config: OrchestratorConfig): envs_to_eval.append(eval_env) if envs_to_eval: - assert student_inference is not None env_names = ", ".join(e.name for e in envs_to_eval) logger.info(f"Running evals at {ckpt_step=} for {env_names}") @@ -496,7 +469,7 @@ async def orchestrate(config: OrchestratorConfig): ) # VLM: offload base64 images to disk immediately to free memory - if is_vlm: + if config.student.model.is_vlm: offload_start = time.perf_counter() num_offloaded = offload_images_to_disk(train_rollouts, config.output_dir) if num_offloaded: @@ -528,7 +501,7 @@ async def _pretokenize_all() -> None: ) ) - if is_vlm: + if config.student.model.is_vlm: mm_token_type_ids_mapping = {} if hasattr(processor, "image_token_id") and processor.image_token_id is not None: mm_token_type_ids_mapping[processor.image_token_id] = 1 @@ -837,7 +810,7 @@ def compute_solve_rates(df): if heart is not None: heart.beat() - if config.eval and eval_envs is not None and student_inference is not None: + if config.eval and eval_envs is not None: logger.info("Running final evals") eval_results = await asyncio.gather( *[ @@ -879,8 +852,7 @@ async def _graceful_shutdown() -> None: await scheduler.stop() if inference_metrics_collector is not None: await inference_metrics_collector.stop() - if student_inference is not None: - await student_inference.stop() + await student_inference.stop() if teacher_inference is not None: await teacher_inference.stop() event_loop_lag_monitor_task.cancel() diff --git a/src/prime_rl/orchestrator/scheduler.py b/src/prime_rl/orchestrator/scheduler.py index fbc4b79efc..53964fa631 100644 --- a/src/prime_rl/orchestrator/scheduler.py +++ b/src/prime_rl/orchestrator/scheduler.py @@ -73,7 +73,7 @@ class Scheduler: def __init__( self, train_envs: TrainEnvs, - student_inference: InferencePool | None, + student_inference: InferencePool, teacher_inference: InferencePool | None, buffer: Buffer, config: OrchestratorConfig, @@ -82,7 +82,6 @@ def __init__( max_off_policy_steps: int, strict_async_level: bool, tasks_per_minute: int | None, - enable_policy_updates: bool = True, lora_name: str | None = None, ): self.logger = get_logger() @@ -100,18 +99,19 @@ def __init__( self.max_async_level = max_async_level self.max_off_policy_steps = max_off_policy_steps self.strict_async_level = strict_async_level - self.enable_policy_updates = enable_policy_updates self.lora_name = lora_name self.json_logging = config.log.json_logging - # student_inference is the weight-sync target (None = no policy updates). - # teacher_inference is set in opd (for logprobs) and sft (for rollouts). - # rollout_inference is whichever pool serves train rollouts for this mode. + # student_inference is the weight-sync target. teacher_inference is set + # in opd (for logprobs) and sft (for rollouts). rollout_inference is + # whichever pool serves train rollouts for this mode. self.student_inference = student_inference self.teacher_inference = teacher_inference - rollout = teacher_inference if config.training_mode == "sft" else student_inference - assert rollout is not None, "rollout_inference resolved to None - config validation should prevent this" - self.rollout_inference: InferencePool = rollout + if config.training_mode == "sft": + assert teacher_inference is not None + self.rollout_inference: InferencePool = teacher_inference + else: + self.rollout_inference = student_inference # model_name is the name to send on rollout requests - matches the rollout pool self.model_name = self.rollout_inference.model_name @@ -328,9 +328,6 @@ async def _apply_policy_update(self, next_ckpt_step: int) -> None: update_weights_start_time = time.perf_counter() weights_path = get_step_path(get_broadcast_dir(self.config.output_dir), next_ckpt_step) - assert self.student_inference is not None, ( - "weight sync requires student_inference - guard with enable_policy_updates" - ) await self.student_inference.update_weights(weights_path, lora_name=self.lora_name, step=next_ckpt_step) self.update_weights_time = time.perf_counter() - update_weights_start_time self.logger.debug(f"Updated weights to step {next_ckpt_step} in {self.update_weights_time:.2f}s") @@ -365,11 +362,6 @@ def _clear_inflight_policy_update(done_task: asyncio.Task) -> None: async def maybe_update_policy(self): """Updates the policy to the latest available checkpoint. Aborts rollout requests that are older than the max retention steps.""" - if not self.enable_policy_updates: - self.ckpt_step = self.step - self.checkpoint_ready.set() - return - while True: next_ckpt_step = self._compute_next_ckpt_step() if next_ckpt_step <= self.ckpt_step: @@ -409,18 +401,14 @@ async def generate_batch(self, step: int) -> list[vf.RolloutOutput]: """Continuously generates a batch of rollouts.""" self.step = step - if self.enable_policy_updates: - # Cancel the previous update policy task to avoid concurrent updates - if self.update_policy_task is not None: - await safe_cancel(self.update_policy_task) + # Cancel the previous update policy task to avoid concurrent updates + if self.update_policy_task is not None: + await safe_cancel(self.update_policy_task) - # Manually check the async barrier before starting the step, then re-create the update policy loop - # This ensures that we respect max_async_level, while still listening for policy updates mid-step - await self.maybe_update_policy() - self.update_policy_task = asyncio.create_task(self.update_policy_loop()) - else: - self.ckpt_step = step - self.checkpoint_ready.set() + # Manually check the async barrier before starting the step, then re-create the update policy loop + # This ensures that we respect max_async_level, while still listening for policy updates mid-step + await self.maybe_update_policy() + self.update_policy_task = asyncio.create_task(self.update_policy_loop()) batch_start_time = time.perf_counter() diff --git a/src/prime_rl/orchestrator/utils.py b/src/prime_rl/orchestrator/utils.py index 9458ea64a5..121cc57c54 100644 --- a/src/prime_rl/orchestrator/utils.py +++ b/src/prime_rl/orchestrator/utils.py @@ -133,59 +133,6 @@ async def _compute_single(client_config: vf.ClientConfig, sample: TrainingSample return await asyncio.gather(*[_compute_single(client, sample) for client, sample in zip(cycle(clients), samples)]) -async def probe_teacher_logprobs(clients: list[vf.ClientConfig], model_name: str) -> None: - """Probe the teacher endpoint for vLLM `prompt_logprobs` support. - - OPD requires the vLLM-specific ``/inference/v1/generate`` endpoint with - ``prompt_logprobs=1`` (see ``compute_teacher_logprobs``). External - OAI-compatible endpoints (PI inference, OpenAI, Anthropic) don't expose - this and would 404 mid-training. Probe once at startup with a tiny - request and raise a clear error if the endpoint can't serve it. - """ - import httpx - from vllm.entrypoints.serve.disagg.protocol import GenerateResponse - - client_config = clients[0] - client = setup_openai_client(client_config) - base = str(client.base_url).rstrip("/").removesuffix("/v1") - url = f"{base}/inference/v1/generate" - - hint = ( - "OPD requires a self-hosted vLLM teacher exposing /inference/v1/generate with " - "prompt_logprobs. External OAI-compatible endpoints (PI inference, OpenAI, Anthropic) " - "are not supported. See docs/training_modes.md." - ) - - try: - http_response = await client.post( - url, - cast_to=httpx.Response, - body={ - "model": model_name, - "token_ids": [1, 2, 3], - "sampling_params": { - "max_tokens": 1, - "temperature": 1.0, - "top_p": 1.0, - "prompt_logprobs": 1, - }, - }, - ) - except Exception as e: - raise RuntimeError(f"OPD teacher probe failed: {url} - {type(e).__name__}: {e}\n{hint}") from e - - try: - response = GenerateResponse.model_validate_json(http_response.content) - except Exception as e: - snippet = http_response.content[:200] if hasattr(http_response, "content") else "" - raise RuntimeError( - f"OPD teacher probe failed: {url} returned a non-vLLM response (got {snippet!r}).\n{hint}" - ) from e - - if not response.prompt_logprobs: - raise RuntimeError(f"OPD teacher probe failed: {url} returned no prompt_logprobs in the response.\n{hint}") - - def get_weight_dir(output_dir: Path, step: int, check_exists: bool = True, wait_timeout: int | None = None) -> Path: """Get the weight directory for a given checkpoint step. diff --git a/tests/unit/orchestrator/test_scheduler.py b/tests/unit/orchestrator/test_scheduler.py index 1d49482362..fd908c434d 100644 --- a/tests/unit/orchestrator/test_scheduler.py +++ b/tests/unit/orchestrator/test_scheduler.py @@ -30,7 +30,6 @@ def make_scheduler() -> Scheduler: scheduler.policy_update_lock = asyncio.Lock() scheduler.inflight_policy_update_task = None scheduler.update_policy_task = None - scheduler.enable_policy_updates = True scheduler.rate_limiter = None return scheduler diff --git a/tests/unit/test_configs.py b/tests/unit/test_configs.py index 731d00c5d1..bdfefbf3ad 100644 --- a/tests/unit/test_configs.py +++ b/tests/unit/test_configs.py @@ -194,26 +194,3 @@ def test_orchestrator_vlm_configs_must_disable_renderer(): def test_selective_activation_checkpointing_requires_custom_impl(): with pytest.raises(ValidationError, match="Selective activation checkpointing requires model.impl='custom'"): TrainerModelConfig.model_validate({"impl": "hf", "ac": {"mode": "selective"}}) - - -def test_sft_training_mode_enables_student_pool_when_inference_configured(): - base_config = { - "training_mode": "sft", - "trainer": {}, - "orchestrator": { - "use_token_client": False, - "use_renderer": False, - "teacher": { - "client": {"base_url": ["http://teacher.example/v1"]}, - "model": {"name": "teacher-model"}, - }, - }, - } - - # Without inference, student client base_url not explicitly set → policy updates disabled - config = RLConfig.model_validate(base_config) - assert "base_url" not in config.orchestrator.student.client.model_fields_set - - # With inference, student client base_url is auto-set → policy updates enabled - config = RLConfig.model_validate({**base_config, "inference": {}}) - assert "base_url" in config.orchestrator.student.client.model_fields_set From 65157c94e5f6c4974dd1e716e7ad9965496b9b28 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Tue, 19 May 2026 01:19:29 +0000 Subject: [PATCH 27/47] feat(configs): default OPD to pure distillation (teacher_tau=1, adv_tau=0) auto_setup_training_mode now seeds trainer.loss defaults when training_mode is opd: teacher_tau=1.0 and adv_tau=0.0 (pure distillation: drop the reward signal, full weight on the teacher KL). User-provided values still win via setdefault. Also accept training_mode set only inside [orchestrator] - same auto-setup applies for both the shared top-level and orchestrator-local forms. Drop the redundant [trainer.loss] teacher_tau = 0.5 block from configs/reverse_text/debug_opd.toml; verified the auto-setup applies end-to-end (3 train steps, teacher logprobs computed each step, clean exit). Also drop the now-redundant InferencePool type annotation on self.rollout_inference - both branches assign a definite InferencePool now that student_inference is non-Optional. Co-Authored-By: Claude Sonnet 4.6 --- configs/reverse_text/debug_opd.toml | 3 --- .../src/prime_rl/configs/rl.py | 23 ++++++++++++++----- src/prime_rl/orchestrator/scheduler.py | 2 +- 3 files changed, 18 insertions(+), 10 deletions(-) diff --git a/configs/reverse_text/debug_opd.toml b/configs/reverse_text/debug_opd.toml index b3d3df41c3..35c41182c4 100644 --- a/configs/reverse_text/debug_opd.toml +++ b/configs/reverse_text/debug_opd.toml @@ -45,9 +45,6 @@ base_url = ["http://localhost:8001/v1"] [trainer.optim] lr = 3e-6 -[trainer.loss] -teacher_tau = 0.5 - [ckpt] [inference] 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 f6716dcc08..815a8606eb 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/rl.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/rl.py @@ -434,15 +434,22 @@ def validate_quantize_in_weight_transfer(self): @model_validator(mode="before") @classmethod def auto_setup_training_mode(cls, data): - """Propagate shared training_mode into orchestrator.training_mode and trainer.loss.type. + """Propagate shared training_mode into orchestrator.training_mode and trainer.loss. Runs before nested validation so that OrchestratorConfig.validate_training_mode - sees the propagated value. Only propagates to components that don't already set - the field explicitly. For 'sft' mode, defaults trainer.loss.type to 'sft'. + sees the propagated value. Only propagates to fields the user didn't set: + + - sft: trainer.loss.type = "sft" + - opd: trainer.loss.teacher_tau = 1.0, trainer.loss.adv_tau = 0.0 (pure distillation) """ if not isinstance(data, dict): return data mode = data.get("training_mode") + if mode is None: + # Also accept training_mode set only inside [orchestrator] + orch_candidate = data.get("orchestrator") + if isinstance(orch_candidate, dict): + mode = orch_candidate.get("training_mode") if mode is None: return data @@ -450,12 +457,16 @@ def auto_setup_training_mode(cls, data): if isinstance(orch, dict) and "training_mode" not in orch: orch["training_mode"] = mode - if mode == "sft": + if mode in ("sft", "opd"): trainer = data.setdefault("trainer", {}) if isinstance(trainer, dict): loss = trainer.setdefault("loss", {}) - if isinstance(loss, dict) and "type" not in loss: - loss["type"] = "sft" + if isinstance(loss, dict): + if mode == "sft": + loss.setdefault("type", "sft") + elif mode == "opd": + loss.setdefault("teacher_tau", 1.0) + loss.setdefault("adv_tau", 0.0) return data diff --git a/src/prime_rl/orchestrator/scheduler.py b/src/prime_rl/orchestrator/scheduler.py index 53964fa631..02840b6443 100644 --- a/src/prime_rl/orchestrator/scheduler.py +++ b/src/prime_rl/orchestrator/scheduler.py @@ -109,7 +109,7 @@ def __init__( self.teacher_inference = teacher_inference if config.training_mode == "sft": assert teacher_inference is not None - self.rollout_inference: InferencePool = teacher_inference + self.rollout_inference = teacher_inference else: self.rollout_inference = student_inference # model_name is the name to send on rollout requests - matches the rollout pool From 69d1f68973f6dbf4ed66e3542b6163c98f99dc56 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Tue, 19 May 2026 01:37:53 +0000 Subject: [PATCH 28/47] chore(configs): rename debug_sft_thinking -> debug_sft_external; drop examples/alphabet_sort/sft_distill_hard.toml debug_sft_external describes the role better - external SFT teacher via OAI-compatible endpoint (PI inference). The "thinking" label was about the specific model the config picks (qwen3-30b-a3b-thinking-2507), not about what the config demonstrates. Drop examples/alphabet_sort/sft_distill_hard.toml: configs/reverse_text/ debug_sft* are the canonical SFT examples now. Co-Authored-By: Claude Sonnet 4.6 --- configs/reverse_text/README.md | 4 +- ..._thinking.toml => debug_sft_external.toml} | 2 +- docs/training_modes.md | 4 +- examples/alphabet_sort/sft_distill_hard.toml | 49 ------------------- 4 files changed, 5 insertions(+), 54 deletions(-) rename configs/reverse_text/{debug_sft_thinking.toml => debug_sft_external.toml} (93%) delete mode 100644 examples/alphabet_sort/sft_distill_hard.toml diff --git a/configs/reverse_text/README.md b/configs/reverse_text/README.md index 9402778a19..d544e54aba 100644 --- a/configs/reverse_text/README.md +++ b/configs/reverse_text/README.md @@ -7,7 +7,7 @@ Minimal end-to-end configs for the three training modes against the `reverse-tex | `debug_rl.toml` | `rl` | none | | `debug_opd.toml` | `opd` | local vLLM (`Qwen3-0.6B-Reverse-Text-RL`) | | `debug_sft.toml` | `sft` | local vLLM (`Qwen3-0.6B-Reverse-Text-RL`) | -| `debug_sft_thinking.toml` | `sft` | PI inference (`qwen/qwen3-30b-a3b-thinking-2507`) | +| `debug_sft_external.toml` | `sft` | PI inference (`qwen/qwen3-30b-a3b-thinking-2507`) | The student inference server is auto-launched on GPU 0 at `http://localhost:8000/v1` with `gpu_memory_utilization=0.5`. The teacher (used by `debug_opd.toml` and `debug_sft.toml`) is **not** auto-launched — start it manually on GPU 1. @@ -35,7 +35,7 @@ uv run rl @ configs/reverse_text/debug_sft.toml # SFT hard distill from qwen3-30b-a3b-thinking via PI inference # (requires PRIME_API_KEY + PRIME_TEAM_ID in env; no local teacher needed) -uv run rl @ configs/reverse_text/debug_sft_thinking.toml +uv run rl @ configs/reverse_text/debug_sft_external.toml ``` See [docs/training_modes.md](../../docs/training_modes.md) for what each mode does. diff --git a/configs/reverse_text/debug_sft_thinking.toml b/configs/reverse_text/debug_sft_external.toml similarity index 93% rename from configs/reverse_text/debug_sft_thinking.toml rename to configs/reverse_text/debug_sft_external.toml index 514615fba4..c4baff223a 100644 --- a/configs/reverse_text/debug_sft_thinking.toml +++ b/configs/reverse_text/debug_sft_external.toml @@ -2,7 +2,7 @@ # X-Prime-Team-ID header is auto-injected from $PRIME_TEAM_ID for pinference.ai URLs. # # Run with: -# uv run rl @ configs/reverse_text/debug_sft_thinking.toml +# uv run rl @ configs/reverse_text/debug_sft_external.toml max_steps = 20 seq_len = 2048 diff --git a/docs/training_modes.md b/docs/training_modes.md index 596211125c..2183efd961 100644 --- a/docs/training_modes.md +++ b/docs/training_modes.md @@ -120,8 +120,8 @@ Notes: ### Reference configs -- `configs/alphabet_sort/sft_distill_hard_qwen4b_lora_prime_teacher.toml` -- `examples/alphabet_sort/sft_distill_hard.toml` +- `configs/reverse_text/debug_sft.toml` (local vLLM teacher) +- `configs/reverse_text/debug_sft_external.toml` (PI inference teacher) ## Parameter reference diff --git a/examples/alphabet_sort/sft_distill_hard.toml b/examples/alphabet_sort/sft_distill_hard.toml deleted file mode 100644 index 7ac6f51bff..0000000000 --- a/examples/alphabet_sort/sft_distill_hard.toml +++ /dev/null @@ -1,49 +0,0 @@ -max_steps = 24 -seq_len = 2048 -training_mode = "sft" - -[deployment] -type = "single_node" -gpus_per_node = 2 -num_train_gpus = 2 -num_infer_gpus = 0 - -[model] -name = "Qwen/Qwen3-4B-Instruct-2507" - -[wandb] -project = "alphabet-sort-sft-distill" -name = "qwen3-4b-hard-distill" - -[ckpt] - -[trainer.optim] -lr = 1e-5 - -[trainer.model.lora] -rank = 16 -alpha = 32 - -[trainer.ckpt.weights] -save_adapter_separately = true - -[orchestrator] -batch_size = 256 -rollouts_per_example = 4 -use_renderer = false - -[orchestrator.train.sampling] -max_completion_tokens = 512 -temperature = 0.7 - -[orchestrator.teacher.client] -base_url = ["https://api.pinference.ai/api/v1"] -api_key_var = "PRIME_API_KEY" - -[orchestrator.teacher.model] -name = "qwen/qwen3-235b-a22b-instruct-2507" - -[[orchestrator.train.env]] -id = "primeintellect/alphabet-sort" -name = "alphabet-sort" -args = { min_turns = 2, max_turns = 2, min_names_per_turn = 1, max_names_per_turn = 3, similarity_power = 4, power_per_turn = false } From 3b402f30b69bff66c7daf09a6c836832a60caa2e Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Tue, 19 May 2026 01:47:36 +0000 Subject: [PATCH 29/47] fix(tests): unbreak CPU tests after student/teacher rename + headers_from_env - tests/unit/train/test_runs.py: orchestrator config attribute access changed from config.model.name to config.student.model.name after the RolloutModelConfig rename. - tests/unit/utils/test_elastic.py: MagicMock client_config needs headers_from_env set to {} - the field was added in this PR and the test's hand-built mock missed it (pydantic rejects MagicMock as a dict). Co-Authored-By: Claude Sonnet 4.6 --- tests/unit/train/test_runs.py | 2 +- tests/unit/utils/test_elastic.py | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/unit/train/test_runs.py b/tests/unit/train/test_runs.py index a8a2ed17a1..b80da9c91e 100644 --- a/tests/unit/train/test_runs.py +++ b/tests/unit/train/test_runs.py @@ -217,7 +217,7 @@ def test_config_loading(tmp_path: Path) -> None: # Access config as OrchestratorConfig object config = multi_run_manager.config[run_idx] - assert config.model.name == "test-model" + assert config.student.model.name == "test-model" assert config.batch_size == 32 assert config.max_steps == 1000 diff --git a/tests/unit/utils/test_elastic.py b/tests/unit/utils/test_elastic.py index a2cd40d1d1..7a68c7413a 100644 --- a/tests/unit/utils/test_elastic.py +++ b/tests/unit/utils/test_elastic.py @@ -415,6 +415,7 @@ def test_elastic_clients_preserve_renderer_model_name_when_model_name_updates(): client_config.connect_timeout = 30.0 client_config.api_key_var = "PRIME_API_KEY" client_config.headers = {} + client_config.headers_from_env = {} client_config.extra_headers_from_state = {} client_config.dp_rank_count = 1 From 2cdcc7684b122b5f6fa8725e621485cf305c6067 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Tue, 19 May 2026 02:03:40 +0000 Subject: [PATCH 30/47] feat(orchestrator): strip logprobs from train sampling in sft mode The trainer in sft mode reconstructs teacher tokens via the student tokenizer - it never reads inference-side logprobs from the teacher endpoint. Stripping logprobs at runtime lets external reasoning-model endpoints (openai/gpt-5*, etc.) accept the rollout request - they reject logprobs with 400 invalid_request "logprobs are not supported with reasoning models". Done at orchestrator-startup time on each TrainEnv's sampling_args rather than as a public config field, since this is a mode-specific runtime concern. Also: bump configs/reverse_text/debug_sft_external.toml to use openai/gpt-5-mini (faster than qwen3-30b-a3b-thinking-2507). Co-Authored-By: Claude Sonnet 4.6 --- configs/reverse_text/README.md | 2 +- configs/reverse_text/debug_sft_external.toml | 6 +++--- src/prime_rl/orchestrator/orchestrator.py | 6 ++++++ 3 files changed, 10 insertions(+), 4 deletions(-) diff --git a/configs/reverse_text/README.md b/configs/reverse_text/README.md index d544e54aba..4981f19e5b 100644 --- a/configs/reverse_text/README.md +++ b/configs/reverse_text/README.md @@ -7,7 +7,7 @@ Minimal end-to-end configs for the three training modes against the `reverse-tex | `debug_rl.toml` | `rl` | none | | `debug_opd.toml` | `opd` | local vLLM (`Qwen3-0.6B-Reverse-Text-RL`) | | `debug_sft.toml` | `sft` | local vLLM (`Qwen3-0.6B-Reverse-Text-RL`) | -| `debug_sft_external.toml` | `sft` | PI inference (`qwen/qwen3-30b-a3b-thinking-2507`) | +| `debug_sft_external.toml` | `sft` | PI inference (`openai/gpt-5-mini`) | The student inference server is auto-launched on GPU 0 at `http://localhost:8000/v1` with `gpu_memory_utilization=0.5`. The teacher (used by `debug_opd.toml` and `debug_sft.toml`) is **not** auto-launched — start it manually on GPU 1. diff --git a/configs/reverse_text/debug_sft_external.toml b/configs/reverse_text/debug_sft_external.toml index c4baff223a..c29229fc78 100644 --- a/configs/reverse_text/debug_sft_external.toml +++ b/configs/reverse_text/debug_sft_external.toml @@ -1,4 +1,4 @@ -# SFT from qwen3-30b-a3b-thinking-2507 via PI inference. +# SFT from openai/gpt-5-mini via PI inference. # X-Prime-Team-ID header is auto-injected from $PRIME_TEAM_ID for pinference.ai URLs. # # Run with: @@ -13,7 +13,7 @@ name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT" [wandb] project = "reverse-text-debug" -name = "debug-sft-thinking" +name = "debug-sft-external" [orchestrator] batch_size = 128 @@ -37,7 +37,7 @@ max_completion_tokens = 128 id = "reverse-text" [orchestrator.teacher.model] -name = "qwen/qwen3-30b-a3b-thinking-2507" +name = "openai/gpt-5-mini" [orchestrator.teacher.client] base_url = ["https://api.pinference.ai/api/v1"] diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index 546ebfedd1..34b7be1ecc 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -200,6 +200,12 @@ async def orchestrate(config: OrchestratorConfig): # Load environments logger.info("Loading training environments") train_envs = TrainEnvs(config.train.env) + if config.training_mode == "sft": + # Teacher rollouts don't need inference-side logprobs (the trainer + # reconstructs teacher tokens), and some external reasoning-model + # endpoints (e.g. openai/gpt-5*) reject the parameter. + for env in train_envs: + env.sampling_args.pop("logprobs", None) logger.info(f"Loaded {len(train_envs)} training environment(s) ({', '.join(train_envs.names)})") await train_envs.start( From c6ba80f8e4e6d5241b7840574ae604b55354253c Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Tue, 19 May 2026 02:07:10 +0000 Subject: [PATCH 31/47] chore(configs): bump debug_sft_external for gpt-5-mini reasoning budget gpt-5-mini's max_completion_tokens budget includes reasoning tokens. At 128 the model burned the whole budget on internal CoT before producing the final answer, returning empty content. Raise to 2048 and set reasoning_effort=minimal so the chat-completions response includes the actual output. Co-Authored-By: Claude Sonnet 4.6 --- configs/reverse_text/debug_sft_external.toml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/configs/reverse_text/debug_sft_external.toml b/configs/reverse_text/debug_sft_external.toml index c29229fc78..349bea4022 100644 --- a/configs/reverse_text/debug_sft_external.toml +++ b/configs/reverse_text/debug_sft_external.toml @@ -21,7 +21,8 @@ rollouts_per_example = 4 use_renderer = false [orchestrator.train.sampling] -max_completion_tokens = 128 +max_completion_tokens = 2048 +extra_body = { reasoning_effort = "minimal" } [[orchestrator.train.env]] id = "reverse-text" From 2b4fd4e5183f831c99f1d0d08f32294916aedc33 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Tue, 19 May 2026 02:10:31 +0000 Subject: [PATCH 32/47] docs: fix stale SFT references after the student/teacher refactor - skills/config/SKILL.md: update the SFT hard distill section - orchestrator.teacher_model -> orchestrator.teacher, drop the "student inference is optional" guidance (the orchestrator now unconditionally sets up student_inference), and add a note that shared training_mode auto-propagates to trainer.loss.type. - rl.py: auto_setup_inference_client docstring referenced the deleted setup_external_rollout_model function. Rewrite to explain what the validator actually does for each mode. Also note the behavior change in the PR description: SFT now requires [inference] (was previously optional). Co-Authored-By: Claude Sonnet 4.6 --- packages/prime-rl-configs/src/prime_rl/configs/rl.py | 5 +++-- skills/config/SKILL.md | 4 ++-- 2 files changed, 5 insertions(+), 4 deletions(-) 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 815a8606eb..1d7a12f2ce 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/rl.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/rl.py @@ -981,8 +981,9 @@ def auto_setup_inference_client(self): """Auto-configure orchestrator student client from the inference server config. For all modes, sets dp_rank_count from inference DP size. For SFT mode, - also sets base_url so setup_external_rollout_model can detect via - model_fields_set whether the student inference server is actually configured. + also sets base_url - rl/opd rely on the ClientConfig default + (``["http://localhost:8000/v1"]``) which already matches the auto-launched + student vLLM at inference.server.port = 8000. """ if self.inference is None: return self diff --git a/skills/config/SKILL.md b/skills/config/SKILL.md index 4901309cef..c332ffbb0d 100644 --- a/skills/config/SKILL.md +++ b/skills/config/SKILL.md @@ -163,9 +163,9 @@ If you wish to configure values of the default variant, you don't need to set th ### SFT hard distill override -Set `orchestrator.training_mode = "sft"` and configure `orchestrator.teacher_model` with the teacher endpoint. The orchestrator stamps each `TrainingSample.sft_loss = True`, which the trainer's `compute_loss` honors by dispatching to `sft_loss_fn` per batch, independent of the trainer's configured default loss. +Set `orchestrator.training_mode = "sft"` (or top-level `training_mode = "sft"`, which auto-propagates) and configure `orchestrator.teacher` with the teacher endpoint. The orchestrator stamps each `TrainingSample.sft_loss = True` and the shared `training_mode` validator sets `trainer.loss.type = "sft"`, which the trainer's `compute_loss` honors by dispatching to `sft_loss_fn` per batch. -When SFT hard distill also needs online evals or policy weight sync against the student model, configure `[inference]` in the RL entrypoint — this starts the student inference server and auto-configures `orchestrator.model.client`, enabling student weight sync. For externally started student inference, set `orchestrator.model.client.base_url` explicitly. If the student client is not configured, SFT keeps teacher-only rollout behavior and skips student policy updates. +`[inference]` is required (same as rl/opd) — it starts the student inference server and auto-configures `orchestrator.student.client.base_url`. The student pool is used for online evals and policy weight sync. For externally started student inference, set `orchestrator.student.client.base_url` explicitly instead. ### RL rollout client defaults From 34f5a6d70cdbe25a58cf8a67229c7546792be0b3 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Tue, 19 May 2026 02:33:49 +0000 Subject: [PATCH 33/47] feat(loss): batch-driven loss dispatch via TrainingSample.training_mode Replace TrainingSample/MicroBatch.sft_loss bool with training_mode literal (rl|opd|sft). The trainer's compute_loss reads training_mode from the micro batch and dispatches: - "sft" -> sft_loss_fn (masked NLL on teacher tokens) - "rl" -> default_loss_fn with (adv_tau=1.0, teacher_tau=0.0) - "opd" -> default_loss_fn with (adv_tau=0.0, teacher_tau=1.0) Taus move from trainer.loss config to compute_loss's mode->tau table and flow into LossInputs per batch; DefaultLossConfig drops the teacher_tau / adv_tau fields. The orchestrator stamps each sample with config.training_mode at the same point it used to set sft_loss. Packer enforces same-mode packing (samples with different training_mode cannot share a micro batch). RLConfig drops the OPD tau auto-setup + "opd requires teacher_tau > 0" validator + the validate_teacher_model teacher_tau gate (all redundant under batch-driven dispatch). Co-Authored-By: Claude Sonnet 4.6 --- .../src/prime_rl/configs/rl.py | 33 ++++------------- .../src/prime_rl/configs/trainer.py | 2 -- src/prime_rl/entrypoints/rl.py | 4 +-- src/prime_rl/orchestrator/orchestrator.py | 3 +- src/prime_rl/trainer/batch.py | 4 +-- src/prime_rl/trainer/rl/data.py | 11 +++--- src/prime_rl/trainer/rl/loss.py | 36 +++++++++++++++---- src/prime_rl/trainer/rl/train.py | 6 ++-- src/prime_rl/transport/types.py | 12 +++++-- tests/unit/orchestrator/test_batch.py | 18 +++++----- tests/unit/train/rl/test_loss.py | 2 +- 11 files changed, 69 insertions(+), 62 deletions(-) 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 1d7a12f2ce..5fdc04b247 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/rl.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/rl.py @@ -437,10 +437,9 @@ def auto_setup_training_mode(cls, data): """Propagate shared training_mode into orchestrator.training_mode and trainer.loss. Runs before nested validation so that OrchestratorConfig.validate_training_mode - sees the propagated value. Only propagates to fields the user didn't set: - - - sft: trainer.loss.type = "sft" - - opd: trainer.loss.teacher_tau = 1.0, trainer.loss.adv_tau = 0.0 (pure distillation) + sees the propagated value. For sft mode, defaults trainer.loss.type to "sft". + The actual loss dispatch is batch-driven (TrainingSample.training_mode); this + only keeps the trainer's static loss_fn consistent with the run's mode. """ if not isinstance(data, dict): return data @@ -457,33 +456,18 @@ def auto_setup_training_mode(cls, data): if isinstance(orch, dict) and "training_mode" not in orch: orch["training_mode"] = mode - if mode in ("sft", "opd"): + if mode == "sft": trainer = data.setdefault("trainer", {}) if isinstance(trainer, dict): loss = trainer.setdefault("loss", {}) if isinstance(loss, dict): - if mode == "sft": - loss.setdefault("type", "sft") - elif mode == "opd": - loss.setdefault("teacher_tau", 1.0) - loss.setdefault("adv_tau", 0.0) + loss.setdefault("type", "sft") return data - @model_validator(mode="after") - def validate_teacher_model(self): - if ( - self.trainer.loss.type == "default" and self.trainer.loss.teacher_tau > 0 - ) and not self.orchestrator.teacher: - raise ValueError( - "orchestrator.teacher must be configured when teacher_tau > 0. " - "Either set teacher_tau = 0, set deployment.num_teacher_gpus, or configure orchestrator.teacher manually." - ) - return self - @model_validator(mode="after") def validate_training_mode_loss_consistency(self): - """Cross-config invariants between orchestrator.training_mode and trainer.loss.""" + """Cross-config invariants between orchestrator.training_mode and trainer.loss.type.""" mode = self.orchestrator.training_mode loss_type = self.trainer.loss.type @@ -497,11 +481,6 @@ def validate_training_mode_loss_consistency(self): f"trainer.loss.type = 'sft' requires training_mode = 'sft' (got '{mode}'). " "The sft loss path expects teacher-generated rollouts." ) - if mode == "opd" and loss_type == "default" and self.trainer.loss.teacher_tau <= 0: - raise ValueError( - "training_mode = 'opd' requires trainer.loss.teacher_tau > 0. " - "Either set teacher_tau > 0 or change training_mode to 'rl'." - ) return self ### Auto-setup and validate shared configs diff --git a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py index a076d1e29c..7ca435a9be 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py @@ -688,8 +688,6 @@ class DefaultLossConfig(BaseModel): dppo_mask_low: Annotated[float, Field(ge=0, description="The low threshold for masking tokens.")] = 0.2 dppo_mask_high: Annotated[float, Field(ge=0, description="The high threshold for masking tokens.")] = 0.2 - adv_tau: Annotated[float, Field(ge=0, description="The tau for advantages.")] = 1.0 - teacher_tau: Annotated[float, Field(ge=0, description="The tau for teacher logprobs.")] = 0.0 kl_tau: Annotated[float, Field(ge=0, description="The tau for KL divergence.")] = 1e-3 diff --git a/src/prime_rl/entrypoints/rl.py b/src/prime_rl/entrypoints/rl.py index 09eec505d1..db5fcafa1a 100644 --- a/src/prime_rl/entrypoints/rl.py +++ b/src/prime_rl/entrypoints/rl.py @@ -232,9 +232,7 @@ def sigterm_handler(signum, frame): ) monitor_thread.start() monitor_threads.append(monitor_thread) - elif ( - config.trainer.loss.type == "default" and config.trainer.loss.teacher_tau > 0 - ) or config.orchestrator.teacher: + elif config.orchestrator.teacher: logger.warning( "No teacher_inference config specified, skipping starting teacher inference server. " "Is your teacher inference server running? Make sure orchestrator.teacher is configured." diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index 34b7be1ecc..807ab374e7 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -557,8 +557,7 @@ def process_rollout(rollout: vf.RolloutOutput, rollout_idx: int) -> list[Trainin sample.advantage = rollout["advantage"] sample.reward = rollout["reward"] sample.env_name = rollout["env_name"] - if config.training_mode == "sft": - sample.sft_loss = True + sample.training_mode = config.training_mode sample_decode_tokens = sum(sample.completion_mask) sample_prefill_tokens = len(sample.prompt_ids) + len(sample.completion_mask) - sample_decode_tokens rollout_decode_tokens += sample_decode_tokens diff --git a/src/prime_rl/trainer/batch.py b/src/prime_rl/trainer/batch.py index e900aa46cc..e5a16ba05b 100644 --- a/src/prime_rl/trainer/batch.py +++ b/src/prime_rl/trainer/batch.py @@ -81,7 +81,7 @@ def prepare_sample(training_example: TrainingSample, seq_len: int) -> MicroBatch pixel_values=training_example.pixel_values, pixel_values_shape=training_example.pixel_values_shape, image_grid_thw=training_example.image_grid_thw, - sft_loss=training_example.sft_loss, + training_mode=training_example.training_mode, ) @@ -123,7 +123,7 @@ def packed_samples_into_micro_bs( # Check if sequence fits in this bin if ( len(bin_content.input_ids) + len(sample.input_ids) <= max_seq_len - and bin_content.sft_loss == sample.sft_loss + and bin_content.training_mode == sample.training_mode ): bin_content.input_ids.extend(sample.input_ids) bin_content.loss_mask.extend(sample.loss_mask) diff --git a/src/prime_rl/trainer/rl/data.py b/src/prime_rl/trainer/rl/data.py index 6342132ea4..e732db3e55 100644 --- a/src/prime_rl/trainer/rl/data.py +++ b/src/prime_rl/trainer/rl/data.py @@ -40,8 +40,9 @@ class TensorMicroBatch(TypedDict): # mm_token_type_ids: token type per token [batch seq], int64 (0=text, 1=image, 2=video) mm_token_type_ids: Int[Tensor, "batch seq"] | None - # When True, trainer uses SFT loss instead of RL loss for this batch - sft_loss: bool + # Selects loss dispatch (rl/opd → default loss with mode-specific taus, + # sft → sft loss). All samples in a micro batch share the same mode. + training_mode: str class FakeDataLoader: @@ -116,7 +117,7 @@ def _get_sample_micro_batch(self, generator: torch.Generator) -> TensorMicroBatc "pixel_values": None, "image_grid_thw": None, "mm_token_type_ids": None, - "sft_loss": False, + "training_mode": "rl", } def _get_micro_batch(self, generator: torch.Generator) -> TensorMicroBatch: @@ -144,7 +145,7 @@ def _get_micro_batch(self, generator: torch.Generator) -> TensorMicroBatch: "pixel_values": None, "image_grid_thw": None, "mm_token_type_ids": None, - "sft_loss": False, + "training_mode": "rl", } @@ -227,5 +228,5 @@ def _micro_batch_to_tensor(self, micro_batch: MicroBatch) -> TensorMicroBatch: ) # [1, seq_len, layers, topk] if micro_batch.routed_experts is not None else None, - sft_loss=micro_batch.sft_loss, + training_mode=micro_batch.training_mode, ) diff --git a/src/prime_rl/trainer/rl/loss.py b/src/prime_rl/trainer/rl/loss.py index dfb13dc9b1..37ae15d876 100644 --- a/src/prime_rl/trainer/rl/loss.py +++ b/src/prime_rl/trainer/rl/loss.py @@ -19,6 +19,10 @@ class LossInputs: teacher_logprobs: Float[Tensor, " seq"] | None advantages: Float[Tensor, " seq"] loss_mask: Bool[Tensor, " seq"] + # adv_tau / teacher_tau are batch-driven (set in compute_loss from the + # MicroBatch's training_mode). rl: (1.0, 0.0); opd: (0.0, 1.0). + adv_tau: float = 1.0 + teacher_tau: float = 0.0 @dataclass @@ -148,10 +152,10 @@ def default_loss_fn(inputs: LossInputs, loss_config: DefaultLossConfig) -> LossO drop_mask = loss_mask & is_masked keep_mask = loss_mask & ~is_masked - advantages = loss_config.adv_tau * advantages + advantages = inputs.adv_tau * advantages if teacher_logprobs is not None: teacher_kl = teacher_logprobs - trainer_logprobs - advantages = advantages + loss_config.teacher_tau * teacher_kl.detach() + advantages = advantages + inputs.teacher_tau * teacher_kl.detach() else: teacher_kl = None @@ -206,6 +210,13 @@ def loss_fn(inputs: LossInputs) -> LossOutputs: return loss_fn +# Mode -> (adv_tau, teacher_tau). sft is dispatched separately to sft_loss_fn. +_MODE_TO_TAUS: dict[str, tuple[float, float]] = { + "rl": (1.0, 0.0), + "opd": (0.0, 1.0), +} + + def compute_loss( trainer_logprobs: list[Float[Tensor, " seq_i"]], inference_logprobs: list[Float[Tensor, " seq_i"]], @@ -214,25 +225,36 @@ def compute_loss( loss_mask: list[Bool[Tensor, " seq_i"]], loss_fn: LossFn, loss_scale: int, - sft_loss: bool = False, + training_mode: str = "rl", ) -> tuple[Float[Tensor, ""], dict[str, Any]]: """ Compute loss for packed sequences (batch size = 1, multiple sequences packed along sequence dimension). + Loss dispatch is batch-driven via ``training_mode``: + + - ``"sft"``: always use ``sft_loss_fn`` (taus ignored). + - ``"rl"`` / ``"opd"``: use the configured ``loss_fn`` with mode-specific taus + (rl: adv_tau=1, teacher_tau=0; opd: adv_tau=0, teacher_tau=1). + Args: trainer_logprobs: Log probabilities for each sequence inference_logprobs: Reference log probabilities for each sequence teacher_logprobs: Teacher log probabilities for each sequence, or None advantages: Advantages for each sequence loss_mask: Loss mask for each sequence - loss_fn: Per-sequence loss function + loss_fn: Per-sequence loss function for non-sft batches loss_scale: Scale factor to normalize the loss - sft_loss: If True, use SFT loss instead of the configured loss_fn for this batch + training_mode: Selects loss dispatch (rl/opd/sft) Returns: Tuple of (scaled_loss, aggregated_metrics) """ - effective_loss_fn = sft_loss_fn if sft_loss else loss_fn + if training_mode == "sft": + effective_loss_fn = sft_loss_fn + adv_tau, teacher_tau = 0.0, 0.0 # unused + else: + effective_loss_fn = loss_fn + adv_tau, teacher_tau = _MODE_TO_TAUS[training_mode] total_loss = 0.0 all_metrics: dict[str, list[Tensor]] = {} @@ -253,6 +275,8 @@ def compute_loss( teacher_logprobs=teach_logp, advantages=adv, loss_mask=mask, + adv_tau=adv_tau, + teacher_tau=teacher_tau, ) result = effective_loss_fn(inputs) diff --git a/src/prime_rl/trainer/rl/train.py b/src/prime_rl/trainer/rl/train.py index fc03e89f3b..b71e2f54b1 100644 --- a/src/prime_rl/trainer/rl/train.py +++ b/src/prime_rl/trainer/rl/train.py @@ -467,7 +467,7 @@ def load_run_checkpoint(_optimizer, idx: int) -> None: loss_mask=loss_mask.squeeze().split(response_lengths), loss_fn=loss_fn, loss_scale=loss_scale, - sft_loss=micro_batch["sft_loss"], + training_mode=micro_batch["training_mode"], ) # Backward pass @@ -488,7 +488,7 @@ def load_run_checkpoint(_optimizer, idx: int) -> None: for env_name, indices in env_to_indices.items(): tensors[f"entropy/{env_name}"].append(entropy[indices]) - if not micro_batch["sft_loss"]: + if micro_batch["training_mode"] != "sft": with torch.no_grad(): _, _, mismatch_kl = compute_importance_ratio_and_mismatch_kl(out["logprobs"], inference_logprobs) mismatch_kl = mismatch_kl[loss_mask].detach().to("cpu") @@ -508,7 +508,7 @@ def load_run_checkpoint(_optimizer, idx: int) -> None: # Debug log with *local, micro step* stats micro_step_message = f"Micro Step {micro_step}/{len(micro_batches)} | Loss: {tensors['loss'][-1].mean().item():.4f} | Entropy: {tensors['entropy/all'][-1].mean().item():.4f}" - if not micro_batch["sft_loss"]: + if micro_batch["training_mode"] != "sft": micro_step_message += f" | Mismatch KL: {tensors['mismatch_kl/all'][-1].mean().item():.4f}" if "max_vio" in tensors: micro_step_message += f" | Max Vio: {tensors['max_vio'][-1].mean().item():.4f}" diff --git a/src/prime_rl/transport/types.py b/src/prime_rl/transport/types.py index f90af1a39d..332d6dc7a3 100644 --- a/src/prime_rl/transport/types.py +++ b/src/prime_rl/transport/types.py @@ -1,5 +1,9 @@ +from typing import Literal + import msgspec +TrainingMode = Literal["rl", "opd", "sft"] + # Orchestrator -> Packer class TrainingSample(msgspec.Struct, array_like=True, gc=False, omit_defaults=True): @@ -27,7 +31,9 @@ class TrainingSample(msgspec.Struct, array_like=True, gc=False, omit_defaults=Tr # mm_token_type_ids: token type ids per token [batch seq], int64 (0=text, 1=image, 2=video) mm_token_type_ids: list[int] | None = None - sft_loss: bool = False # When True, trainer uses SFT loss instead of RL loss for this sample + # Loss dispatch is batch-driven: rl/opd use default_loss_fn (with mode-specific + # taus), sft uses sft_loss_fn. Stamped by the orchestrator from training_mode. + training_mode: TrainingMode = "rl" class TrainingBatch(msgspec.Struct, array_like=True, gc=False, omit_defaults=True): @@ -61,4 +67,6 @@ class MicroBatch(msgspec.Struct, array_like=True, gc=False, omit_defaults=True): # mm_token_type_ids: token type ids per token [batch seq], int64 (0=text, 1=image, 2=video) mm_token_type_ids: list[int] | None = None - sft_loss: bool = False # When True, trainer uses SFT loss instead of RL loss for this batch + # Loss dispatch is batch-driven (rl/opd → default loss with mode-specific taus, + # sft → sft loss). All samples packed into a micro batch share the same mode. + training_mode: TrainingMode = "rl" diff --git a/tests/unit/orchestrator/test_batch.py b/tests/unit/orchestrator/test_batch.py index 5e4eb65164..e01089ccf4 100644 --- a/tests/unit/orchestrator/test_batch.py +++ b/tests/unit/orchestrator/test_batch.py @@ -8,7 +8,7 @@ def make_training_example(): def _make_training_example( temperature: float = 1.0, - sft_loss: bool = False, + training_mode: str = "rl", env_name: str = "test-env", ) -> TrainingSample: return TrainingSample( @@ -21,7 +21,7 @@ def _make_training_example( teacher_logprobs=[0.0, 0.0, 0.0, 0.0], advantage=1.0, env_name=env_name, - sft_loss=sft_loss, + training_mode=training_mode, ) return _make_training_example @@ -99,17 +99,17 @@ def test_prepare_batch_packs_different_temperatures(make_training_example): assert flat_batches[0].env_names == ["env-a"] * 4 + ["env-b"] * 4 -def test_prepare_sample_propagates_sft_loss(make_training_example): - example = make_training_example(sft_loss=True) +def test_prepare_sample_propagates_training_mode(make_training_example): + example = make_training_example(training_mode="sft") micro_batch = prepare_sample(example, seq_len=16) - assert micro_batch.sft_loss is True + assert micro_batch.training_mode == "sft" -def test_prepare_batch_does_not_pack_mixed_sft_loss(make_training_example): - rl_example = make_training_example(sft_loss=False) - sft_example = make_training_example(sft_loss=True) +def test_prepare_batch_does_not_pack_mixed_training_mode(make_training_example): + rl_example = make_training_example(training_mode="rl") + sft_example = make_training_example(training_mode="sft") batches_per_gpu = prepare_batch( rollouts=[rl_example, sft_example], @@ -121,7 +121,7 @@ def test_prepare_batch_does_not_pack_mixed_sft_loss(make_training_example): flat_batches = [batch for worker_batches in batches_per_gpu for batch in worker_batches] assert len(flat_batches) == 2 - assert {batch.sft_loss for batch in flat_batches} == {False, True} + assert {batch.training_mode for batch in flat_batches} == {"rl", "sft"} def test_prepare_sample_with_routed_experts(): diff --git a/tests/unit/train/rl/test_loss.py b/tests/unit/train/rl/test_loss.py index 696897e368..7b4f401194 100644 --- a/tests/unit/train/rl/test_loss.py +++ b/tests/unit/train/rl/test_loss.py @@ -112,7 +112,7 @@ def test_sft_loss_override_uses_masked_nll_with_default_loss_config(): loss_mask=loss_mask, loss_fn=loss_fn, loss_scale=2, - sft_loss=True, + training_mode="sft", ) assert torch.isclose(loss, torch.tensor(0.15, device=loss.device), atol=1e-6) From 17645e992ce69f796ae3d9f66349d5c0f0216b83 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Tue, 19 May 2026 02:45:47 +0000 Subject: [PATCH 34/47] refactor(loss): dedicated opd_loss_fn instead of routing taus through inputs Replace the per-batch adv_tau/teacher_tau routing through LossInputs with a dedicated opd_loss_fn that lives alongside default_loss_fn and sft_loss_fn. A bit of code dup in opd_loss_fn (clone of default's DPPO math) makes the two paths read as separate losses rather than the same function with magic multipliers. - setup_loss_fn -> setup_loss_fns: returns a dict {training_mode: LossFn}. rl uses the configured loss (default/custom). opd uses opd_loss_fn (only available when loss_config is DefaultLossConfig). sft is always available. - compute_loss takes loss_fns + training_mode, looks up the right fn. - opd_loss_fn requires teacher_logprobs (raises otherwise). - LossInputs no longer carries adv_tau/teacher_tau - those values are baked into each loss fn. - Add validator: opd mode requires trainer.loss.type = "default". Co-Authored-By: Claude Sonnet 4.6 --- .../src/prime_rl/configs/rl.py | 5 + src/prime_rl/trainer/rl/loss.py | 135 ++++++++++++------ src/prime_rl/trainer/rl/train.py | 6 +- tests/unit/train/rl/test_loss.py | 27 ++-- 4 files changed, 111 insertions(+), 62 deletions(-) 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 5fdc04b247..ec5f78b8dc 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/rl.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/rl.py @@ -481,6 +481,11 @@ def validate_training_mode_loss_consistency(self): f"trainer.loss.type = 'sft' requires training_mode = 'sft' (got '{mode}'). " "The sft loss path expects teacher-generated rollouts." ) + if mode == "opd" and loss_type != "default": + raise ValueError( + f"training_mode = 'opd' requires trainer.loss.type = 'default' (got '{loss_type}'). " + "opd_loss_fn reuses the dppo_mask_* and kl_tau knobs from DefaultLossConfig." + ) return self ### Auto-setup and validate shared configs diff --git a/src/prime_rl/trainer/rl/loss.py b/src/prime_rl/trainer/rl/loss.py index 37ae15d876..c725a3fbb3 100644 --- a/src/prime_rl/trainer/rl/loss.py +++ b/src/prime_rl/trainer/rl/loss.py @@ -19,10 +19,6 @@ class LossInputs: teacher_logprobs: Float[Tensor, " seq"] | None advantages: Float[Tensor, " seq"] loss_mask: Bool[Tensor, " seq"] - # adv_tau / teacher_tau are batch-driven (set in compute_loss from the - # MicroBatch's training_mode). rl: (1.0, 0.0); opd: (0.0, 1.0). - adv_tau: float = 1.0 - teacher_tau: float = 0.0 @dataclass @@ -119,7 +115,7 @@ def compute_importance_ratio_and_mismatch_kl( def default_loss_fn(inputs: LossInputs, loss_config: DefaultLossConfig) -> LossOutputs: """ - DPPO+KL loss, combining: + DPPO+KL loss for RL training, combining: - DPPO-Binary TV Loss (https://arxiv.org/pdf/2602.04879) - Kimi-K2.5 KL Loss (https://arxiv.org/pdf/2602.02276) @@ -131,7 +127,6 @@ def default_loss_fn(inputs: LossInputs, loss_config: DefaultLossConfig) -> LossO """ trainer_logprobs = inputs.trainer_logprobs inference_logprobs = inputs.inference_logprobs - teacher_logprobs = inputs.teacher_logprobs advantages = inputs.advantages loss_mask = inputs.loss_mask @@ -152,13 +147,6 @@ def default_loss_fn(inputs: LossInputs, loss_config: DefaultLossConfig) -> LossO drop_mask = loss_mask & is_masked keep_mask = loss_mask & ~is_masked - advantages = inputs.adv_tau * advantages - if teacher_logprobs is not None: - teacher_kl = teacher_logprobs - trainer_logprobs - advantages = advantages + inputs.teacher_tau * teacher_kl.detach() - else: - teacher_kl = None - pg_loss = keep_mask * advantages * importance_ratio kl_loss = loss_mask * log_importance_ratio**2 loss = (-pg_loss + loss_config.kl_tau * kl_loss).sum() @@ -172,8 +160,59 @@ def default_loss_fn(inputs: LossInputs, loss_config: DefaultLossConfig) -> LossO "masked_advantage_positive": _safe_mean(positive_advantages, drop_mask), "masked_advantage_negative": _safe_mean(negative_advantages, drop_mask), } - if teacher_kl is not None: - metrics["teacher_kl"] = _safe_mean(teacher_kl, loss_mask) + + return LossOutputs(loss=loss, metrics=metrics) + + +def opd_loss_fn(inputs: LossInputs, loss_config: DefaultLossConfig) -> LossOutputs: + """ + On-policy distillation loss. Same DPPO+KL machinery as ``default_loss_fn``, + but the per-token policy-gradient signal is the teacher KL (teacher minus + student logprobs over the student's own tokens) instead of the reward + advantage. Equivalent to setting ``adv_tau=0, teacher_tau=1`` on the + classical mixed loss - we just inline it to keep the two paths separate. + """ + trainer_logprobs = inputs.trainer_logprobs + inference_logprobs = inputs.inference_logprobs + teacher_logprobs = inputs.teacher_logprobs + advantages = inputs.advantages # used only for the dppo sign mask + loss_mask = inputs.loss_mask + + if teacher_logprobs is None: + raise ValueError("opd_loss_fn requires teacher_logprobs - configure a teacher for opd mode.") + + log_importance_ratio, importance_ratio, mismatch_kl = compute_importance_ratio_and_mismatch_kl( + trainer_logprobs, inference_logprobs + ) + + probs_diff = torch.exp(trainer_logprobs) - torch.exp(inference_logprobs) + dppo_invalid_mask_high = probs_diff > loss_config.dppo_mask_high + dppo_invalid_mask_low = probs_diff < -loss_config.dppo_mask_low + positive_advantages = advantages > 0 + negative_advantages = advantages < 0 + dppo_invalid_mask = torch.where(positive_advantages, dppo_invalid_mask_high, dppo_invalid_mask_low) + + is_masked = dppo_invalid_mask + is_masked_high = positive_advantages & dppo_invalid_mask_high + is_masked_low = negative_advantages & dppo_invalid_mask_low + drop_mask = loss_mask & is_masked + keep_mask = loss_mask & ~is_masked + + teacher_kl = teacher_logprobs - trainer_logprobs + pg_loss = keep_mask * teacher_kl.detach() * importance_ratio + kl_loss = loss_mask * log_importance_ratio**2 + loss = (-pg_loss + loss_config.kl_tau * kl_loss).sum() + + metrics = { + "masked_mismatch_kl": _safe_mean(mismatch_kl, loss_mask & is_masked), + "unmasked_mismatch_kl": _safe_mean(mismatch_kl, keep_mask), + "is_masked": _safe_mean(is_masked, loss_mask), + "is_masked_low": _safe_mean(is_masked_low, loss_mask), + "is_masked_high": _safe_mean(is_masked_high, loss_mask), + "masked_advantage_positive": _safe_mean(positive_advantages, drop_mask), + "masked_advantage_negative": _safe_mean(negative_advantages, drop_mask), + "teacher_kl": _safe_mean(teacher_kl, loss_mask), + } return LossOutputs(loss=loss, metrics=metrics) @@ -190,31 +229,38 @@ def sft_loss_fn(inputs: LossInputs) -> LossOutputs: return LossOutputs(loss=loss, metrics=metrics) -def setup_loss_fn(loss_config: LossConfig) -> LossFn: - """Setup the loss function based on config.""" +def setup_loss_fns(loss_config: LossConfig) -> dict[str, LossFn]: + """Build the per-training-mode loss fn dispatch table. + + - ``"sft"`` is always available (sft_loss_fn, ignores loss_config). + - ``"rl"`` uses the configured loss (default / custom). + - ``"opd"`` is available only with DefaultLossConfig and uses opd_loss_fn. + """ + fns: dict[str, LossFn] = {"sft": sft_loss_fn} + if isinstance(loss_config, CustomLossConfig): custom_fn = import_object(loss_config.import_path) kwargs = loss_config.kwargs - def loss_fn(inputs: LossInputs) -> LossOutputs: + def rl_fn(inputs: LossInputs) -> LossOutputs: return custom_fn(inputs, **kwargs) - return loss_fn + fns["rl"] = rl_fn + elif isinstance(loss_config, SFTLossConfig): + # sft loss type is only compatible with training_mode = "sft" (validated upstream). + pass + else: # DefaultLossConfig - if isinstance(loss_config, SFTLossConfig): - return sft_loss_fn + def rl_fn(inputs: LossInputs) -> LossOutputs: + return default_loss_fn(inputs, loss_config) - def loss_fn(inputs: LossInputs) -> LossOutputs: - return default_loss_fn(inputs, loss_config) + def opd_fn(inputs: LossInputs) -> LossOutputs: + return opd_loss_fn(inputs, loss_config) - return loss_fn + fns["rl"] = rl_fn + fns["opd"] = opd_fn - -# Mode -> (adv_tau, teacher_tau). sft is dispatched separately to sft_loss_fn. -_MODE_TO_TAUS: dict[str, tuple[float, float]] = { - "rl": (1.0, 0.0), - "opd": (0.0, 1.0), -} + return fns def compute_loss( @@ -223,18 +269,16 @@ def compute_loss( teacher_logprobs: list[Float[Tensor, " seq_i"]] | None, advantages: list[Float[Tensor, " seq_i"]], loss_mask: list[Bool[Tensor, " seq_i"]], - loss_fn: LossFn, + loss_fns: dict[str, LossFn], loss_scale: int, training_mode: str = "rl", ) -> tuple[Float[Tensor, ""], dict[str, Any]]: """ Compute loss for packed sequences (batch size = 1, multiple sequences packed along sequence dimension). - Loss dispatch is batch-driven via ``training_mode``: - - - ``"sft"``: always use ``sft_loss_fn`` (taus ignored). - - ``"rl"`` / ``"opd"``: use the configured ``loss_fn`` with mode-specific taus - (rl: adv_tau=1, teacher_tau=0; opd: adv_tau=0, teacher_tau=1). + Loss dispatch is batch-driven: ``training_mode`` selects the loss fn from + ``loss_fns`` (built by ``setup_loss_fns``). sft → sft_loss_fn, opd → + opd_loss_fn, rl → the configured default/custom loss. Args: trainer_logprobs: Log probabilities for each sequence @@ -242,19 +286,20 @@ def compute_loss( teacher_logprobs: Teacher log probabilities for each sequence, or None advantages: Advantages for each sequence loss_mask: Loss mask for each sequence - loss_fn: Per-sequence loss function for non-sft batches + loss_fns: Per-mode loss fn dispatch table from setup_loss_fns() loss_scale: Scale factor to normalize the loss - training_mode: Selects loss dispatch (rl/opd/sft) + training_mode: Selects which loss fn to apply Returns: Tuple of (scaled_loss, aggregated_metrics) """ - if training_mode == "sft": - effective_loss_fn = sft_loss_fn - adv_tau, teacher_tau = 0.0, 0.0 # unused - else: - effective_loss_fn = loss_fn - adv_tau, teacher_tau = _MODE_TO_TAUS[training_mode] + try: + effective_loss_fn = loss_fns[training_mode] + except KeyError: + raise ValueError( + f"No loss fn available for training_mode={training_mode!r} " + f"(available: {sorted(loss_fns)}). Check trainer.loss.type." + ) total_loss = 0.0 all_metrics: dict[str, list[Tensor]] = {} @@ -275,8 +320,6 @@ def compute_loss( teacher_logprobs=teach_logp, advantages=adv, loss_mask=mask, - adv_tau=adv_tau, - teacher_tau=teacher_tau, ) result = effective_loss_fn(inputs) diff --git a/src/prime_rl/trainer/rl/train.py b/src/prime_rl/trainer/rl/train.py index b71e2f54b1..4e75111907 100644 --- a/src/prime_rl/trainer/rl/train.py +++ b/src/prime_rl/trainer/rl/train.py @@ -31,7 +31,7 @@ compute_loss, compute_importance_ratio_and_mismatch_kl, selective_log_softmax, - setup_loss_fn, + setup_loss_fns, shift_tensor_left, shift_tensor_right, ) @@ -150,7 +150,7 @@ def train(config: TrainerConfig): # Set up the loss function logger.info(f"Setting up loss function ({config.loss})") - loss_fn = setup_loss_fn(config.loss) + loss_fns = setup_loss_fns(config.loss) # Set up the optimizer logger.info(f"Initializing optimizer ({config.optim})") @@ -465,7 +465,7 @@ def load_run_checkpoint(_optimizer, idx: int) -> None: else None, advantages=advantages.squeeze().split(response_lengths), loss_mask=loss_mask.squeeze().split(response_lengths), - loss_fn=loss_fn, + loss_fns=loss_fns, loss_scale=loss_scale, training_mode=micro_batch["training_mode"], ) diff --git a/tests/unit/train/rl/test_loss.py b/tests/unit/train/rl/test_loss.py index 7b4f401194..bc7c010415 100644 --- a/tests/unit/train/rl/test_loss.py +++ b/tests/unit/train/rl/test_loss.py @@ -2,7 +2,7 @@ import torch from prime_rl.configs.trainer import CustomLossConfig, DefaultLossConfig, SFTLossConfig -from prime_rl.trainer.rl.loss import LossInputs, LossOutputs, compute_entropy, compute_loss, setup_loss_fn +from prime_rl.trainer.rl.loss import LossInputs, LossOutputs, compute_entropy, compute_loss, setup_loss_fns pytestmark = [pytest.mark.gpu] @@ -14,14 +14,14 @@ def test_grpo_loss(): advantages = [torch.randn(50).cuda(), torch.randn(30).cuda()] loss_mask = [torch.ones(50, dtype=torch.bool).cuda(), torch.ones(30, dtype=torch.bool).cuda()] - loss_fn = setup_loss_fn(DefaultLossConfig(dppo_mask_high=10.0)) + loss_fns = setup_loss_fns(DefaultLossConfig(dppo_mask_high=10.0)) loss, _ = compute_loss( trainer_logprobs, inference_logprobs, teacher_logprobs, advantages, loss_mask=loss_mask, - loss_fn=loss_fn, + loss_fns=loss_fns, loss_scale=1.0, ) assert loss.shape == () @@ -34,14 +34,14 @@ def test_gspo_loss(): advantages = [torch.randn(40).cuda(), torch.randn(60).cuda()] loss_mask = [torch.ones(40, dtype=torch.bool).cuda(), torch.ones(60, dtype=torch.bool).cuda()] - loss_fn = setup_loss_fn(DefaultLossConfig(dppo_mask_high=10.0)) + loss_fns = setup_loss_fns(DefaultLossConfig(dppo_mask_high=10.0)) loss, _ = compute_loss( trainer_logprobs, inference_logprobs, teacher_logprobs, advantages, loss_mask=loss_mask, - loss_fn=loss_fn, + loss_fns=loss_fns, loss_scale=1.0, ) assert loss.shape == () @@ -53,13 +53,13 @@ def test_entropy_loss(): assert entropy.shape == (10, 10) -def test_setup_loss_fn_with_custom_config(): - """Test setup_loss_fn with CustomLossConfig importing a custom loss.""" +def test_setup_loss_fns_with_custom_config(): + """Test setup_loss_fns with CustomLossConfig importing a custom loss.""" loss_config = CustomLossConfig( import_path="tests.unit.train.rl.test_loss._dummy_custom_loss", kwargs={"multiplier": 2.0}, ) - loss_fn = setup_loss_fn(loss_config) + loss_fns = setup_loss_fns(loss_config) inputs = LossInputs( trainer_logprobs=torch.randn(50, dtype=torch.float32).cuda(), @@ -69,7 +69,7 @@ def test_setup_loss_fn_with_custom_config(): loss_mask=torch.ones(50, dtype=torch.bool).cuda(), ) - result = loss_fn(inputs) + result = loss_fns["rl"](inputs) assert isinstance(result, LossOutputs) assert result.loss.shape == () assert "custom_metric" in result.metrics @@ -81,15 +81,16 @@ def test_sft_loss_matches_masked_nll(): advantages = [torch.zeros(3, dtype=torch.float32).cuda()] loss_mask = [torch.tensor([True, False, True], dtype=torch.bool).cuda()] - loss_fn = setup_loss_fn(SFTLossConfig()) + loss_fns = setup_loss_fns(SFTLossConfig()) loss, metrics = compute_loss( trainer_logprobs=trainer_logprobs, inference_logprobs=inference_logprobs, teacher_logprobs=None, advantages=advantages, loss_mask=loss_mask, - loss_fn=loss_fn, + loss_fns=loss_fns, loss_scale=2, + training_mode="sft", ) # loss = -sum(masked logprobs) / loss_scale = -(-0.1 - 0.2) / 2 = 0.15 @@ -103,14 +104,14 @@ def test_sft_loss_override_uses_masked_nll_with_default_loss_config(): advantages = [torch.ones(3, dtype=torch.float32).cuda()] loss_mask = [torch.tensor([True, False, True], dtype=torch.bool).cuda()] - loss_fn = setup_loss_fn(DefaultLossConfig()) + loss_fns = setup_loss_fns(DefaultLossConfig()) loss, metrics = compute_loss( trainer_logprobs=trainer_logprobs, inference_logprobs=inference_logprobs, teacher_logprobs=None, advantages=advantages, loss_mask=loss_mask, - loss_fn=loss_fn, + loss_fns=loss_fns, loss_scale=2, training_mode="sft", ) From e4a7a4944eff1f44ee9edd414d913c50e49b901e Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Tue, 19 May 2026 02:57:33 +0000 Subject: [PATCH 35/47] refactor(loss): drop trainer.loss type discriminator + shared training_mode Now that loss dispatch is fully batch-driven (TrainingSample.training_mode), the trainer.loss type union and the shared RLConfig.training_mode propagator are redundant. Flatten: - Delete SFTLossConfig (the mode implies it; no fields). - trainer.loss: DefaultLossConfig = DefaultLossConfig() (flat, not a union). - trainer.custom_loss: CustomLossConfig | None = None - optional override for rl-mode loss only; opd and sft always use opd_loss_fn / sft_loss_fn. - setup_loss_fns(loss_config, custom_loss) always returns all three keys - mixed-mode batches dispatch correctly at the trainer level (orchestrator still stamps a single mode per run; that's a separate concern). - Drop RLConfig.training_mode shared field, the auto_setup_training_mode before-validator, and validate_training_mode_loss_consistency. The orchestrator-level [orchestrator] training_mode is the single source. Move training_mode = "sft" out of top-level into [orchestrator] in the three SFT debug configs. Co-Authored-By: Claude Sonnet 4.6 --- configs/reverse_text/debug_sft.toml | 2 +- configs/reverse_text/debug_sft_external.toml | 2 +- configs/reverse_text/debug_sft_lora.toml | 58 ++++++++++++++++ .../src/prime_rl/configs/rl.py | 66 ------------------- .../src/prime_rl/configs/trainer.py | 27 +++----- src/prime_rl/trainer/rl/loss.py | 42 ++++++------ src/prime_rl/trainer/rl/train.py | 2 +- tests/unit/train/rl/test_loss.py | 4 +- 8 files changed, 93 insertions(+), 110 deletions(-) create mode 100644 configs/reverse_text/debug_sft_lora.toml diff --git a/configs/reverse_text/debug_sft.toml b/configs/reverse_text/debug_sft.toml index a9c6c96214..8620fccf7f 100644 --- a/configs/reverse_text/debug_sft.toml +++ b/configs/reverse_text/debug_sft.toml @@ -7,7 +7,6 @@ max_steps = 20 seq_len = 2048 -training_mode = "sft" [model] name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT" @@ -17,6 +16,7 @@ project = "reverse-text-debug" name = "debug-sft" [orchestrator] +training_mode = "sft" batch_size = 128 rollouts_per_example = 4 use_renderer = false diff --git a/configs/reverse_text/debug_sft_external.toml b/configs/reverse_text/debug_sft_external.toml index 349bea4022..bf6fd3219c 100644 --- a/configs/reverse_text/debug_sft_external.toml +++ b/configs/reverse_text/debug_sft_external.toml @@ -6,7 +6,6 @@ max_steps = 20 seq_len = 2048 -training_mode = "sft" [model] name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT" @@ -16,6 +15,7 @@ project = "reverse-text-debug" name = "debug-sft-external" [orchestrator] +training_mode = "sft" batch_size = 128 rollouts_per_example = 4 use_renderer = false diff --git a/configs/reverse_text/debug_sft_lora.toml b/configs/reverse_text/debug_sft_lora.toml new file mode 100644 index 0000000000..16288ab7c3 --- /dev/null +++ b/configs/reverse_text/debug_sft_lora.toml @@ -0,0 +1,58 @@ +# Start the teacher inference server first (on a separate GPU): +# CUDA_VISIBLE_DEVICES=1 uv run inference \ +# --model.name PrimeIntellect/Qwen3-0.6B-Reverse-Text-RL \ +# --server.port 8001 --gpu-memory-utilization 0.5 --model.enforce-eager +# Then: +# uv run rl @ configs/reverse_text/debug_sft_lora.toml + +max_steps = 20 +seq_len = 2048 + +[model] +name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT" + +[wandb] +project = "reverse-text-debug" +name = "debug-sft-lora" + +[orchestrator] +training_mode = "sft" +batch_size = 128 +rollouts_per_example = 4 +use_renderer = false + +[orchestrator.train.sampling] +max_completion_tokens = 128 + +[[orchestrator.train.env]] +id = "reverse-text" + +[orchestrator.eval] +interval = 1 +num_examples = 128 + +[orchestrator.eval.sampling] +max_completion_tokens = 128 + +[[orchestrator.eval.env]] +id = "reverse-text" + +[orchestrator.teacher.model] +name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-RL" + +[orchestrator.teacher.client] +base_url = ["http://localhost:8001/v1"] + +[trainer.optim] +lr = 5e-5 + +[trainer.model.lora] +rank = 8 + +[trainer.ckpt.weights] +save_adapter_separately = true + +[ckpt] + +[inference] +gpu_memory_utilization = 0.5 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 ec5f78b8dc..9421c3957e 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/rl.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/rl.py @@ -313,15 +313,6 @@ class RLConfig(BaseConfig): ), ] = None - training_mode: Annotated[ - Literal["rl", "opd", "sft"] | None, - Field( - description="Shared training mode. Propagates to orchestrator.training_mode and, " - "for 'sft', switches trainer.loss to SFTLossConfig. " - "Explicitly set per-component values always take precedence." - ), - ] = None - max_steps: Annotated[ int | None, Field( @@ -431,63 +422,6 @@ def validate_quantize_in_weight_transfer(self): return self - @model_validator(mode="before") - @classmethod - def auto_setup_training_mode(cls, data): - """Propagate shared training_mode into orchestrator.training_mode and trainer.loss. - - Runs before nested validation so that OrchestratorConfig.validate_training_mode - sees the propagated value. For sft mode, defaults trainer.loss.type to "sft". - The actual loss dispatch is batch-driven (TrainingSample.training_mode); this - only keeps the trainer's static loss_fn consistent with the run's mode. - """ - if not isinstance(data, dict): - return data - mode = data.get("training_mode") - if mode is None: - # Also accept training_mode set only inside [orchestrator] - orch_candidate = data.get("orchestrator") - if isinstance(orch_candidate, dict): - mode = orch_candidate.get("training_mode") - if mode is None: - return data - - orch = data.setdefault("orchestrator", {}) - if isinstance(orch, dict) and "training_mode" not in orch: - orch["training_mode"] = mode - - if mode == "sft": - trainer = data.setdefault("trainer", {}) - if isinstance(trainer, dict): - loss = trainer.setdefault("loss", {}) - if isinstance(loss, dict): - loss.setdefault("type", "sft") - - return data - - @model_validator(mode="after") - def validate_training_mode_loss_consistency(self): - """Cross-config invariants between orchestrator.training_mode and trainer.loss.type.""" - mode = self.orchestrator.training_mode - loss_type = self.trainer.loss.type - - if mode == "sft" and loss_type != "sft": - raise ValueError( - f"training_mode = 'sft' requires trainer.loss.type = 'sft' (got '{loss_type}'). " - "Either set trainer.loss.type = 'sft' or change training_mode." - ) - if mode in ("rl", "opd") and loss_type == "sft": - raise ValueError( - f"trainer.loss.type = 'sft' requires training_mode = 'sft' (got '{mode}'). " - "The sft loss path expects teacher-generated rollouts." - ) - if mode == "opd" and loss_type != "default": - raise ValueError( - f"training_mode = 'opd' requires trainer.loss.type = 'default' (got '{loss_type}'). " - "opd_loss_fn reuses the dppo_mask_* and kl_tau knobs from DefaultLossConfig." - ) - return self - ### Auto-setup and validate shared configs @model_validator(mode="after") diff --git a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py index 7ca435a9be..b98f3ce908 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py @@ -682,33 +682,22 @@ class CheckpointConfig(BaseConfig): class DefaultLossConfig(BaseModel): - """Config for the default loss.""" - - type: Literal["default"] = "default" + """Knobs for the default DPPO+KL loss math, shared by ``default_loss_fn`` (rl) + and ``opd_loss_fn`` (opd). The actual loss fn is selected per batch from + ``TrainingSample.training_mode``.""" dppo_mask_low: Annotated[float, Field(ge=0, description="The low threshold for masking tokens.")] = 0.2 dppo_mask_high: Annotated[float, Field(ge=0, description="The high threshold for masking tokens.")] = 0.2 kl_tau: Annotated[float, Field(ge=0, description="The tau for KL divergence.")] = 1e-3 -class SFTLossConfig(BaseModel): - """Config for SFT-style masked negative log-likelihood loss.""" - - type: Literal["sft"] = "sft" - - class CustomLossConfig(BaseModel): - """Config for a custom external loss function.""" - - type: Literal["custom"] = "custom" + """Optional override for the rl-mode loss fn. opd and sft are unaffected.""" import_path: Annotated[str, Field(description="Import path to the loss function (e.g., 'my_module.my_loss')")] kwargs: Annotated[dict[str, Any], Field(default_factory=dict, description="Kwargs to pass to the loss function")] -LossConfig: TypeAlias = Annotated[DefaultLossConfig | SFTLossConfig | CustomLossConfig, Field(discriminator="type")] - - class FakeDataLoaderConfig(BaseConfig): """Configures a fake data loader sampling random micro batches for debugging.""" @@ -781,8 +770,12 @@ class TrainerConfig(BaseConfig): # The data configuration data: DataLoaderConfig = DataLoaderConfig() - # The loss configuration - loss: LossConfig = DefaultLossConfig() + # DPPO+KL knobs (shared by default_loss_fn and opd_loss_fn); selection of + # which loss fn runs is driven by TrainingSample.training_mode, not by config. + loss: DefaultLossConfig = DefaultLossConfig() + + # Optional override for rl-mode loss only. opd and sft are unaffected. + custom_loss: CustomLossConfig | None = None # The optimizer configuration optim: OptimizerConfig = AdamWConfig() diff --git a/src/prime_rl/trainer/rl/loss.py b/src/prime_rl/trainer/rl/loss.py index c725a3fbb3..eae27e4c95 100644 --- a/src/prime_rl/trainer/rl/loss.py +++ b/src/prime_rl/trainer/rl/loss.py @@ -6,7 +6,7 @@ from jaxtyping import Bool, Float, Int, jaxtyped from torch import Tensor -from prime_rl.configs.trainer import CustomLossConfig, DefaultLossConfig, LossConfig, SFTLossConfig +from prime_rl.configs.trainer import CustomLossConfig, DefaultLossConfig from prime_rl.utils.utils import import_object @@ -229,38 +229,36 @@ def sft_loss_fn(inputs: LossInputs) -> LossOutputs: return LossOutputs(loss=loss, metrics=metrics) -def setup_loss_fns(loss_config: LossConfig) -> dict[str, LossFn]: +def setup_loss_fns( + loss_config: DefaultLossConfig, + custom_loss: CustomLossConfig | None = None, +) -> dict[str, LossFn]: """Build the per-training-mode loss fn dispatch table. - - ``"sft"`` is always available (sft_loss_fn, ignores loss_config). - - ``"rl"`` uses the configured loss (default / custom). - - ``"opd"`` is available only with DefaultLossConfig and uses opd_loss_fn. + Always returns all three modes - the trainer is mode-agnostic and routes + per batch from ``TrainingSample.training_mode``: + + - ``"sft"`` → ``sft_loss_fn`` (masked NLL on teacher tokens) + - ``"opd"`` → ``opd_loss_fn`` (teacher KL as gradient signal, DPPO + KL machinery) + - ``"rl"`` → ``default_loss_fn`` with the DPPO+KL knobs, or the + ``custom_loss`` override if configured. """ - fns: dict[str, LossFn] = {"sft": sft_loss_fn} - if isinstance(loss_config, CustomLossConfig): - custom_fn = import_object(loss_config.import_path) - kwargs = loss_config.kwargs + def opd_fn(inputs: LossInputs) -> LossOutputs: + return opd_loss_fn(inputs, loss_config) + + if custom_loss is not None: + custom_fn = import_object(custom_loss.import_path) + kwargs = custom_loss.kwargs def rl_fn(inputs: LossInputs) -> LossOutputs: return custom_fn(inputs, **kwargs) - - fns["rl"] = rl_fn - elif isinstance(loss_config, SFTLossConfig): - # sft loss type is only compatible with training_mode = "sft" (validated upstream). - pass - else: # DefaultLossConfig + else: def rl_fn(inputs: LossInputs) -> LossOutputs: return default_loss_fn(inputs, loss_config) - def opd_fn(inputs: LossInputs) -> LossOutputs: - return opd_loss_fn(inputs, loss_config) - - fns["rl"] = rl_fn - fns["opd"] = opd_fn - - return fns + return {"sft": sft_loss_fn, "opd": opd_fn, "rl": rl_fn} def compute_loss( diff --git a/src/prime_rl/trainer/rl/train.py b/src/prime_rl/trainer/rl/train.py index 4e75111907..af7084af6f 100644 --- a/src/prime_rl/trainer/rl/train.py +++ b/src/prime_rl/trainer/rl/train.py @@ -150,7 +150,7 @@ def train(config: TrainerConfig): # Set up the loss function logger.info(f"Setting up loss function ({config.loss})") - loss_fns = setup_loss_fns(config.loss) + loss_fns = setup_loss_fns(config.loss, config.custom_loss) # Set up the optimizer logger.info(f"Initializing optimizer ({config.optim})") diff --git a/tests/unit/train/rl/test_loss.py b/tests/unit/train/rl/test_loss.py index bc7c010415..1585dac7bd 100644 --- a/tests/unit/train/rl/test_loss.py +++ b/tests/unit/train/rl/test_loss.py @@ -1,7 +1,7 @@ import pytest import torch -from prime_rl.configs.trainer import CustomLossConfig, DefaultLossConfig, SFTLossConfig +from prime_rl.configs.trainer import CustomLossConfig, DefaultLossConfig from prime_rl.trainer.rl.loss import LossInputs, LossOutputs, compute_entropy, compute_loss, setup_loss_fns pytestmark = [pytest.mark.gpu] @@ -81,7 +81,7 @@ def test_sft_loss_matches_masked_nll(): advantages = [torch.zeros(3, dtype=torch.float32).cuda()] loss_mask = [torch.tensor([True, False, True], dtype=torch.bool).cuda()] - loss_fns = setup_loss_fns(SFTLossConfig()) + loss_fns = setup_loss_fns(DefaultLossConfig()) loss, metrics = compute_loss( trainer_logprobs=trainer_logprobs, inference_logprobs=inference_logprobs, From 665bf7f6de7306c70db3460508d9907bfc875023 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Tue, 19 May 2026 03:11:14 +0000 Subject: [PATCH 36/47] fix(rl entrypoint): warn for missing [inference] in all modes (incl. sft) After the always-require-student refactor, sft runs need a student inference pool too (orchestrator unconditionally creates it for evals + weight sync). The old "sft mode, using teacher for rollouts" info log was correct under the old semantics but now silently puts the user on a path where the orchestrator hangs waiting for a non-existent student server at the default localhost:8000. Replace with a single warning that prints the configured student base_url and flags the hang risk. Co-Authored-By: Claude Sonnet 4.6 --- src/prime_rl/entrypoints/rl.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/src/prime_rl/entrypoints/rl.py b/src/prime_rl/entrypoints/rl.py index db5fcafa1a..b4d7ccb068 100644 --- a/src/prime_rl/entrypoints/rl.py +++ b/src/prime_rl/entrypoints/rl.py @@ -191,12 +191,13 @@ def sigterm_handler(signum, frame): monitor_thread.start() monitor_threads.append(monitor_thread) else: - if config.orchestrator.training_mode != "sft": - logger.warning( - "No inference config specified, skipping starting inference server. Make sure your inference server is running." - ) - else: - logger.info("No inference config specified, using teacher model for rollout generation (sft mode).") + logger.warning( + "No [inference] block configured - the student inference server will not be started here. " + "All training modes (rl/opd/sft) require a student inference pool for evals + weight sync; " + "make sure one is running at orchestrator.student.client.base_url " + f"({', '.join(config.orchestrator.student.client.base_url)}), otherwise the orchestrator " + "will hang waiting for it." + ) # Optionally, start teacher inference process if config.teacher_inference: From 03c3a696d322e97bced52f0e37d86e136771a8cc Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Tue, 19 May 2026 03:13:05 +0000 Subject: [PATCH 37/47] refactor(loss): make opd_loss_fn's tau values explicit (adv_tau=0, teacher_tau=1) Mirror default_loss_fn's structure literally: declare the two tau locals at the top, then use them in the same `adv_tau * advantages + teacher_tau * teacher_kl.detach()` expression. The "this is the default loss with those specific knobs baked in" relationship is now visible in the code rather than hidden in a substituted expression. Co-Authored-By: Claude Sonnet 4.6 --- src/prime_rl/trainer/rl/loss.py | 19 ++++++++++++------- 1 file changed, 12 insertions(+), 7 deletions(-) diff --git a/src/prime_rl/trainer/rl/loss.py b/src/prime_rl/trainer/rl/loss.py index eae27e4c95..79df278a81 100644 --- a/src/prime_rl/trainer/rl/loss.py +++ b/src/prime_rl/trainer/rl/loss.py @@ -166,16 +166,19 @@ def default_loss_fn(inputs: LossInputs, loss_config: DefaultLossConfig) -> LossO def opd_loss_fn(inputs: LossInputs, loss_config: DefaultLossConfig) -> LossOutputs: """ - On-policy distillation loss. Same DPPO+KL machinery as ``default_loss_fn``, - but the per-token policy-gradient signal is the teacher KL (teacher minus - student logprobs over the student's own tokens) instead of the reward - advantage. Equivalent to setting ``adv_tau=0, teacher_tau=1`` on the - classical mixed loss - we just inline it to keep the two paths separate. + On-policy distillation loss: the default DPPO+KL math with the tau knobs + hardcoded to drop the reward signal and use the teacher KL as the + per-token policy-gradient signal. Equivalent to ``default_loss_fn`` with + ``adv_tau = 0`` and ``teacher_tau = 1``; we inline both here so the two + paths read as distinct losses. """ + adv_tau = 0.0 + teacher_tau = 1.0 + trainer_logprobs = inputs.trainer_logprobs inference_logprobs = inputs.inference_logprobs teacher_logprobs = inputs.teacher_logprobs - advantages = inputs.advantages # used only for the dppo sign mask + advantages = inputs.advantages loss_mask = inputs.loss_mask if teacher_logprobs is None: @@ -199,7 +202,9 @@ def opd_loss_fn(inputs: LossInputs, loss_config: DefaultLossConfig) -> LossOutpu keep_mask = loss_mask & ~is_masked teacher_kl = teacher_logprobs - trainer_logprobs - pg_loss = keep_mask * teacher_kl.detach() * importance_ratio + advantages = adv_tau * advantages + teacher_tau * teacher_kl.detach() + + pg_loss = keep_mask * advantages * importance_ratio kl_loss = loss_mask * log_importance_ratio**2 loss = (-pg_loss + loss_config.kl_tau * kl_loss).sum() From dda0d916bacbe60210044446e25fb9e83d57d82b Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Tue, 19 May 2026 03:20:30 +0000 Subject: [PATCH 38/47] refactor(loss): restore DefaultLossConfig.adv_tau + LossConfig union (rl-only) - DefaultLossConfig regains adv_tau (= 1.0). default_loss_fn reads it. - trainer.loss is back to a discriminated union of DefaultLossConfig | CustomLossConfig (drop SFTLossConfig stays dropped). type discriminator restored. - trainer.loss only applies to rl-mode batches. opd and sft batches dispatch unconditionally to opd_loss_fn / sft_loss_fn - they don't read trainer.loss. - opd_loss_fn is self-contained: the DPPO/KL knobs (dppo_mask_*, kl_tau) are baked in as module-level _OPD_* constants matching DefaultLossConfig's defaults. No loss_config parameter. - Drop the trainer.custom_loss sibling field (subsumed by the union). Also add debug_opd_lora.toml + bump lr to 1e-4 in both lora debug configs. Co-Authored-By: Claude Sonnet 4.6 --- configs/reverse_text/debug_opd_lora.toml | 60 +++++++++++++++++++ configs/reverse_text/debug_sft_lora.toml | 2 +- .../src/prime_rl/configs/trainer.py | 24 ++++---- src/prime_rl/trainer/rl/loss.py | 48 ++++++++------- src/prime_rl/trainer/rl/train.py | 2 +- 5 files changed, 104 insertions(+), 32 deletions(-) create mode 100644 configs/reverse_text/debug_opd_lora.toml diff --git a/configs/reverse_text/debug_opd_lora.toml b/configs/reverse_text/debug_opd_lora.toml new file mode 100644 index 0000000000..9150a6b3fe --- /dev/null +++ b/configs/reverse_text/debug_opd_lora.toml @@ -0,0 +1,60 @@ +# Start the teacher inference server first (on a separate GPU): +# CUDA_VISIBLE_DEVICES=1 uv run inference \ +# --model.name PrimeIntellect/Qwen3-0.6B-Reverse-Text-RL \ +# --server.port 8001 --gpu-memory-utilization 0.5 --model.enforce-eager +# Then: +# uv run rl @ configs/reverse_text/debug_opd_lora.toml + +max_steps = 20 +seq_len = 2048 + +[model] +name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT" + +[wandb] +project = "reverse-text-debug" +name = "debug-opd-lora" + +[orchestrator] +training_mode = "opd" +batch_size = 128 +rollouts_per_example = 16 + +[orchestrator.renderer] +name = "qwen3" + +[orchestrator.train.sampling] +max_completion_tokens = 128 + +[[orchestrator.train.env]] +id = "reverse-text" + +[orchestrator.eval] +interval = 1 +num_examples = 128 + +[orchestrator.eval.sampling] +max_completion_tokens = 128 + +[[orchestrator.eval.env]] +id = "reverse-text" + +[orchestrator.teacher.model] +name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-RL" + +[orchestrator.teacher.client] +base_url = ["http://localhost:8001/v1"] + +[trainer.optim] +lr = 1e-4 + +[trainer.model.lora] +rank = 8 + +[trainer.ckpt.weights] +save_adapter_separately = true + +[ckpt] + +[inference] +gpu_memory_utilization = 0.5 diff --git a/configs/reverse_text/debug_sft_lora.toml b/configs/reverse_text/debug_sft_lora.toml index 16288ab7c3..64a048b34f 100644 --- a/configs/reverse_text/debug_sft_lora.toml +++ b/configs/reverse_text/debug_sft_lora.toml @@ -44,7 +44,7 @@ name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-RL" base_url = ["http://localhost:8001/v1"] [trainer.optim] -lr = 5e-5 +lr = 1e-4 [trainer.model.lora] rank = 8 diff --git a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py index b98f3ce908..9e15a1a7ee 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py @@ -682,22 +682,29 @@ class CheckpointConfig(BaseConfig): class DefaultLossConfig(BaseModel): - """Knobs for the default DPPO+KL loss math, shared by ``default_loss_fn`` (rl) - and ``opd_loss_fn`` (opd). The actual loss fn is selected per batch from - ``TrainingSample.training_mode``.""" + """Knobs for the default DPPO+KL loss. Only consumed by rl-mode batches + (``default_loss_fn``); opd and sft have their own self-contained loss fns.""" + + type: Literal["default"] = "default" dppo_mask_low: Annotated[float, Field(ge=0, description="The low threshold for masking tokens.")] = 0.2 dppo_mask_high: Annotated[float, Field(ge=0, description="The high threshold for masking tokens.")] = 0.2 + adv_tau: Annotated[float, Field(ge=0, description="The tau for advantages.")] = 1.0 kl_tau: Annotated[float, Field(ge=0, description="The tau for KL divergence.")] = 1e-3 class CustomLossConfig(BaseModel): - """Optional override for the rl-mode loss fn. opd and sft are unaffected.""" + """Override for the rl-mode loss fn. opd and sft are unaffected.""" + + type: Literal["custom"] = "custom" import_path: Annotated[str, Field(description="Import path to the loss function (e.g., 'my_module.my_loss')")] kwargs: Annotated[dict[str, Any], Field(default_factory=dict, description="Kwargs to pass to the loss function")] +LossConfig: TypeAlias = Annotated[DefaultLossConfig | CustomLossConfig, Field(discriminator="type")] + + class FakeDataLoaderConfig(BaseConfig): """Configures a fake data loader sampling random micro batches for debugging.""" @@ -770,12 +777,9 @@ class TrainerConfig(BaseConfig): # The data configuration data: DataLoaderConfig = DataLoaderConfig() - # DPPO+KL knobs (shared by default_loss_fn and opd_loss_fn); selection of - # which loss fn runs is driven by TrainingSample.training_mode, not by config. - loss: DefaultLossConfig = DefaultLossConfig() - - # Optional override for rl-mode loss only. opd and sft are unaffected. - custom_loss: CustomLossConfig | None = None + # Loss config for the rl-mode batches only. opd and sft batches dispatch to + # opd_loss_fn / sft_loss_fn unconditionally - they don't read trainer.loss. + loss: LossConfig = DefaultLossConfig() # The optimizer configuration optim: OptimizerConfig = AdamWConfig() diff --git a/src/prime_rl/trainer/rl/loss.py b/src/prime_rl/trainer/rl/loss.py index 79df278a81..3f68e939ee 100644 --- a/src/prime_rl/trainer/rl/loss.py +++ b/src/prime_rl/trainer/rl/loss.py @@ -6,7 +6,7 @@ from jaxtyping import Bool, Float, Int, jaxtyped from torch import Tensor -from prime_rl.configs.trainer import CustomLossConfig, DefaultLossConfig +from prime_rl.configs.trainer import CustomLossConfig, DefaultLossConfig, LossConfig from prime_rl.utils.utils import import_object @@ -147,6 +147,7 @@ def default_loss_fn(inputs: LossInputs, loss_config: DefaultLossConfig) -> LossO drop_mask = loss_mask & is_masked keep_mask = loss_mask & ~is_masked + advantages = loss_config.adv_tau * advantages pg_loss = keep_mask * advantages * importance_ratio kl_loss = loss_mask * log_importance_ratio**2 loss = (-pg_loss + loss_config.kl_tau * kl_loss).sum() @@ -164,13 +165,24 @@ def default_loss_fn(inputs: LossInputs, loss_config: DefaultLossConfig) -> LossO return LossOutputs(loss=loss, metrics=metrics) -def opd_loss_fn(inputs: LossInputs, loss_config: DefaultLossConfig) -> LossOutputs: +# OPD knobs are baked into opd_loss_fn (not user-configurable via trainer.loss, +# which only applies to rl). Defaults mirror DefaultLossConfig's so the dppo/kl +# behavior matches the rl path. +_OPD_DPPO_MASK_HIGH = 0.2 +_OPD_DPPO_MASK_LOW = 0.2 +_OPD_KL_TAU = 1e-3 + + +def opd_loss_fn(inputs: LossInputs) -> LossOutputs: """ On-policy distillation loss: the default DPPO+KL math with the tau knobs hardcoded to drop the reward signal and use the teacher KL as the per-token policy-gradient signal. Equivalent to ``default_loss_fn`` with ``adv_tau = 0`` and ``teacher_tau = 1``; we inline both here so the two paths read as distinct losses. + + Self-contained: doesn't read ``trainer.loss`` (which is rl-only). The + dppo/kl knobs are baked in. """ adv_tau = 0.0 teacher_tau = 1.0 @@ -189,8 +201,8 @@ def opd_loss_fn(inputs: LossInputs, loss_config: DefaultLossConfig) -> LossOutpu ) probs_diff = torch.exp(trainer_logprobs) - torch.exp(inference_logprobs) - dppo_invalid_mask_high = probs_diff > loss_config.dppo_mask_high - dppo_invalid_mask_low = probs_diff < -loss_config.dppo_mask_low + dppo_invalid_mask_high = probs_diff > _OPD_DPPO_MASK_HIGH + dppo_invalid_mask_low = probs_diff < -_OPD_DPPO_MASK_LOW positive_advantages = advantages > 0 negative_advantages = advantages < 0 dppo_invalid_mask = torch.where(positive_advantages, dppo_invalid_mask_high, dppo_invalid_mask_low) @@ -206,7 +218,7 @@ def opd_loss_fn(inputs: LossInputs, loss_config: DefaultLossConfig) -> LossOutpu pg_loss = keep_mask * advantages * importance_ratio kl_loss = loss_mask * log_importance_ratio**2 - loss = (-pg_loss + loss_config.kl_tau * kl_loss).sum() + loss = (-pg_loss + _OPD_KL_TAU * kl_loss).sum() metrics = { "masked_mismatch_kl": _safe_mean(mismatch_kl, loss_mask & is_masked), @@ -234,27 +246,23 @@ def sft_loss_fn(inputs: LossInputs) -> LossOutputs: return LossOutputs(loss=loss, metrics=metrics) -def setup_loss_fns( - loss_config: DefaultLossConfig, - custom_loss: CustomLossConfig | None = None, -) -> dict[str, LossFn]: +def setup_loss_fns(loss_config: LossConfig) -> dict[str, LossFn]: """Build the per-training-mode loss fn dispatch table. Always returns all three modes - the trainer is mode-agnostic and routes per batch from ``TrainingSample.training_mode``: - ``"sft"`` → ``sft_loss_fn`` (masked NLL on teacher tokens) - - ``"opd"`` → ``opd_loss_fn`` (teacher KL as gradient signal, DPPO + KL machinery) - - ``"rl"`` → ``default_loss_fn`` with the DPPO+KL knobs, or the - ``custom_loss`` override if configured. - """ + - ``"opd"`` → ``opd_loss_fn`` (teacher KL as gradient signal, hardcoded + DPPO + KL knobs) + - ``"rl"`` → ``default_loss_fn(loss_config)`` for ``DefaultLossConfig``, + or the imported function for ``CustomLossConfig``. - def opd_fn(inputs: LossInputs) -> LossOutputs: - return opd_loss_fn(inputs, loss_config) - - if custom_loss is not None: - custom_fn = import_object(custom_loss.import_path) - kwargs = custom_loss.kwargs + ``trainer.loss`` only affects the rl path - opd and sft are independent. + """ + if isinstance(loss_config, CustomLossConfig): + custom_fn = import_object(loss_config.import_path) + kwargs = loss_config.kwargs def rl_fn(inputs: LossInputs) -> LossOutputs: return custom_fn(inputs, **kwargs) @@ -263,7 +271,7 @@ def rl_fn(inputs: LossInputs) -> LossOutputs: def rl_fn(inputs: LossInputs) -> LossOutputs: return default_loss_fn(inputs, loss_config) - return {"sft": sft_loss_fn, "opd": opd_fn, "rl": rl_fn} + return {"sft": sft_loss_fn, "opd": opd_loss_fn, "rl": rl_fn} def compute_loss( diff --git a/src/prime_rl/trainer/rl/train.py b/src/prime_rl/trainer/rl/train.py index af7084af6f..4e75111907 100644 --- a/src/prime_rl/trainer/rl/train.py +++ b/src/prime_rl/trainer/rl/train.py @@ -150,7 +150,7 @@ def train(config: TrainerConfig): # Set up the loss function logger.info(f"Setting up loss function ({config.loss})") - loss_fns = setup_loss_fns(config.loss, config.custom_loss) + loss_fns = setup_loss_fns(config.loss) # Set up the optimizer logger.info(f"Initializing optimizer ({config.optim})") From b3a69225ab49d0d00eda91dfb18788a34fac5373 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Tue, 19 May 2026 03:23:03 +0000 Subject: [PATCH 39/47] refactor(loss): inline opd_loss_fn knobs (no named locals/constants) Drop the module-level _OPD_* constants and the local adv_tau/teacher_tau bindings - inline the literal values into the math. Keeps the parallel structure to default_loss_fn visible: `0.0 * advantages + 1.0 * teacher_kl` in the same place default_loss_fn has `loss_config.adv_tau * advantages`. Co-Authored-By: Claude Sonnet 4.6 --- src/prime_rl/trainer/rl/loss.py | 21 +++++---------------- 1 file changed, 5 insertions(+), 16 deletions(-) diff --git a/src/prime_rl/trainer/rl/loss.py b/src/prime_rl/trainer/rl/loss.py index 3f68e939ee..d7c5b655b2 100644 --- a/src/prime_rl/trainer/rl/loss.py +++ b/src/prime_rl/trainer/rl/loss.py @@ -165,14 +165,6 @@ def default_loss_fn(inputs: LossInputs, loss_config: DefaultLossConfig) -> LossO return LossOutputs(loss=loss, metrics=metrics) -# OPD knobs are baked into opd_loss_fn (not user-configurable via trainer.loss, -# which only applies to rl). Defaults mirror DefaultLossConfig's so the dppo/kl -# behavior matches the rl path. -_OPD_DPPO_MASK_HIGH = 0.2 -_OPD_DPPO_MASK_LOW = 0.2 -_OPD_KL_TAU = 1e-3 - - def opd_loss_fn(inputs: LossInputs) -> LossOutputs: """ On-policy distillation loss: the default DPPO+KL math with the tau knobs @@ -182,11 +174,8 @@ def opd_loss_fn(inputs: LossInputs) -> LossOutputs: paths read as distinct losses. Self-contained: doesn't read ``trainer.loss`` (which is rl-only). The - dppo/kl knobs are baked in. + dppo/kl knobs are inlined to match DefaultLossConfig's defaults. """ - adv_tau = 0.0 - teacher_tau = 1.0 - trainer_logprobs = inputs.trainer_logprobs inference_logprobs = inputs.inference_logprobs teacher_logprobs = inputs.teacher_logprobs @@ -201,8 +190,8 @@ def opd_loss_fn(inputs: LossInputs) -> LossOutputs: ) probs_diff = torch.exp(trainer_logprobs) - torch.exp(inference_logprobs) - dppo_invalid_mask_high = probs_diff > _OPD_DPPO_MASK_HIGH - dppo_invalid_mask_low = probs_diff < -_OPD_DPPO_MASK_LOW + dppo_invalid_mask_high = probs_diff > 0.2 + dppo_invalid_mask_low = probs_diff < -0.2 positive_advantages = advantages > 0 negative_advantages = advantages < 0 dppo_invalid_mask = torch.where(positive_advantages, dppo_invalid_mask_high, dppo_invalid_mask_low) @@ -214,11 +203,11 @@ def opd_loss_fn(inputs: LossInputs) -> LossOutputs: keep_mask = loss_mask & ~is_masked teacher_kl = teacher_logprobs - trainer_logprobs - advantages = adv_tau * advantages + teacher_tau * teacher_kl.detach() + advantages = 0.0 * advantages + 1.0 * teacher_kl.detach() pg_loss = keep_mask * advantages * importance_ratio kl_loss = loss_mask * log_importance_ratio**2 - loss = (-pg_loss + _OPD_KL_TAU * kl_loss).sum() + loss = (-pg_loss + 1e-3 * kl_loss).sum() metrics = { "masked_mismatch_kl": _safe_mean(mismatch_kl, loss_mask & is_masked), From 5c0f8bd47bb3ffc9f9eb4d1c7401cc8f2e9e5aad Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Tue, 19 May 2026 03:23:41 +0000 Subject: [PATCH 40/47] docs(loss): trim opd_loss_fn docstring Co-Authored-By: Claude Sonnet 4.6 --- src/prime_rl/trainer/rl/loss.py | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/src/prime_rl/trainer/rl/loss.py b/src/prime_rl/trainer/rl/loss.py index d7c5b655b2..9a9eb25a63 100644 --- a/src/prime_rl/trainer/rl/loss.py +++ b/src/prime_rl/trainer/rl/loss.py @@ -169,12 +169,7 @@ def opd_loss_fn(inputs: LossInputs) -> LossOutputs: """ On-policy distillation loss: the default DPPO+KL math with the tau knobs hardcoded to drop the reward signal and use the teacher KL as the - per-token policy-gradient signal. Equivalent to ``default_loss_fn`` with - ``adv_tau = 0`` and ``teacher_tau = 1``; we inline both here so the two - paths read as distinct losses. - - Self-contained: doesn't read ``trainer.loss`` (which is rl-only). The - dppo/kl knobs are inlined to match DefaultLossConfig's defaults. + per-token policy-gradient signal. """ trainer_logprobs = inputs.trainer_logprobs inference_logprobs = inputs.inference_logprobs From 1365901c07b0142ccaac1541853e19397c016074 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Tue, 19 May 2026 03:52:38 +0000 Subject: [PATCH 41/47] =?UTF-8?q?docs:=20update=20reverse=5Ftext=20README?= =?UTF-8?q?=20=E2=80=94=20add=20lora=20configs?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-Authored-By: Claude Sonnet 4.6 --- configs/reverse_text/README.md | 24 +++++++++++++++--------- 1 file changed, 15 insertions(+), 9 deletions(-) diff --git a/configs/reverse_text/README.md b/configs/reverse_text/README.md index 4981f19e5b..a115792695 100644 --- a/configs/reverse_text/README.md +++ b/configs/reverse_text/README.md @@ -2,16 +2,20 @@ Minimal end-to-end configs for the three training modes against the `reverse-text` env using `PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT` as the student. -| Config | Mode | Teacher | -|---|---|---| -| `debug_rl.toml` | `rl` | none | -| `debug_opd.toml` | `opd` | local vLLM (`Qwen3-0.6B-Reverse-Text-RL`) | -| `debug_sft.toml` | `sft` | local vLLM (`Qwen3-0.6B-Reverse-Text-RL`) | -| `debug_sft_external.toml` | `sft` | PI inference (`openai/gpt-5-mini`) | +| Config | Mode | Teacher | Notes | +|---|---|---|---| +| `debug_rl.toml` | `rl` | none | | +| `debug_opd.toml` | `opd` | local vLLM (`Qwen3-0.6B-Reverse-Text-RL`) | | +| `debug_opd_lora.toml` | `opd` | local vLLM (`Qwen3-0.6B-Reverse-Text-RL`) | trains a LoRA adapter (rank 8) | +| `debug_sft.toml` | `sft` | local vLLM (`Qwen3-0.6B-Reverse-Text-RL`) | | +| `debug_sft_lora.toml` | `sft` | local vLLM (`Qwen3-0.6B-Reverse-Text-RL`) | trains a LoRA adapter (rank 8) | +| `debug_sft_external.toml` | `sft` | PI inference (`openai/gpt-5-mini`) | external OAI endpoint; no local teacher | -The student inference server is auto-launched on GPU 0 at `http://localhost:8000/v1` with `gpu_memory_utilization=0.5`. The teacher (used by `debug_opd.toml` and `debug_sft.toml`) is **not** auto-launched — start it manually on GPU 1. +The student inference server is auto-launched on GPU 0 at `http://localhost:8000/v1` with `gpu_memory_utilization=0.5`. The local teacher (used by everything except `debug_rl.toml` and `debug_sft_external.toml`) is **not** auto-launched — start it manually on GPU 1. -## Start the teacher (only needed for opd/sft) +## Start the local teacher + +Needed for `debug_opd*.toml` and `debug_sft.toml` / `debug_sft_lora.toml`: ```bash CUDA_VISIBLE_DEVICES=1 uv run inference \ @@ -29,11 +33,13 @@ uv run rl @ configs/reverse_text/debug_rl.toml # OPD (needs teacher on port 8001) uv run rl @ configs/reverse_text/debug_opd.toml +uv run rl @ configs/reverse_text/debug_opd_lora.toml # SFT hard distill (needs teacher on port 8001) uv run rl @ configs/reverse_text/debug_sft.toml +uv run rl @ configs/reverse_text/debug_sft_lora.toml -# SFT hard distill from qwen3-30b-a3b-thinking via PI inference +# SFT hard distill from openai/gpt-5-mini via PI inference # (requires PRIME_API_KEY + PRIME_TEAM_ID in env; no local teacher needed) uv run rl @ configs/reverse_text/debug_sft_external.toml ``` From 90e75d94e6f6858b62bb8d1e34fe509d22188d1b Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Tue, 19 May 2026 04:08:25 +0000 Subject: [PATCH 42/47] polish: auto-disable renderer in sft + CHANGELOG + log/shim cleanup - OrchestratorConfig auto-sets use_renderer=False when training_mode='sft' via an after-validator placed before the renderer validators in declaration order, so the user doesn't have to set it. Drop the rejecting branch in validate_training_mode + the explicit use_renderer=false lines in the three sft debug configs. - CHANGELOG: add a "first-class training_mode + batch-driven loss dispatch" entry covering the breaking removals (orchestrator.use_sft_loss, orchestrator.teacher_rollout_model, trainer.loss.type='sft' / SFTLossConfig, trainer.loss.teacher_tau) with migration notes. - _accept_legacy_student_layout: use set(ModelConfig.model_fields) instead of hardcoding the field list, so new ModelConfig fields are picked up automatically. - Trim the orchestrator startup log to just '()'; the prose description was redundant with docs/training_modes.md. Co-Authored-By: Claude Sonnet 4.6 --- CHANGELOG.md | 7 +++++++ configs/reverse_text/debug_sft.toml | 1 - configs/reverse_text/debug_sft_external.toml | 1 - configs/reverse_text/debug_sft_lora.toml | 1 - .../src/prime_rl/configs/orchestrator.py | 17 +++++++++++------ src/prime_rl/orchestrator/orchestrator.py | 8 +------- 6 files changed, 19 insertions(+), 16 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 59b0b126f6..34197e7490 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,13 @@ Documenting **breaking** configuration changes — renamed, removed, or moved fields that require users to update existing configs. +- **First-class `training_mode` + batch-driven loss dispatch** (collection of removals/renames). Loss selection is now driven by `TrainingSample.training_mode` (`rl` / `opd` / `sft`), set under `[orchestrator]`. The trainer is mode-agnostic and dispatches per batch. + - **`orchestrator.use_sft_loss` removed**: replaced by `[orchestrator] training_mode = "sft"`. The orchestrator stamps each sample's `training_mode` and the trainer dispatches to `sft_loss_fn` per batch. + - **`orchestrator.teacher_rollout_model` removed (no alias)**: configs must rename to `[orchestrator.teacher]` (`teacher_model` is also accepted as a back-compat alias for `teacher`). Used in both opd and sft modes; the role is determined by `training_mode`. + - **`orchestrator.model` → `orchestrator.student`**: renamed for symmetry with `orchestrator.teacher`. Legacy `[orchestrator.model]` and a flat `[orchestrator.model.lora]`-style layout are aliased via a back-compat shim, so most existing configs still parse. New code should use `[orchestrator.student.*]`. + - **`trainer.loss.type = "sft"` removed (`SFTLossConfig` deleted)**: zero-field discriminator that is now implied by `training_mode = "sft"`. `trainer.loss` collapsed to a `DefaultLossConfig | CustomLossConfig` union. Existing configs setting `[trainer.loss] type = "sft"` must remove that line; the orchestrator-level `training_mode = "sft"` is enough. + - **`trainer.loss.teacher_tau` removed from `DefaultLossConfig`**: the teacher-KL term is now exclusive to `opd_loss_fn` (which inlines `teacher_tau = 1.0`). `default_loss_fn` only handles reward-driven RL and no longer accepts a `teacher_tau`. Configs that set `teacher_tau` on `[trainer.loss]` will fail validation — switch to `[orchestrator] training_mode = "opd"` (which uses `opd_loss_fn` with the value baked in). + - **`trainer.loss` only applies to rl-mode batches**: opd and sft don't read `trainer.loss`. `opd_loss_fn`'s knobs (dppo_mask_*, kl_tau) are inlined as literals; tweak them by editing the function. `CustomLossConfig` overrides the rl-mode loss only — opd and sft always use their dedicated fns. (2026-05-19) - **`orchestrator.use_token_client` removed**: The server-tokenized TITO path has been deprecated end-to-end. The orchestrator-side config flag (`use_token_client`), the verifiers client_type (`openai_chat_completions_token`), and the inference server's `/v1/chat/completions/tokens` route (along with the `OpenAIServingChatWithTokens` wrapper) are all gone. The orchestrator now picks between renderer-backed TITO (`use_renderer = true`, default) and MITO (`use_renderer = false`, fallback). Existing configs with `use_token_client = true` must migrate to `use_renderer = true` (or `use_renderer = false` for MITO); configs with `use_token_client = false` can simply drop the field. (2026-05-19) - **`orchestrator.advantage.length_penalty` → discriminated sub-config**: The scalar `length_penalty: Literal["tokens","turns"] | None` is replaced by a `LengthPenaltyConfig | None` discriminated on `type`. Token shaping now takes weighted completion + tool-response token costs. Migration: `length_penalty = "tokens"` becomes `[orchestrator.advantage.length_penalty]\ntype = "tokens"` (default weights `completion_weight = 1.0`, `tool_response_weight = 1.0` — total context). `length_penalty = "turns"` becomes `[orchestrator.advantage.length_penalty]\ntype = "turns"`. (2026-05-06) - **`orchestrator.advantage.length_shaping` → `orchestrator.advantage.length_penalty`**: The boolean `length_shaping` flag has been replaced by `length_penalty: Literal["tokens", "turns"] | None` (default: `None`). `length_shaping = true` becomes `length_penalty = "tokens"`; `length_shaping = false` becomes `length_penalty = None`. The new `"turns"` option applies the same correctness-gated efficiency shaping using trajectory turn count instead of completion-token count. (2026-05-01) diff --git a/configs/reverse_text/debug_sft.toml b/configs/reverse_text/debug_sft.toml index 8620fccf7f..af5366b1db 100644 --- a/configs/reverse_text/debug_sft.toml +++ b/configs/reverse_text/debug_sft.toml @@ -19,7 +19,6 @@ name = "debug-sft" training_mode = "sft" batch_size = 128 rollouts_per_example = 4 -use_renderer = false [orchestrator.train.sampling] max_completion_tokens = 128 diff --git a/configs/reverse_text/debug_sft_external.toml b/configs/reverse_text/debug_sft_external.toml index bf6fd3219c..4d74bf1736 100644 --- a/configs/reverse_text/debug_sft_external.toml +++ b/configs/reverse_text/debug_sft_external.toml @@ -18,7 +18,6 @@ name = "debug-sft-external" training_mode = "sft" batch_size = 128 rollouts_per_example = 4 -use_renderer = false [orchestrator.train.sampling] max_completion_tokens = 2048 diff --git a/configs/reverse_text/debug_sft_lora.toml b/configs/reverse_text/debug_sft_lora.toml index 64a048b34f..26a50020f8 100644 --- a/configs/reverse_text/debug_sft_lora.toml +++ b/configs/reverse_text/debug_sft_lora.toml @@ -19,7 +19,6 @@ name = "debug-sft-lora" training_mode = "sft" batch_size = 128 rollouts_per_example = 4 -use_renderer = false [orchestrator.train.sampling] max_completion_tokens = 128 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 328ef8dd2c..e81e5706d9 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py @@ -1165,7 +1165,7 @@ def _accept_legacy_student_layout(cls, data: Any) -> Any: data["model"] = legacy_model # 3. Re-nest flat ModelConfig keys under student.model. - model_only_keys = {"name", "trust_remote_code", "vlm", "lora"} + model_only_keys = set(ModelConfig.model_fields) student = data.get("student") if isinstance(student, dict): flat = {k: student.pop(k) for k in list(student) if k in model_only_keys} @@ -1216,6 +1216,16 @@ def validate_unique_filter_types(self): raise ValueError(f"Duplicate filter types: {types}. Each filter type may only appear once.") return self + @model_validator(mode="after") + def _force_no_renderer_for_sft(self): + """SFT rolls out via the teacher's plain chat-completions endpoint; the + renderer client doesn't apply. Force use_renderer=False so the user + doesn't have to remember to set it. Declared before the renderer + validators below so they see the corrected value.""" + if self.training_mode == "sft": + self.use_renderer = False + return self + @model_validator(mode="after") def validate_training_mode(self): """Enforce training mode invariants that involve only orchestrator fields.""" @@ -1224,11 +1234,6 @@ def validate_training_mode(self): raise ValueError("orchestrator.teacher must not be set when training_mode = 'rl'.") if self.training_mode in ("opd", "sft") and not has_teacher: raise ValueError(f"orchestrator.teacher must be configured when training_mode = '{self.training_mode}'.") - if self.training_mode == "sft" and self.use_renderer: - raise ValueError( - "orchestrator.use_renderer must be false when training_mode = 'sft' " - "(teacher rollout uses the plain OpenAI chat-completions client)." - ) return self @model_validator(mode="after") diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index 807ab374e7..1b0cb4b3ee 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -93,13 +93,7 @@ async def orchestrate(config: OrchestratorConfig): ) intercept_vf_logging(logger="verifiers.serve", level="WARN") # show logs from env clients - # Print start message - mode_descriptions = { - "rl": "student generates rollouts, no teacher", - "opd": "student generates rollouts, teacher judges", - "sft": "teacher generates rollouts, student trains on teacher tokens", - } - logger.info(f"Starting orchestrator in {config.training_mode} mode ({mode_descriptions[config.training_mode]})") + logger.info(f"Starting orchestrator ({config.training_mode})") set_default_executor() event_loop_lag_monitor = EventLoopLagMonitor() From 13c74cb4938b16544ed4d5628cd3015aadf3be9a Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Tue, 19 May 2026 04:25:21 +0000 Subject: [PATCH 43/47] fix(configs): unbreak OPD auto-wiring when num_teacher_gpus is set Add a mode="before" validator on RLConfig that stubs an empty [orchestrator.teacher] block when deployment.num_teacher_gpus is set and the user didn't write one. Without this, OrchestratorConfig.validate_training_mode ("opd requires orchestrator.teacher") fires during nested validation - before auto_setup_teacher_inference can wire the teacher from teacher_inference. Repro that now passes: RLConfig.model_validate({ "deployment": {"num_teacher_gpus": 1}, "trainer": {}, "orchestrator": {"training_mode": "opd"}, "inference": {}, }) -> auto_setup_teacher_inference fills in orchestrator.teacher.client.base_url = ["http://localhost:8001/v1"] and model.name from teacher_inference. Co-Authored-By: Claude Sonnet 4.6 --- .../src/prime_rl/configs/rl.py | 21 +++++++++++++++++++ 1 file changed, 21 insertions(+) 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 9421c3957e..7f8116feaf 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/rl.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/rl.py @@ -367,6 +367,27 @@ class RLConfig(BaseConfig): ### Validate configs (e.g. raise for unsupported (combinations of) configs) + @model_validator(mode="before") + @classmethod + def _stub_orchestrator_teacher_for_auto_setup(cls, data): + """When `deployment.num_teacher_gpus > 0` and the user didn't write an + `[orchestrator.teacher]` block, inject an empty one so + `OrchestratorConfig.validate_training_mode` (which fires during nested + validation) doesn't reject `training_mode = "opd"` for "missing teacher". + `auto_setup_teacher_inference` (after) then fills in client.base_url and + model.name from the auto-launched teacher_inference server.""" + if not isinstance(data, dict): + return data + deployment = data.get("deployment") + if not isinstance(deployment, dict): + return data + if not deployment.get("num_teacher_gpus"): + return data + orch = data.setdefault("orchestrator", {}) + if isinstance(orch, dict) and "teacher" not in orch and "teacher_model" not in orch: + orch["teacher"] = {} + return data + @model_validator(mode="after") def validate_deployment(self): if self.deployment.type == "multi_node": From a96d2aa525bebe8fd9bf258010e9ef16e6731c47 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Tue, 19 May 2026 04:26:53 +0000 Subject: [PATCH 44/47] docs(training_modes): shrink mode-comparison table + add sft/opd teacher note Co-Authored-By: Claude Sonnet 4.6 --- docs/training_modes.md | 31 +++++++++++++------------------ 1 file changed, 13 insertions(+), 18 deletions(-) diff --git a/docs/training_modes.md b/docs/training_modes.md index 2183efd961..a61360e22c 100644 --- a/docs/training_modes.md +++ b/docs/training_modes.md @@ -1,27 +1,22 @@ # Training Modes -prime-rl supports three training modes, selected via `training_mode`: +PRIME-RL supports three training modes through our RL trainer, selected via `training_mode`: -- **`rl`** — standard reinforcement learning from rewards -- **`opd`** — on-policy distillation: RL with an extra KL term toward a teacher's logprobs ([Thinking Machines blog post](https://thinkingmachines.ai/blog/on-policy-distillation/)) -- **`sft`** — hard distillation: supervised fine-tuning on teacher-generated rollouts +- **`rl`** — reinforcement learning: student generates rollouts, no teacher +- **`opd`** — [on-policy distillation](https://thinkingmachines.ai/blog/on-policy-distillation/): students generates rollouts, train to minimize the KL divergence between the student and teacher's logprobs for each token in the rollout +- **`sft`** — supervised fine-tuning on teacher-generated rollouts + +> Note: PRIME-RL also has a dedicated `sft` entrypoint for more traditional supervised fine-tuning from a HF dataset. When using the `sft` training mode on the orchestrator, teacher rollouts are generated on-the-fly and used for training. The mode determines who generates rollouts, what role the teacher plays, and what must be configured. -## Mode comparison - -| | **rl** | **opd** | **sft** | -|---|---|---|---| -| **Student does** | generate rollouts → get trained on them | generate rollouts → get trained on them | serve inference (for evals + weight sync); get trained on teacher's rollouts | -| **Teacher does** | nothing (must be unset) | score student rollouts (token-level logprobs) | generate rollouts | -| **Loss** | reward-based (advantage) | reward + KL to teacher logprobs (`teacher_tau > 0`) | pure NLL on teacher tokens (hard distill) | -| **Student inference** (`[inference]`) | **required** | **required** | **required** | -| **Teacher inference** (`[teacher_inference]`) | forbidden | **required, must be vLLM** | not used (teacher is external) | -| **`[orchestrator.teacher]`** | must be `None` | auto-wired from `[teacher_inference]` | **required** — `client.base_url` + `model.name` of external endpoint | -| **`num_teacher_gpus`** | unset | **required** (`> 0`) | unset (teacher is external) | -| **Teacher endpoint type** | n/a | **local vLLM only** | **any OpenAI-compatible** (PI inference, OpenAI, Anthropic, local vLLM…) | -| **Weight sync (trainer → ?)** | → student inference | → student inference (teacher frozen) | → student inference (teacher frozen) | -| **Evals** | student | student | student | +| Mode | Student | Teacher | +|---|---|---| +| `rl` | required | forbidden | +| `opd` | required | required (local vLLM) | +| `sft` | required | required (any OAI-compatible endpoint) | + +**SFT vs OPD teachers** differ in what the orchestrator asks of them. SFT only calls `/v1/chat/completions` to generate rollouts — any OpenAI-compatible endpoint works (PI inference, OpenAI, Anthropic, a local vLLM). OPD additionally needs token-level logprobs scored over the student's tokens, which today only vLLM's `/inference/v1/generate` with `prompt_logprobs` exposes — so the OPD teacher must be a vLLM server. ## Key implications From 6a0169b4c809af9daadb23957376c107304eb5fb Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Tue, 19 May 2026 04:39:22 +0000 Subject: [PATCH 45/47] chore(configs): move training-mode debug configs to configs/debug/training_modes/ The reverse-text debug configs are about exercising the training modes, not about reverse-text specifically. Move to configs/debug/training_modes/ and drop the debug_ prefix from each filename (the directory name carries it). configs/reverse_text/debug_rl.toml -> configs/debug/training_modes/rl.toml configs/reverse_text/debug_opd.toml -> configs/debug/training_modes/opd.toml configs/reverse_text/debug_opd_lora.toml -> configs/debug/training_modes/opd_lora.toml configs/reverse_text/debug_sft.toml -> configs/debug/training_modes/sft.toml configs/reverse_text/debug_sft_lora.toml -> configs/debug/training_modes/sft_lora.toml configs/reverse_text/debug_sft_external.toml -> configs/debug/training_modes/sft_external.toml configs/reverse_text/README.md -> configs/debug/training_modes/README.md Update path references in each TOML's header comment, the README, and docs/training_modes.md. Co-Authored-By: Claude Sonnet 4.6 --- configs/debug/training_modes/README.md | 47 ++++++++ .../training_modes/opd.toml} | 2 +- .../training_modes/opd_lora.toml} | 2 +- .../training_modes/rl.toml} | 0 .../training_modes/sft.toml} | 2 +- .../training_modes/sft_external.toml} | 2 +- .../training_modes/sft_lora.toml} | 2 +- configs/reverse_text/README.md | 47 -------- docs/training_modes.md | 102 ++---------------- 9 files changed, 58 insertions(+), 148 deletions(-) create mode 100644 configs/debug/training_modes/README.md rename configs/{reverse_text/debug_opd.toml => debug/training_modes/opd.toml} (94%) rename configs/{reverse_text/debug_opd_lora.toml => debug/training_modes/opd_lora.toml} (94%) rename configs/{reverse_text/debug_rl.toml => debug/training_modes/rl.toml} (100%) rename configs/{reverse_text/debug_sft.toml => debug/training_modes/sft.toml} (94%) rename configs/{reverse_text/debug_sft_external.toml => debug/training_modes/sft_external.toml} (93%) rename configs/{reverse_text/debug_sft_lora.toml => debug/training_modes/sft_lora.toml} (94%) delete mode 100644 configs/reverse_text/README.md diff --git a/configs/debug/training_modes/README.md b/configs/debug/training_modes/README.md new file mode 100644 index 0000000000..67c5450947 --- /dev/null +++ b/configs/debug/training_modes/README.md @@ -0,0 +1,47 @@ +# Training Mode — Debug Configs + +Minimal end-to-end configs for the three training modes (`rl` / `opd` / `sft`) against the `reverse-text` env, using `PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT` as the student. + +| Config | Mode | Teacher | Notes | +|---|---|---|---| +| `rl.toml` | `rl` | none | | +| `opd.toml` | `opd` | local vLLM (`Qwen3-0.6B-Reverse-Text-RL`) | | +| `opd_lora.toml` | `opd` | local vLLM (`Qwen3-0.6B-Reverse-Text-RL`) | trains a LoRA adapter (rank 8) | +| `sft.toml` | `sft` | local vLLM (`Qwen3-0.6B-Reverse-Text-RL`) | | +| `sft_lora.toml` | `sft` | local vLLM (`Qwen3-0.6B-Reverse-Text-RL`) | trains a LoRA adapter (rank 8) | +| `sft_external.toml` | `sft` | PI inference (`openai/gpt-5-mini`) | external OAI endpoint; no local teacher | + +The student inference server is auto-launched on GPU 0 at `http://localhost:8000/v1` with `gpu_memory_utilization=0.5`. The local teacher (used by everything except `rl.toml` and `sft_external.toml`) is **not** auto-launched — start it manually on GPU 1. + +## Start the local teacher + +Needed for `opd*.toml` and `sft.toml` / `sft_lora.toml`: + +```bash +CUDA_VISIBLE_DEVICES=1 uv run inference \ + --model.name PrimeIntellect/Qwen3-0.6B-Reverse-Text-RL \ + --server.port 8001 \ + --gpu-memory-utilization 0.5 \ + --model.enforce-eager +``` + +## Run the debug configs + +```bash +# RL (no teacher) +uv run rl @ configs/debug/training_modes/rl.toml + +# OPD (needs teacher on port 8001) +uv run rl @ configs/debug/training_modes/opd.toml +uv run rl @ configs/debug/training_modes/opd_lora.toml + +# SFT hard distill (needs teacher on port 8001) +uv run rl @ configs/debug/training_modes/sft.toml +uv run rl @ configs/debug/training_modes/sft_lora.toml + +# SFT hard distill from openai/gpt-5-mini via PI inference +# (requires PRIME_API_KEY + PRIME_TEAM_ID in env; no local teacher needed) +uv run rl @ configs/debug/training_modes/sft_external.toml +``` + +See [docs/training_modes.md](../../docs/training_modes.md) for what each mode does. diff --git a/configs/reverse_text/debug_opd.toml b/configs/debug/training_modes/opd.toml similarity index 94% rename from configs/reverse_text/debug_opd.toml rename to configs/debug/training_modes/opd.toml index 85c98a9848..b24cf3fe91 100644 --- a/configs/reverse_text/debug_opd.toml +++ b/configs/debug/training_modes/opd.toml @@ -3,7 +3,7 @@ # --model.name PrimeIntellect/Qwen3-0.6B-Reverse-Text-RL \ # --server.port 8001 --gpu-memory-utilization 0.5 --model.enforce-eager # Then: -# uv run rl @ configs/reverse_text/debug_opd.toml +# uv run rl @ configs/debug/training_modes/opd.toml max_steps = 20 seq_len = 2048 diff --git a/configs/reverse_text/debug_opd_lora.toml b/configs/debug/training_modes/opd_lora.toml similarity index 94% rename from configs/reverse_text/debug_opd_lora.toml rename to configs/debug/training_modes/opd_lora.toml index 9150a6b3fe..135f083936 100644 --- a/configs/reverse_text/debug_opd_lora.toml +++ b/configs/debug/training_modes/opd_lora.toml @@ -3,7 +3,7 @@ # --model.name PrimeIntellect/Qwen3-0.6B-Reverse-Text-RL \ # --server.port 8001 --gpu-memory-utilization 0.5 --model.enforce-eager # Then: -# uv run rl @ configs/reverse_text/debug_opd_lora.toml +# uv run rl @ configs/debug/training_modes/opd_lora.toml max_steps = 20 seq_len = 2048 diff --git a/configs/reverse_text/debug_rl.toml b/configs/debug/training_modes/rl.toml similarity index 100% rename from configs/reverse_text/debug_rl.toml rename to configs/debug/training_modes/rl.toml diff --git a/configs/reverse_text/debug_sft.toml b/configs/debug/training_modes/sft.toml similarity index 94% rename from configs/reverse_text/debug_sft.toml rename to configs/debug/training_modes/sft.toml index af5366b1db..3d583e6185 100644 --- a/configs/reverse_text/debug_sft.toml +++ b/configs/debug/training_modes/sft.toml @@ -3,7 +3,7 @@ # --model.name PrimeIntellect/Qwen3-0.6B-Reverse-Text-RL \ # --server.port 8001 --gpu-memory-utilization 0.5 --model.enforce-eager # Then: -# uv run rl @ configs/reverse_text/debug_sft.toml +# uv run rl @ configs/debug/training_modes/sft.toml max_steps = 20 seq_len = 2048 diff --git a/configs/reverse_text/debug_sft_external.toml b/configs/debug/training_modes/sft_external.toml similarity index 93% rename from configs/reverse_text/debug_sft_external.toml rename to configs/debug/training_modes/sft_external.toml index 4d74bf1736..3e42d7c8f3 100644 --- a/configs/reverse_text/debug_sft_external.toml +++ b/configs/debug/training_modes/sft_external.toml @@ -2,7 +2,7 @@ # X-Prime-Team-ID header is auto-injected from $PRIME_TEAM_ID for pinference.ai URLs. # # Run with: -# uv run rl @ configs/reverse_text/debug_sft_external.toml +# uv run rl @ configs/debug/training_modes/sft_external.toml max_steps = 20 seq_len = 2048 diff --git a/configs/reverse_text/debug_sft_lora.toml b/configs/debug/training_modes/sft_lora.toml similarity index 94% rename from configs/reverse_text/debug_sft_lora.toml rename to configs/debug/training_modes/sft_lora.toml index 26a50020f8..560f94a321 100644 --- a/configs/reverse_text/debug_sft_lora.toml +++ b/configs/debug/training_modes/sft_lora.toml @@ -3,7 +3,7 @@ # --model.name PrimeIntellect/Qwen3-0.6B-Reverse-Text-RL \ # --server.port 8001 --gpu-memory-utilization 0.5 --model.enforce-eager # Then: -# uv run rl @ configs/reverse_text/debug_sft_lora.toml +# uv run rl @ configs/debug/training_modes/sft_lora.toml max_steps = 20 seq_len = 2048 diff --git a/configs/reverse_text/README.md b/configs/reverse_text/README.md deleted file mode 100644 index a115792695..0000000000 --- a/configs/reverse_text/README.md +++ /dev/null @@ -1,47 +0,0 @@ -# Reverse Text — Debug Configs - -Minimal end-to-end configs for the three training modes against the `reverse-text` env using `PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT` as the student. - -| Config | Mode | Teacher | Notes | -|---|---|---|---| -| `debug_rl.toml` | `rl` | none | | -| `debug_opd.toml` | `opd` | local vLLM (`Qwen3-0.6B-Reverse-Text-RL`) | | -| `debug_opd_lora.toml` | `opd` | local vLLM (`Qwen3-0.6B-Reverse-Text-RL`) | trains a LoRA adapter (rank 8) | -| `debug_sft.toml` | `sft` | local vLLM (`Qwen3-0.6B-Reverse-Text-RL`) | | -| `debug_sft_lora.toml` | `sft` | local vLLM (`Qwen3-0.6B-Reverse-Text-RL`) | trains a LoRA adapter (rank 8) | -| `debug_sft_external.toml` | `sft` | PI inference (`openai/gpt-5-mini`) | external OAI endpoint; no local teacher | - -The student inference server is auto-launched on GPU 0 at `http://localhost:8000/v1` with `gpu_memory_utilization=0.5`. The local teacher (used by everything except `debug_rl.toml` and `debug_sft_external.toml`) is **not** auto-launched — start it manually on GPU 1. - -## Start the local teacher - -Needed for `debug_opd*.toml` and `debug_sft.toml` / `debug_sft_lora.toml`: - -```bash -CUDA_VISIBLE_DEVICES=1 uv run inference \ - --model.name PrimeIntellect/Qwen3-0.6B-Reverse-Text-RL \ - --server.port 8001 \ - --gpu-memory-utilization 0.5 \ - --model.enforce-eager -``` - -## Run the debug configs - -```bash -# RL (no teacher) -uv run rl @ configs/reverse_text/debug_rl.toml - -# OPD (needs teacher on port 8001) -uv run rl @ configs/reverse_text/debug_opd.toml -uv run rl @ configs/reverse_text/debug_opd_lora.toml - -# SFT hard distill (needs teacher on port 8001) -uv run rl @ configs/reverse_text/debug_sft.toml -uv run rl @ configs/reverse_text/debug_sft_lora.toml - -# SFT hard distill from openai/gpt-5-mini via PI inference -# (requires PRIME_API_KEY + PRIME_TEAM_ID in env; no local teacher needed) -uv run rl @ configs/reverse_text/debug_sft_external.toml -``` - -See [docs/training_modes.md](../../docs/training_modes.md) for what each mode does. diff --git a/docs/training_modes.md b/docs/training_modes.md index a61360e22c..e1787711b1 100644 --- a/docs/training_modes.md +++ b/docs/training_modes.md @@ -18,105 +18,15 @@ The mode determines who generates rollouts, what role the teacher plays, and wha **SFT vs OPD teachers** differ in what the orchestrator asks of them. SFT only calls `/v1/chat/completions` to generate rollouts — any OpenAI-compatible endpoint works (PI inference, OpenAI, Anthropic, a local vLLM). OPD additionally needs token-level logprobs scored over the student's tokens, which today only vLLM's `/inference/v1/generate` with `prompt_logprobs` exposes — so the OPD teacher must be a vLLM server. -## Key implications - -**OPD's teacher cannot be an external API.** `compute_teacher_logprobs` (`src/prime_rl/orchestrator/utils.py`) calls vLLM's `/inference/v1/generate` with `prompt_logprobs=1`. That endpoint is vLLM-specific; PI inference, OpenAI, etc. return 404. For OPD, set `num_teacher_gpus` and let `[teacher_inference]` spin up a local vLLM. - -**SFT's teacher is just chat completions.** It only needs `/v1/chat/completions`. Point `[orchestrator.teacher.client]` at anything OpenAI-compatible. No local GPU needed for the teacher. - -**RL forbids any teacher.** Even a stray `[orchestrator.teacher]` block fails validation. - -**Student model name is always the model being trained.** In SFT this is *not* the rollout-generating model — that's the teacher. The student model field still determines tokenizer, trainer init weights, and what gets saved as checkpoints. - -## Minimal config per mode - -```toml -# rl -training_mode = "rl" -[inference] -``` - -```toml -# opd — auto-launched local teacher -training_mode = "opd" -[deployment] -num_teacher_gpus = 1 # spin up a teacher vLLM -[orchestrator.teacher] # empty block; client + model auto-wired -[trainer.loss] -teacher_tau = 0.5 -[inference] -# Override [teacher_inference.model] to use a different teacher model than the student. -``` - -```toml -# sft — external teacher (PI inference) -training_mode = "sft" -[orchestrator.teacher.client] -base_url = ["https://api.pinference.ai/api/v1"] -[orchestrator.teacher.model] -name = "qwen/qwen3-30b-a3b-instruct-2507" -[inference] -``` - -## OPD details - -### Using an external (already-running) teacher - -Skip `num_teacher_gpus` and point at the existing endpoint. The teacher **must** be a vLLM server (for the `/inference/v1/generate` + `prompt_logprobs` endpoint): - -```toml -training_mode = "opd" -[trainer.loss] -teacher_tau = 0.5 - -[orchestrator.teacher.client] -base_url = ["http://teacher-server:8000/v1"] - -[orchestrator.teacher.model] -name = "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B" -``` - -### Pure distillation (no verification) - -For agentic environments where verification is expensive (code execution, tool use, multi-turn interactions), skip verification and use only the teacher signal: - -```toml -training_mode = "opd" -[deployment] -num_teacher_gpus = 2 - -[trainer.loss] -teacher_tau = 1.0 -adv_tau = 0.0 # disable reward-based learning - -[orchestrator.verification] -enabled = false # skip expensive verification -``` - -The student learns to match the teacher without needing any reward signal. - -### Monitoring - -The `teacher_kl` metric shows the KL divergence from teacher to student. Lower means the student is closer to the teacher. - -## SFT details - -### VLM support - -Image input is supported in SFT mode when the student is a VLM: - -- Prompts can include OpenAI-style image items in `message.content`, e.g. `{"type": "image_url", "image_url": {"url": "data:image/..."}}` -- The orchestrator extracts and preprocesses images from trajectory prompts and attaches `pixel_values` / `image_grid_thw` to training samples -- No teacher token IDs / logprobs are required; reconstruction still happens from messages +### Reference configs -Notes: -- This path currently expects `data:image/...` payloads in message content -- The teacher rollout endpoint must also handle the same multimodal prompts during generation +Debug-scale configs for all three modes (and LoRA variants) live in [`configs/debug/training_modes/`](../configs/debug/training_modes/): -### Reference configs +- `rl.toml` / `opd.toml` / `opd_lora.toml` +- `sft.toml` / `sft_lora.toml` (local vLLM teacher) +- `sft_external.toml` (PI inference teacher) -- `configs/reverse_text/debug_sft.toml` (local vLLM teacher) -- `configs/reverse_text/debug_sft_external.toml` (PI inference teacher) +See [`configs/debug/training_modes/README.md`](../configs/debug/training_modes/README.md) for run commands. ## Parameter reference From bd82ed50eb8478fa91af97364a288d2c9b8f87f8 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Tue, 19 May 2026 04:44:07 +0000 Subject: [PATCH 46/47] refactor(configs): drop auto pinference team-id, rename student-shortcut shim - ClientConfig.auto_setup_pinference_team_header dropped. The hostname-sniff was implicit magic. Configs that need the X-Prime-Team-ID header against PI inference set it explicitly under [...client.headers_from_env] now. Update configs/debug/training_modes/sft_external.toml accordingly and drop the auto-injection note from the file header. - Rename OrchestratorConfig._accept_legacy_student_layout (with leading underscore + "backward-compat" framing) to fold_student_shortcuts. The shim is more than a back-compat: it's the ergonomic path for rl configs too, letting users write [orchestrator.model.name] / [orchestrator.client] rather than [orchestrator.student.model.name] / [orchestrator.student.client]. Docstring reframed accordingly. Teacher has no equivalent shortcut on purpose: rl mode forbids teacher, so the same shortcut routing to two roles would be ambiguous. Co-Authored-By: Claude Sonnet 4.6 --- .../debug/training_modes/sft_external.toml | 5 ++++- .../src/prime_rl/configs/orchestrator.py | 19 +++++++++---------- .../src/prime_rl/configs/shared.py | 6 ------ 3 files changed, 13 insertions(+), 17 deletions(-) diff --git a/configs/debug/training_modes/sft_external.toml b/configs/debug/training_modes/sft_external.toml index 3e42d7c8f3..7fa5a478d9 100644 --- a/configs/debug/training_modes/sft_external.toml +++ b/configs/debug/training_modes/sft_external.toml @@ -1,5 +1,5 @@ # SFT from openai/gpt-5-mini via PI inference. -# X-Prime-Team-ID header is auto-injected from $PRIME_TEAM_ID for pinference.ai URLs. +# Requires PRIME_API_KEY + PRIME_TEAM_ID in the environment. # # Run with: # uv run rl @ configs/debug/training_modes/sft_external.toml @@ -43,6 +43,9 @@ name = "openai/gpt-5-mini" base_url = ["https://api.pinference.ai/api/v1"] api_key_var = "PRIME_API_KEY" +[orchestrator.teacher.client.headers_from_env] +X-Prime-Team-ID = "PRIME_TEAM_ID" + [trainer.optim] lr = 3e-6 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 e81e5706d9..d928e2cacf 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py @@ -1126,20 +1126,19 @@ class OrchestratorConfig(BaseConfig): @model_validator(mode="before") @classmethod - def _accept_legacy_student_layout(cls, data: Any) -> Any: - """Backward-compat shims for the pre-refactor student layout. - - Pre-refactor OrchestratorConfig had top-level `model: ModelConfig` and - `client: ClientConfig` fields. The student/teacher rename consolidated - both under `student: RolloutModelConfig` (with `model` as a legacy alias - for `student`). Re-nest legacy keys so old configs still parse: + def fold_student_shortcuts(cls, data: Any) -> Any: + """Accept top-level ``[orchestrator.model]`` / ``[orchestrator.client]`` + as shorthand for the student sub-config. Useful for ergonomic rl configs + where ``[orchestrator.student.*]`` is overkill, and required for + pre-refactor configs that used the flat layout to keep parsing: - [orchestrator.client.*] -> [orchestrator.student.client.*] - [orchestrator.model.] -> [orchestrator.student.model.] - (where is a ModelConfig field: name, trust_remote_code, vlm, lora) + (where is any ModelConfig field) - Teacher was always nested pre-refactor (teacher_model.model + - teacher_model.client), so we don't touch it. + Teacher must always be configured under [orchestrator.teacher.*] + (no equivalent shortcut), because rl mode forbids a teacher and we + don't want the same shortcut to silently route to two different roles. """ if not isinstance(data, dict): return data diff --git a/packages/prime-rl-configs/src/prime_rl/configs/shared.py b/packages/prime-rl-configs/src/prime_rl/configs/shared.py index 7af041e540..fc1f7187dc 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/shared.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/shared.py @@ -366,12 +366,6 @@ class ClientConfig(BaseConfig): ), ] = None - @model_validator(mode="after") - def auto_setup_pinference_team_header(self): - if any("pinference.ai" in url for url in self.base_url): - self.headers_from_env.setdefault("X-Prime-Team-ID", "PRIME_TEAM_ID") - return self - @property def is_elastic(self) -> bool: """Check if elastic mode is enabled.""" From cfb85455c328ce31898029395f690345834d40ef Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Tue, 19 May 2026 05:18:14 +0000 Subject: [PATCH 47/47] fix(tests): update multi_run CLI args after student rename test_reverse_text_multi_run.start_orchestrator passes --client.base-url and --model.lora.name as CLI flags to `uv run orchestrator`. After the student rename, those tyro args don't exist anymore (they're now --student.client.base-url and --student.model.lora.name). tyro rejects them before the fold_student_shortcuts before-validator can re-nest the dict, so the orchestrator process exits immediately and the test times out waiting for "Step 11" to appear in the log. Rename both CLI args to their new student-scoped form. Co-Authored-By: Claude Sonnet 4.6 --- tests/integration/test_reverse_text_multi_run.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/integration/test_reverse_text_multi_run.py b/tests/integration/test_reverse_text_multi_run.py index 5f57d4f8d8..345f43c6d2 100644 --- a/tests/integration/test_reverse_text_multi_run.py +++ b/tests/integration/test_reverse_text_multi_run.py @@ -205,14 +205,14 @@ def start_orchestrator( run_dir.as_posix(), "--max-steps", str(max_steps), - "--model.lora.name", + "--student.model.lora.name", name, "--wandb.project", wandb_project, "--wandb.name", f"{wandb_name}-{proc_name}", ] - cmd.append("--client.base-url") + cmd.append("--student.client.base-url") cmd.extend(INFERENCE_BASE_URLS) with open(orch_log_dir / "orchestrator.log", "w") as f: