feat(orchestrator): add static dataset SFT

[tim0120]

Summary

Rebased onto main (post-#2746) and adds static-dataset SFT as a sampling.source for the existing sft algorithm.

sft remains one CE-only algorithm. The source chooses the flavor:

• a frozen teacher source (an inline frozen sampling.source) samples teacher rollouts through the normal Verifiers env path and trains CE on those tokens
• a dataset source (type = "dataset") replays stored supervised rows locally — no Verifiers env server, model client, sampling logprobs, or off-policy aging

Config surface

[orchestrator.algo]
type = "sft"

[orchestrator.algo.sampling.source]
type = "dataset"
name = "org/my-sft-dataset"
split = "train"
# optional: subset, max_examples, messages_column, prompt_column, completion_column, tools_column

[[orchestrator.train.env]]
id = "static-sft"
name = "static-sft"
group_size = 1

Rows carry either a messages column or prompt + completion columns, with optional tools.

Implementation notes

• Adds StaticDatasetConfig to SamplingConfig.source and Sampler.source_kind (policy / frozen_model / dataset) for the liveness plumbing.
• Replay is v1-native: a dataset env loads its rows at start() (no env server; num_tasks = dataset length, so TrainSource weighted sampling works unchanged) and builds a renderer from the policy tokenizer — the same pattern as opsd's hint renderer. row_to_rollout renders the row's messages once and constructs a typed Rollout (vf.Trace) whose MessageNodes carry the exact rendered token sequence; the renderer's sampled_mask marks assistant emission tokens as the CE targets, so trace_to_samples and the trainer hot path are untouched.
• The dispatcher schedules dataset groups as local replay tasks (task_idx-based, same inflight accounting); non-policy sources skip off-policy aging.
• Validation: sft rejects the policy as its source; non-policy sources can only feed CE algorithms; dataset sources require group_size = 1 (rows are fixed targets); an empty split and a renderer without sampled-token attribution (the default Jinja renderer) fail loudly at startup.
• SFTDistillAlgorithm → SFTAlgorithm (the class now covers both flavors; it declares CE routing and no scoring hooks either way).
• Fixes Rollout.is_trainable to count stamped nonzero ce/ref_kl weight streams as a training signal (cached after stamping) — previously every sft/opd/opsd step displayed Trainable 0/N (0.0%) and warned about task difficulty, since those algorithms ship advantages=None by design. Display/metrics only; what ships to the trainer is unchanged.

Review hardening

A deep review pass surfaced and fixed:

• StaticDatasetConfig.type is required: the undiscriminated SamplingSource union used to silently validate a frozen-model table missing base_url as a dataset, swallowing the real error.
• The gibberish/repetition filters exempt rollouts that carry no sampling logprobs — frozen-teacher and static-replay rollouts zero-fill branch.logprobs, which the filters previously misread as prob-1.0 confidence on every token (bogus detection metrics; run-killing with enforce = true).
• The replay renderer is built around the orchestrator's tokenizer (honoring [orchestrator.tokenizer]) instead of re-loading from the served model name; renders are serialized behind a lock (HF fast tokenizers aren't safe for concurrent single-instance use), and an unattributing renderer (e.g. default) fails at env start with a clear message instead of per-row.
• Tool-calling rows normalize to the OAI wire shape (struct/dict arguments → JSON string, missing ids filled), and tool definitions accept the verifiers tool_defs column and bare shapes — mirroring the offline SFT trainer's normalization.
• Dataset rows stay Arrow-backed and are indexed per task_idx (no up-front materialization of the whole split).
• Dataset-sourced envs default group_size to 1 instead of failing on the inherited global value, so mixed RL + replay configs validate; only an explicit conflicting value raises.
• JSON-looking text completions (e.g. a bare "42") stay content — only message-shaped JSON cells are decoded.
• Local replays take an inflight permit but no tasks_per_minute rate-limiter slot (that budget paces outbound inference requests).

Known limitations (by design, fail loud): rows whose tool_calls can't be coerced to the OAI shape error at replay time and surface as errored rollouts; static-replay rollouts carry no rewards, so train/agg reward panels include their 0.0 in mixed runs.

Validation

10-step smoke runs on an 8×H100 pod (Qwen3-0.6B, qwen3 renderer, willcb/R1-reverse-wikipedia-paragraphs-v1-1000 capped at 512 rows, batch_size 32), re-run after the hardening batch: dataset env starts with no env server (num_tasks=512), CE loss 0.46 → 0.40 over 10 steps, 0 errors/truncations, Trainable 32/32 (100.0%), clean exit. Inspected a rendered rollout end-to-end: system/user nodes fully masked, assistant node trains on exactly its emission span (746/750 tokens — the 4 masked tokens are the <|im_start|>assistant\n scaffold), first CE target is the <think> open, last is </answer><|im_end|>. Also verified on the pod: a name-only (no base_url) source now fails with the real diagnosis; a mixed grpo+dataset config with global group_size = 16 resolves to 16/1 per env and an explicit conflicting value still raises; a tool-calling row with struct arguments, missing call ids, and verifiers-shaped tools renders into a correctly masked multi-turn rollout; scalar text completions ("42") train as content.

🤖 Generated with Claude Code

---

Note

Medium Risk
Touches rollout scheduling, SFT config validation, and training-signal detection; mistakes could mis-route batches or misreport trainability, but scope is mostly additive on the CE path with loud startup failures for bad renderers/datasets.

Overview
Static dataset SFT extends sft so orchestrator.algo.sampling.source can be a Hugging Face dataset (type = "dataset") in addition to a frozen teacher endpoint. Rows replay locally into tokenized Rollouts (no Verifiers env server, model client, sampling logprobs, or off-policy aging); CE targets come from the renderer’s sampled_mask on assistant tokens.

Config adds StaticDatasetConfig, SamplingSource, and Sampler.source_kind (policy / frozen_model / dataset). Dataset envs default group_size to 1 (explicit group_size != 1 errors); StaticDatasetConfig.type is required so a frozen model missing base_url does not validate as a dataset. SFTDistillAlgorithm is renamed to SFTAlgorithm.

The dispatcher routes dataset train groups to TrainEnv.run_static_rollout (inflight permit, no tasks_per_minute slot). static_sft.py loads HF rows (messages or prompt/completion, optional tools) with tokenizer/renderer alignment and tool-call normalization. Gibberish/repetition filters skip rollouts without sampling logprobs. Rollout.is_trainable treats stamped nonzero ce/ref_kl weights as trainable and invalidates cache after stamping.

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

[Copilot]

Pull request overview

Adds “static dataset SFT” as a new rollout source for the existing sft algorithm, allowing SFT training to replay supervised rows locally (no Verifiers env server/client/model sampling), while keeping SFT CE-only and enforcing config compatibility constraints (e.g., dataset SFT requires group_size = 1).

Changes:

• Introduces StaticDatasetConfig as a sampling.source option and plumbs Sampler.source_kind through dispatcher/env logic to support local static replay.
• Adds static_sft.py for dataset row loading + rollout construction, plus token-usage backfilling from finalized samples.
• Updates validation/docs/skills and extends unit tests for dataset SFT acceptance/rejection paths.

Reviewed changes

Copilot reviewed 18 out of 18 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
tests/unit/test_configs.py	Adds config validation test for dataset SFT requiring group_size = 1.
tests/unit/orchestrator/test_static_sft.py	New unit tests for static SFT rollout shaping (messages, prompt/completion, tool_defs).
tests/unit/orchestrator/test_algorithms.py	Updates algorithm tests for dataset source acceptance and CE-only constraints.
src/prime_rl/orchestrator/types.py	Updates dispatcher state docstrings to reflect per-example scheduling semantics.
src/prime_rl/orchestrator/trajectories.py	Adjusts SFT assertion wording and improves tool-def normalization for OAI-format passthrough.
src/prime_rl/orchestrator/train_sink.py	Adds token-usage derivation/backfill from finalized samples during rollout processing.
src/prime_rl/orchestrator/static_sft.py	Adds dataset loading + rollout construction for static/dataset-sourced SFT.
src/prime_rl/orchestrator/sampler.py	Exposes Sampler.source_kind and makes sampling args conditional on policy vs non-policy sources.
src/prime_rl/orchestrator/envs.py	Allows dataset-sourced train envs to skip env startup and run local static rollouts.
src/prime_rl/orchestrator/dispatcher.py	Routes dataset-sourced train rollouts to local replay scheduling (no model/env).
src/prime_rl/orchestrator/algo/sft.py	Renames/clarifies SFT algorithm class and docstring around source-provided targets.
src/prime_rl/orchestrator/algo/init.py	Wires sft type to SFTAlgorithm symbol rename.
skills/configs/SKILL.md	Updates config skill to the new top-level [orchestrator.algo] surface and static SFT notes.
packages/prime-rl-configs/src/prime_rl/configs/rl.py	Updates inference client auto-setup docstring to include dataset/frozen-teacher SFT.
packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py	Adds validation enforcing dataset SFT uses group_size = 1.
packages/prime-rl-configs/src/prime_rl/configs/algorithm.py	Adds StaticDatasetConfig, expands sampling source union, and updates validation/error text.
docs/training.md	Updates training docs table to the new algo config surface and describes dataset SFT option.
docs/algorithms.md	Documents dataset vs frozen-teacher SFT sources and updates algorithm abstraction text accordingly.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[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.}

Want reviews to match your repository better? Bugbot Learning can learn team-specific rules from PR activity. A team admin can enable Learning in the Cursor dashboard.

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