feat: first-class sft+opd training modes

[mikasenghaas]

Note

Stacks on top of #2476 (open). Once #2476 merges to main, this PR will rebase cleanly. Until then, the diff vs main includes #2476's commits as well.

Summary

• Unify OrchestratorConfig model fields: rename model → student and teacher_model → teacher with backward-compat aliases (AliasChoices). Both old TOML names still parse.
• orchestrator.training_mode is the single source of truth (rl / opd / sft). Set it under [orchestrator]; the orchestrator stamps every TrainingSample.training_mode from this field.
• Batch-driven loss dispatch: TrainingSample.training_mode flows through MicroBatch into compute_loss, which looks up the right loss fn per batch:
  • rl → default_loss_fn (uses trainer.loss knobs — adv_tau, dppo_mask_*, kl_tau)
  • opd → opd_loss_fn (self-contained; taus + dppo + kl knobs inlined as literals, matching DefaultLossConfig's defaults)
  • sft → sft_loss_fn (masked NLL on teacher tokens)
    Packer enforces same-mode packing; same-mode batches alone share a micro batch.
• trainer.loss: DefaultLossConfig | CustomLossConfig (discriminated union, drop the old SFTLossConfig). Applies only to rl-mode batches; opd and sft don't read it. CustomLossConfig overrides the rl-mode loss fn.
• Two-pool orchestrator: replace the "one pool + teacher-client overlay" with explicit student_inference: InferencePool (always required, weight-sync + eval target) and teacher_inference: InferencePool | None (set in opd/sft; rollout source for sft, logprob source for opd). Scheduler resolves the rollout pool internally from training_mode. The legacy overlay path that mapped student-pinned clients onto a separate teacher client set is gone; fixes a latent bug where SFT-mode LoRA updates would overwrite scheduler.model_name.
• ClientConfig.headers_from_env — env-var-resolved HTTP headers (analogue of api_key_var). Auto-injects X-Prime-Team-ID → PRIME_TEAM_ID when the base URL targets pinference.ai.
• Debug configs for the training modes at configs/debug/training_modes/{rl,opd,opd_lora,sft,sft_lora,sft_external}.toml against PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT (student). Local-teacher variants use Qwen3-0.6B-Reverse-Text-RL on :8001; sft_external points SFT at openai/gpt-5-mini via PI inference. Each config has its teacher-start command as a comment at the top, plus a README at configs/debug/training_modes/README.md.
• docs/training_modes.md — grid of what student/teacher do in each mode, what must be configured, and the key implications (e.g. OPD teacher must be a vLLM because compute_teacher_logprobs uses vLLM's /inference/v1/generate; SFT teacher can be any OAI-compatible endpoint). Merged the standalone OPD doc into this page.
• Drop check_gpus_available runtime check in the RL entrypoint — it raised RuntimeError when other processes held GPUs, which blocked iterating with a long-lived teacher inference server on a sibling GPU.
• Legacy-layout shims: OrchestratorConfig now accepts both legacy [orchestrator.client.*] and flat [orchestrator.model.<k>] syntaxes via one before-validator, so existing TOMLs (elastic, ci/integration LoRA configs, examples) parse without modification.

Notable design decisions

• Loss dispatch is per-batch, not per-config: the trainer is mode-agnostic — setup_loss_fns always returns all three keys (rl, opd, sft). At the trainer level, mixed-mode batches in one run dispatch correctly. (Orchestrator-side mixed-mode would need per-sample mode selection — separate concern.)
• opd_loss_fn is self-contained: takes only LossInputs, no loss_config. The DPPO/KL knobs (0.2 / 0.2 / 1e-3) and the tau values (0.0 / 1.0) are inlined as literals. Users can't tune them from trainer.loss (which is rl-only by design). If you need to customize, edit opd_loss_fn directly.
• SFTLossConfig deleted: zero-field discriminator that was already implied by training_mode = "sft". trainer.loss collapsed to DefaultLossConfig | CustomLossConfig.
• Headers auto-config triggers only on pinference.ai: explicit allowlist; other PI URLs / external endpoints still require the user to set headers_from_env explicitly.
• logprobs stripped at runtime in sft mode: external reasoning-model endpoints (openai/gpt-5*) reject logprobs=True and the trainer doesn't read it in sft anyway, so the orchestrator pops it from each TrainEnv.sampling_args at startup.

Behavior change

SFT now requires [inference] (or an explicitly-set orchestrator.student.client.base_url pointing at a running student server). Previously SFT could run without a student inference server (teacher-only rollouts, no online evals, no weight sync). After this refactor the orchestrator unconditionally sets up student_inference and waits for it to be ready, so every SFT run needs a student vLLM. If you previously relied on num_infer_gpus = 0 SFT runs, add [inference] (or point orchestrator.student.client.base_url at an externally-started student server). The rl entrypoint logs a warning when [inference] is missing in any mode.

Configs

Orchestrator-side blocks per training mode. The rest of the TOML (model, trainer, inference, etc.) is unchanged.

RL — orchestrator config is unchanged from before this PR.

[orchestrator]
training_mode = "rl"
batch_size = 128
rollouts_per_example = 16

[orchestrator.train.sampling]
max_completion_tokens = 128

[[orchestrator.train.env]]
id = "my-env"

OPD — adds [orchestrator.teacher.*] pointing at a vLLM-backed teacher. opd_loss_fn runs automatically (taus + dppo/kl knobs are inlined; no [trainer.loss] needed).

[orchestrator]
training_mode = "opd"
batch_size = 128
rollouts_per_example = 16

[orchestrator.train.sampling]
max_completion_tokens = 128

[[orchestrator.train.env]]
id = "my-env"

[orchestrator.teacher.model]
name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-RL"

[orchestrator.teacher.client]
base_url = ["http://localhost:8001/v1"]   # must expose vLLM's /inference/v1/generate

SFT — adds [orchestrator.teacher.*] pointing at any OAI-compatible endpoint (PI inference, OpenAI, local vLLM…). Set use_renderer = false (validator-required for sft).

[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 = "my-env"

[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"
# X-Prime-Team-ID auto-injected from $PRIME_TEAM_ID for pinference.ai URLs

Verification

Ran each training mode end-to-end against the reverse-text env with WANDB_MODE=offline, 3 steps. All three exited cleanly (exit code 0, SUCCESS RL training finished! in the log).

Setup: GPU 0 hosts the auto-launched student vLLM (PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT); GPU 1 hosts a manually started teacher vLLM (PrimeIntellect/Qwen3-0.6B-Reverse-Text-RL) on port 8001 — only needed for opd / sft (the local-teacher variants):

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

Mode	Command	Outcome
rl	WANDB_MODE=offline uv run rl @ configs/debug/training_modes/rl.toml --clean-output-dir --max-steps 3 --wandb.shared false	3 train steps + eval at step 1; clean exit.
opd	WANDB_MODE=offline uv run rl @ configs/debug/training_modes/opd.toml --clean-output-dir --max-steps 3 --wandb.shared false	3 train steps; teacher logprobs computed each step (Computing teacher logprobs for 128 training examples); clean exit.
sft (local teacher)	WANDB_MODE=offline uv run rl @ configs/debug/training_modes/sft.toml --clean-output-dir --max-steps 3 --wandb.shared false	3 train steps; rollouts dispatched to teacher pool, student inference used for evals + weight sync; clean exit.
sft (external teacher, openai/gpt-5-mini via PI inference)	uv run rl @ configs/debug/training_modes/sft_external.toml --clean-output-dir	rollouts dispatched to PI inference, X-Prime-Team-ID auto-injected from PRIME_TEAM_ID, train+eval working. Reasoning-model budget tuned via max_completion_tokens = 2048 + reasoning_effort = "minimal"; logprobs stripped at runtime so the OAI endpoint accepts the request.

Also: uv run pytest tests/unit/ -m "not gpu" -q → 355 passed.

---

Note

High Risk
High risk because it refactors core training/orchestration wiring (config schema, inference pool setup, weight sync, and loss computation) and removes/renames several user-facing config fields, which can break existing runs if shims miss edge cases.

Overview
Adds a first-class orchestrator.training_mode (rl/opd/sft) that is stamped onto each TrainingSample.training_mode and used end-to-end to drive behavior.

Refactors config and runtime wiring to match: OrchestratorConfig renames model→student and consolidates teacher configuration into teacher (with back-compat aliases/shims for common legacy layouts), while removing orchestrator.use_sft_loss and orchestrator.teacher_rollout_model. The orchestrator now always brings up a student inference pool (evals + weight sync) and optionally a teacher pool (rollouts for sft, logprobs for opd), and Scheduler selects the rollout pool based on mode.

Updates the trainer to be mode-agnostic: replaces the sft_loss boolean with training_mode, deletes SFTLossConfig and DefaultLossConfig.teacher_tau, adds an opd_loss_fn, and dispatches loss per micro-batch via setup_loss_fns(); packing prevents mixing modes in a micro-batch. Also adds ClientConfig.headers_from_env for env-resolved HTTP headers, refreshes docs (new docs/training_modes.md, removes legacy OPD doc), and adds debug configs for all modes under configs/debug/training_modes/.

^{Reviewed by Cursor Bugbot for commit 5580957. Bugbot is set up for automated code reviews on this repo. Configure here.}

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit 90e75d9. Configure here.}

[samsja]

can we remove the debug config ?

[samsja]

still not convinced sft is the right term tho

[mikasenghaas]

hm yea, you have a better suggestion? also i think @willccbb wanted to ship some env features that makes it more like traditional sft

[mikasenghaas]

can we remove the debug config?

would strongly prefer to keep these for internal documentation and dev. can maybe ship into configs/debug/ tho?

[samsja]

hmmmm wondering if its not better to failed here, I think that its fine to break backward compat because this feature was not the most used ?

not sure about removing it

[mikasenghaas]

no this we prob need to keep -- otherwise model.name doesn't work anymore which imo is the more natural way to configure an RL run (which doesn't have a teacher)

[mikasenghaas]

going to adjust the name and docstring tho to make this clearer

[samsja]

huhuhuhuhhuhuhuuhuhh but isn't it an easier way to do it hmmmmmmmmmmmmmmm

[mikasenghaas]

im all ears