research/design-F1-systems-framing.md

# 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:

1. **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_loss` with `alpha_sdpo=0,
   beta_replay=0` reduces 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 in `loss.py`. It establishes a floor so
   GRPO has a non-degenerate starting policy. (For a clean separation, an SFT
   `Trainer` over the same corpus is equivalent; the point is the *corpus* —
   winning leaves — and the *masking* — response tokens only.)
2. **THEN RL.** `ComposerReplicationTrainer` runs GRPO/Dr.GRPO (`make_po_config`)
   on the `FeatureDeletionEnv.reward_fn` execution 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 (`runsc` RuntimeClass)
  sandbox pods running `FeatureDeletionEnv._grade()`. Writes all six dataset
  prefixes to S3 via IRSA.
- **SFT job** between loops: a SageMaker Training Job (`File` mode for the
  `sft_corpus/` < 100 GB; `FastFile` if it grows past ~50–100 GB) OR an EKS GPU
  pod; reads `sft_corpus/`, writes a competence-floor checkpoint to
  `s3://…/runs/<id>/ckpt_sft/`.
- **INNER loop**: `ComposerReplicationTrainer` on a Karpenter p5/g6e NodePool,
  swappable to a HyperPod-attached node-group (1:1 EKS↔HyperPod mapping, already
  provisioned per the `sagemaker-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. `SageMakerExecutor` is the bursty fallback
  (N single-instance Training Jobs, `EnableNetworkIsolation=False` so 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

1. `composer_replication/datagen/tree_controller.py` (NEW, ~250-350 LOC) — the
   recursion: apply each model's candidate action via `FeatureDeletionEnv.step()`,
   branch again, grade leaves, emit the six S3 prefixes. The core OUTER delta.
2. `composer_replication/pipeline/s3_layout.py` (NEW, ~80 LOC) — typed writers
   for `sft_corpus/ | dpo_pairs/ | rl_task_pool/ | divergence_pairs/ | wm_tuples/
   | holdout/`; the OUTER→INNER contract in one place.
3. `composer_replication/pipeline/sft_floor.py` (NEW, ~60 LOC) — SFT-first
   driver: read `sft_corpus/`, run `compose_loss` (`alpha_sdpo=0, beta_replay=0`)
   or an SFT `Trainer`, write `ckpt_sft/`. Wraps existing `_lm_response_ce`.
4. `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), reading `wm_tuples/`; gated OFF by default.
5. `composer_replication/diloco/serverless/eks.py` (EDIT/FINISH) — `EKSExecutor`
   Indexed Job mapping for the inner loop (a few hundred LOC; sibling of
   `ModalSpawnExecutor`).
6. `[serverless]` extra: add `s3fs`/`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 (`HeldOutGuard` wiring)
- composer_replication/loss.py:71-261, :277-304 (`compose_loss`, `_lm_response_ce`)
- composer_replication/safety/holdout.py (`HeldoutSplit` disjointness)
- 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)