F1 — Systems Framing: Dataset-Gen+SFT, RL, or BOTH?
Facet question: Is this a dataset-generation+SFT system, an RL system, or BOTH?
Committed answer: BOTH — and decisively so — structured as TWO LOOPS at two timescales, not two phases. The repo already physically contains both halves. The OUTER (slow) loop is dataset/curriculum construction; the INNER (fast) loop is RL. They feed each other continuously: the inner loop's improved student generates the outer loop's next seed traces, and the inner loop's learned deliberation-confidence becomes the outer loop's branch gate. This is the report §5 verdict ("two loops at different timescales, not two phases") made concrete against the code.
This is not an architectural opinion — it is forced by what the modules are:
| Repo component | What it is, mechanically | Loop |
|---|---|---|
ingestion/claude_code.py |
Claude Code JSONL → TraceState (one node per assistant turn; tool_error flag; strip_thinking=False) — produces seed traces |
OUTER (dataset) |
teacher_replay.py |
N-teacher OpenRouter replay + extract_dpo_pairs — OFFLINE dataset gen, flat depth-1 stars hung off a frozen trace |
OUTER (dataset) |
datagen/substrates.py (SweBenchAdapter) |
SWE-bench tuple → FeatureDeletionTask (revert gold patch = synthesize broken repo) — task SYNTHESIS |
OUTER (curriculum) |
datagen/env.py FeatureDeletionEnv.step()/_grade() |
execution oracle: runs action in sandbox, returns observation; _grade() = masked FAIL_TO_PASS pass-fraction — this is the RL env reward kernel |
OUTER (fitness) AND consumed by INNER (reward_fn) |
datagen/curriculum.py DifficultyCurriculum |
p̂(1−p̂) frontier weighting, retire >0.95, quarantine <0.02 — selection/sampling for the next round | OUTER (selection) |
| The multi-model MCTS tree controller (to build) | recursion: apply each model's action via env.step(), branch again — depth-1 stars → tree |
OUTER (the core delta) |
trainer/composer_trainer.py ComposerReplicationTrainer |
a real trl.GRPOTrainer subclass; total = grpo + α·sdpo + β·trace_replay_dpo — this is RL |
INNER (fast RL) |
datagen/env.py FeatureDeletionEnv.reward_fn |
TRL RewardFunc adapter (prompts, completions → list[float]) — the env wearing its RL face |
INNER (RL reward) |
loss.py compose_loss |
TRL-free 3-channel harness; Channel 1 = LM cross-entropy on response tokens — this is the SFT-able face (BC limit GRPO converges to) | INNER, also the SFT-first floor |
safety/holdout.py + safety/kill_switch.py (HeldOutGuard/HeldoutSplit), wired into the trainer |
run-level collapse/reward-hacking tripwire on proxy-minus-realeval gap — an RL-run safeguard | INNER (safety) |
So: FeatureDeletionEnv.reward_fn is an RL env. teacher_replay is offline
dataset gen. ComposerReplicationTrainer(trl.GRPOTrainer) is RL.
compose_loss's lm_ce channel is an SFT harness. HeldOutGuard is an RL-run
safeguard. The system is both, by construction.
The SFT-first competence floor, THEN RL (mirrors Cursor's CPT+SFT→RL)
docs/COMPOSER_RECIPE_MAPPING.md confirms Cursor's ordering is Continued
Pretraining → SFT → RL, and that the repo deliberately skips CPT (starts
from an already-code-tuned base, e.g. Qwen3-Coder-7B / Qwen3-Coder-30B-A3B).
That leaves a two-step ordering this facet commits to:
- SFT-first (competence floor). Take the OUTER loop's clean winning
trajectories (oracle-clean
_grade()passes only — Gate 1) and run standard SFT. The carrier already exists:compose_losswithalpha_sdpo=0, beta_replay=0reduces to_lm_response_ce— next-token cross-entropy masked to assistant-response tokens. This is the "GRPO converges to BC under deterministic rewards" limit stated inloss.py. It establishes a floor so GRPO has a non-degenerate starting policy. (For a clean separation, an SFTTrainerover the same corpus is equivalent; the point is the corpus — winning leaves — and the masking — response tokens only.) - THEN RL.
ComposerReplicationTrainerruns GRPO/Dr.GRPO (make_po_config) on theFeatureDeletionEnv.reward_fnexecution oracle, with the optional SDPO (α) and trace-replay-DPO (β) channels, plus the contested world-model next-state head as a second SDPO mode (parameter-isolated; report §2/§4).
This is exactly the report §5 line: "SFT-first establishes a competence floor on
clean winning trajectories before RL — mirroring Cursor's CPT+SFT→RL ordering and
the repo's own outer (datagen/teacher_replay) / inner
(ComposerReplicationTrainer) split."
The single diagram-able data flow
┌──────────────────────────────────────────────────────────────┐
│ OUTER LOOP (slow: hours→days; bursty, Spot-friendly, EKS) │
│ │
raw base model ───►│ ingestion/claude_code.py (seed traces, TraceState) │
+ Claude traces │ │ │
│ ▼ │
│ multi-model MCTS tree controller ──► N models branch │
│ (divergence-gated; teacher_replay generalized flat→tree) │
│ │ │
│ ▼ apply each action │
│ FeatureDeletionEnv.step() ─► sandbox exec ─► new state │
│ │ │
│ ▼ leaf grade │
│ FeatureDeletionEnv._grade() (masked FAIL_TO_PASS pass-frac) │
│ │ + DifficultyCurriculum.update (p(1-p) selection) │
│ ▼ HARVEST + TYPE the divergence │
└────────────┼───────────────────────────────────────────────────┘
▼
┌──────────────────── S3 ────────────────────────────────────────────────┐
│ s3://<bucket>/runs/<run_id>/ │
│ sft_corpus/ ← clean WINNING trajectories (SFT-first) │
│ dpo_pairs/ ← near-miss (chosen=sibling winner, rejected=loser)│
│ rl_task_pool/ ← FeatureDeletionTask registry + curriculum priors │
│ wm_tuples/ ← (state, action, next_state, outcome) — ALL branches│
│ divergence_pairs/ ← divergence-annotated nodes (where siblings forked)│
│ holdout/ ← DISJOINT held-out eval anchor (never fed back) │
│ diloco_rendezvous/ ← round_{NNNNNN}/rank_{RRRR}.pt (ObjectStoreAllReduce)│
└────────────┬──────────────────────────────────────────────────────────────┘
▼
┌───────── SFT JOB ───────────────┐ (compose_loss lm_ce / SFT Trainer)
│ sft_corpus → competence-floor ckpt│ ──► seeds the RL init policy
└────────────┬──────────────────────┘
▼
┌──────────────────────────────────────────────────────────────┐
│ INNER LOOP (fast: minutes→steps; resilience-bound, GPU) │
│ ComposerReplicationTrainer (trl.GRPOTrainer subclass) │
│ total = grpo(reward_fn on _grade) │
│ + α·sdpo (hint-distill, JSD on aligned post-hint tokens) │
│ + β·trace_replay_dpo (dpo_pairs) │
│ + [world-model next-state head — 2nd SDPO mode, gated] │
│ HeldOutGuard tripwire on proxy−realeval gap (kill-switch) │
│ DiLoCo outer-sync every ~500-1000 steps via S3 diloco_rendezvous │
└────────────┬──────────────────────────────────────────────────┘
▼
improved student model
│
┌────────────────────────┴────────────────────────────────────────────────┐
│ FEEDBACK (why loops, not phases): │
│ 1. improved student generates the next round's seed traces (back to OUTER)│
│ 2. its learned deliberation-confidence becomes the next round's branch │
│ gate (the §3 bootstrap: cross-model disagreement early → learned │
│ deliberation-confidence later — same lever, two levels) │
└────────────────────────────────────────────────────────────────────────────┘
What goes in S3 between the loops (crisp)
The boundary between OUTER and INNER is S3 — and on AWS S3 is the DiLoCo
rendezvous backend with zero new code (the ObjectStoreAllReduce fsspec path,
round_{NNNNNN}/rank_{RRRR}.pt). The same bucket carries the dataset hand-off.
Live target bucket in this account/region:
s3://amazon-sagemaker-386931836011-us-west-2-7597bf4d9a3d/ (us-west-2;
a sagemaker-dynamo-on-eks-hyperpod-* bucket also already exists, evidence
HyperPod-on-EKS is provisioned here).
| S3 prefix | Contents | Producer (OUTER) | Consumer (INNER) |
|---|---|---|---|
sft_corpus/ |
SFT corpus — clean winning trajectories (oracle-clean _grade() passes), masked to assistant-response tokens; JSONL |
tree controller + _grade() Gate 1 |
SFT job → compose_loss lm_ce (or SFT Trainer) |
dpo_pairs/ |
DPO pairs — {chosen=sibling-winner-or-teacher-consensus, rejected=student/losing-branch, state_messages} (the DPOPair schema from teacher_replay.py) |
extract_dpo_pairs generalized to sibling winners |
Channel 3 (β·trace_replay_dpo) |
rl_task_pool/ |
RL task pool — FeatureDeletionTask registry (repo, broken_image, fail_to_pass, pass_to_pass, deleted_symbols, test_command) + DifficultyCurriculum priors |
SweBenchAdapter + curriculum |
FeatureDeletionEnv registry → reward_fn |
divergence_pairs/ |
Divergence-annotated pairs — nodes where sibling next-action distributions disagreed (the SDPO privileged-info conditioning variable + which sibling subtree separated first) | tree controller (records pre-expansion divergence) | Channel 2 SDPO ctx_teacher splice / SiblingBootstrapGenerator |
wm_tuples/ |
World-model next-state tuples — (state, action, next_state, outcome) from ALL branches incl. failures (CWM "train-on-all" target; the safe home for failed-branch signal) |
every env.step() observation, every leaf grade |
world-model next-state head (2nd SDPO mode), gated |
holdout/ |
Disjoint held-out eval anchor — HeldoutSplit partition; NEVER fed back to the generator |
HeldoutSplit.split(seed=…) once |
HeldOutGuard.heldout_eval_fn() |
diloco_rendezvous/ |
DiLoCo pseudo-gradient exchange, round_{NNNNNN}/rank_{RRRR}.pt |
inner replicas | inner replicas (ObjectStoreAllReduce) |
Note the keystone (report §4/§2): a failed branch is poison for the policy
gradient but gold for the world model — so a loser goes to wm_tuples/
(safe, no policy penalty) and optionally to dpo_pairs/ as a contrastive
rejected against a sibling winner (never as raw negative gradient), but never
to sft_corpus/. That is "type the signal and route it" realized as S3 prefixes.
AWS-native realization (concrete, this account)
- OUTER loop on EKS (single control plane): Argo Workflows DAG, one node =
one divergence-gated branch; vLLM RayService pods for open-weight model
families + API-egress pods for hosted teachers; gVisor (
runscRuntimeClass) sandbox pods runningFeatureDeletionEnv._grade(). Writes all six dataset prefixes to S3 via IRSA. - SFT job between loops: a SageMaker Training Job (
Filemode for thesft_corpus/< 100 GB;FastFileif it grows past ~50–100 GB) OR an EKS GPU pod; readssft_corpus/, writes a competence-floor checkpoint tos3://…/runs/<id>/ckpt_sft/. - INNER loop:
ComposerReplicationTraineron a Karpenter p5/g6e NodePool, swappable to a HyperPod-attached node-group (1:1 EKS↔HyperPod mapping, already provisioned per thesagemaker-dynamo-on-eks-hyperpod-*bucket) for resilience-bound long runs. DiLoCo replicas rendezvous only through S3 — no cross-job NCCL — so a straggler blocks at the poll loop (timeout_s=1800) instead of deadlocking a gang.SageMakerExecutoris the bursty fallback (N single-instance Training Jobs,EnableNetworkIsolation=Falseso the container can S3 PUT/GET the rendezvous).
The DiLoCo math, MockManager, make_diloco_outer_loop, the trainer, the loss,
and the env are untouched across all of this — the ServerlessExecutor
Protocol + ObjectStoreAllReduce are the entire portability contract.
Repo delta to realize the framing
composer_replication/datagen/tree_controller.py(NEW, ~250-350 LOC) — the recursion: apply each model's candidate action viaFeatureDeletionEnv.step(), branch again, grade leaves, emit the six S3 prefixes. The core OUTER delta.composer_replication/pipeline/s3_layout.py(NEW, ~80 LOC) — typed writers forsft_corpus/ | dpo_pairs/ | rl_task_pool/ | divergence_pairs/ | wm_tuples/ | holdout/; the OUTER→INNER contract in one place.composer_replication/pipeline/sft_floor.py(NEW, ~60 LOC) — SFT-first driver: readsft_corpus/, runcompose_loss(alpha_sdpo=0, beta_replay=0) or an SFTTrainer, writeckpt_sft/. Wraps existing_lm_response_ce.composer_replication/trainer/composer_trainer.py(EDIT, ~40 LOC) — add the world-model next-state head as a 2nd SDPO mode (parameter-isolated adapter +<deliberate>token), readingwm_tuples/; gated OFF by default.composer_replication/diloco/serverless/eks.py(EDIT/FINISH) —EKSExecutorIndexed Job mapping for the inner loop (a few hundred LOC; sibling ofModalSpawnExecutor).[serverless]extra: adds3fs/boto3/kubernetes(the documented dep gap).
Open questions
- Where exactly does SFT-first run — a dedicated SageMaker Training Job, or an EKS pod sharing the inner NodePool? (Cost/latency tradeoff; corpus size gates File vs FastFile.)
- Should the SFT-floor checkpoint be re-derived each outer generation (full SFT→RL re-warm) or only once at gen-0 (RL-only thereafter)? The flywheel feedback suggests gen-0-only, with RL carrying subsequent gains.
- Is the world-model head trained inside the inner GRPO trainer (shared step) or
as a separate offline pass over
wm_tuples/? Report §2 favors parameter-isolation; a separate pass is the strongest isolation.
Citations
- research/notes/final_report_socratic-mcts-swe-worldmodel-8f6dea.md §5 (two loops), §2 (world-model head as 2nd SDPO mode), §4 (type-the-signal routing), §6 (reuse/build table)
- composer_replication/datagen/env.py:63-94 (
step/_grade/reward_fn) - composer_replication/teacher_replay.py:162-262 (
replay_trace,extract_dpo_pairs,DPOPair) - composer_replication/trainer/composer_trainer.py:54-178 (3-channel
_compute_loss), :184-251 (HeldOutGuardwiring) - composer_replication/loss.py:71-261, :277-304 (
compose_loss,_lm_response_ce) - composer_replication/safety/holdout.py (
HeldoutSplitdisjointness) - composer_replication/diloco/serverless/{executor.py,allreduce.py,sagemaker.py}
(Protocol +
ObjectStoreAllReduce+ S3 rendezvous) - composer_replication/ingestion/claude_code.py:6-21 (one-node-per-turn,
strip_thinking) - docs/COMPOSER_RECIPE_MAPPING.md:48-137 (CPT+SFT→RL ordering; repo skips CPT)
- Live:
aws s3 ls→ amazon-sagemaker-386931836011-us-west-2-7597bf4d9a3d, sagemaker-dynamo-on-eks-hyperpod-* (us-west-2, acct 386931836011)