docs/USER_GUIDE.md

# Composer Replication Framework — User Guide

A zero-to-training walkthrough for the open replication of Cursor Composer 2.5.
Pace: an ML engineer who knows GRPO/DPO at a textbook level but has never
opened this repo. Every step references real code, and every kwarg name
listed below has been imported and verified against
`composer_replication/` source.

---

## 1. What is this framework?

A pure-PyTorch replication of the **3-channel composer loss** that powers
agentic-coding model training. One model, one optimizer, three additive
loss terms — composed every step:

```
                      ┌────────────────────────────────────────────┐
                      │            compose_loss(model, batch)       │
                      └────────────────────────────────────────────┘
                                          │
        ┌─────────────────────────────────┼─────────────────────────────────┐
        ▼                                 ▼                                 ▼
┌───────────────────┐          ┌──────────────────────┐         ┌──────────────────────┐
│  Channel 1 (RL)   │          │  Channel 2 (SDPO)    │         │  Channel 3 (replay)  │
│  GRPO            │          │  hint-distillation   │         │  multi-teacher DPO   │
│  → lm_ce stub in │          │  generalized JSD     │         │  on (chosen,         │
│  verification    │          │  student vs teacher  │         │  rejected) pairs     │
│  harness         │          │  (hint-conditioned)  │         │  from N teachers     │
└─────────┬─────────┘          └──────────┬───────────┘         └──────────┬───────────┘
          │  weight = 1 (always on)       │ alpha_sdpo            beta_replay │
          └────────────────┬──────────────┴─────────────────┬──────────────────┘
                           ▼                                ▼
                   total = lm_ce + α·sdpo_jsd + β·trace_replay_dpo
                   (channel auto-disables if its weight=0 OR its inputs are missing)
```

Two API surfaces, on purpose:

- **Verification harness** — `compose_loss(model, batch, ...)` is a free
  function (channel 1 = LM cross-entropy, the GRPO limit under deterministic
  rewards). Use it for CPU smokes, unit tests, and gradient-flow debugging.
- **Production trainer** — `ComposerReplicationTrainer` is a `trl.GRPOTrainer`
  subclass that overrides `_compute_loss` with the same 3 channels on top of
  TRL's real reward + advantage machinery.

The verification harness is what you'll use for sections 2–6; the production
trainer (and its alternates VeRL/PRIME-RL/Monarch) is section 8.

Source of truth: `composer_replication/loss.py` for `compose_loss`,
`composer_replication/trainer/composer_trainer.py` for the trainer subclass.

---

## 2. Install — which extras to pick

Always start with the core install:

```bash
git clone https://huggingface.co/Codeseys/composer-replication-framework
cd composer-replication-framework
pip install -e .
```

That gets you `torch>=2.0` + `transformers>=4.46` and is enough for the
verification harness on CPU (sections 3, 5, 6).

The seven optional extras are declared in `pyproject.toml` `[project.optional-dependencies]`:

```
                              Do you need …
                                    │
        ┌──────────────────────────┼──────────────────────────┐
        ▼                          ▼                          ▼
   real teacher calls          DiLoCo on                  production
   over OpenRouter?            >1 GPU?                    GRPO training?
        │                          │                          │
        │ yes                      │ yes                      │ yes
        ▼                          ▼                          ▼
   pip install -e ".[replay]"  pip install -e ".[diloco]"  pip install -e ".[train]"
   (httpx)                     (torchft-nightly)           (trl, peft, accelerate, datasets)
        │                          │                          │
        │ + want CPU-side          │ + scaling beyond a        │ + want PRIME-RL
        │ DPO normalization?       │ single host?              │ (Recipe C)?
        ▼                          ▼                          ▼
   pip install -e \".[replaysim]\"  pip install -e \".[serverless]\"  pip install -e \".[prime-rl]\"
   (data-juicer; depends         (fsspec, huggingface_hub)    (prime-rl>=0.5)
    on [replay])
                                                              │ + Monarch actor mesh?
                                                              ▼
                                                          pip install -e \".[monarch]\"
                                                          (monarch>=0.4.1)
```

Quick decision table:

| Goal                                                  | Install                                  |
|-------------------------------------------------------|------------------------------------------|
| CPU smoke / verification (sections 3, 5, 6)           | `pip install -e .`                       |
| Section 4 (replaysim DJNormalizer)                    | `pip install -e ".[replaysim]"`          |
| Section 7 dev loop (LocalProcessExecutor + file://)   | `pip install -e ".[serverless]"`         |
| Real DiLoCo outer-loop                                | `pip install -e ".[diloco,serverless]"`  |
| Section 8 Recipe A (TRL GRPO)                         | `pip install -e ".[train]"`              |
| Section 8 Recipe C (PRIME-RL)                         | `pip install -e ".[prime-rl]"`           |
| Section 8 Recipe C+D (PRIME-RL + Monarch)             | `pip install -e ".[prime-rl,monarch]"`   |
| Everything for development                            | `pip install -e ".[dev]"`                |

---

## 3. Quickstart: `examples/qwen_05b_quickstart` end-to-end on CPU

The fastest way to convince yourself the framework works on a real HF model.
~3–5 min wall-clock on CPU, ~1 GB disk for Qwen2.5-0.5B weights.

```bash
pip install -e .
python examples/qwen_05b_quickstart/run.py
```

What the script does (read the source at
`examples/qwen_05b_quickstart/run.py`):

1. Pin RNG (`random.seed(42)`, `torch.manual_seed(42)`) so the per-step
   numbers below are reproducible.
2. Load `Qwen/Qwen2.5-0.5B-Instruct` on CPU in fp32, set `model.train()`.
3. `batch = build_batch(tokenizer, device="cpu")` — a real chat-template-formatted
   batch with all keys the 3-channel composer might consume.
4. Five backward steps with `compose_loss(model, batch, alpha_sdpo=0.1,
   beta_replay=0.05)`; an `AdamW(lr=1e-5)` optimizer; finite-grad check
   after each step.

Expected output (transcribed from `examples/qwen_05b_quickstart/run.log`):

```
[quickstart] loading Qwen/Qwen2.5-0.5B-Instruct (CPU, fp32) ...
[quickstart] loaded — 0.494B params
[quickstart] building real chat-template batch ...
[quickstart] running 5 backward steps ...
  step 0: total=0.7390  lm_ce=0.7358  sdpo=0.0000  dpo=0.0639  finite=True
  step 1: total=0.0379  lm_ce=0.0351  sdpo=0.0000  dpo=0.0563  finite=True
  step 2: total=0.0122  lm_ce=0.0110  sdpo=0.0000  dpo=0.0240  finite=True
  step 3: total=0.0060  lm_ce=0.0055  sdpo=0.0000  dpo=0.0098  finite=True
  step 4: total=0.0031  lm_ce=0.0029  sdpo=0.0000  dpo=0.0044  finite=True
========================================================
  Initial loss: 0.7390  →  Final loss: 0.0031  →  Reduction: 99.6%
  Verdict: PASS
========================================================
```

How to read this:

- **`total` collapses by ~99%.** The model successfully memorizes the
  single batch — exactly what you expect from an SGD pass on a 0.5B model
  with one fixed input. This is a wiring check, not a generalization claim.
- **`lm_ce` carries almost all the magnitude.** Channel 1 (the GRPO stub)
  is doing the work — the response tokens are short and have low entropy
  under the trained model.
- **`sdpo=0.0000` on every step.** Channel 2 has auto-disabled because the
  default `build_batch` does not include `ctx_teacher_input_ids`. Compare
  the conditional in `compose_loss`:
  ```python
  if (alpha_sdpo > 0.0
      and "ctx_teacher_input_ids" in inputs
      and inputs["ctx_teacher_input_ids"].numel() > 0):
  ```
  — channel auto-off if either the weight or the inputs are missing.
- **`dpo > 0` and trending down.** The batch *does* include
  `dpo_chosen_input_ids`, `dpo_chosen_response_mask`,
  `dpo_chosen_ref_logprobs` (and the rejected counterparts), so channel 3
  is live.
- **`finite=True`** — every step's `p.grad` was finite for every parameter.
  This is the wiring contract; if it ever flips to `False` the smoke fails.

If you see `Verdict: PASS`, the framework is correctly installed and
gradients flow through all live channels. You are ready for section 4.

---

## 4. Adding the trace-replay channel

The quickstart batch *had* DPO inputs, but they were synthetic — the
`build_batch` helper bakes them in. To get **real** DPO pairs from
multi-teacher disagreement, use the replaysim package.

### 4a. Spin up `replay_trace`

```python
import asyncio
from composer_replication import (
    DEFAULT_TEACHERS, replay_trace, extract_dpo_pairs,
)

# Trace must be a list[TraceState]; see composer_replication/teacher_replay.py
# for the exact TypedDict shape. Each state holds a chat-messages prefix +
# the student's actual action at that step.
states = [...]   # your frozen agentic trace; see spike 001 for a 50-step example

teacher_actions = asyncio.run(
    replay_trace(
        states=states,
        teachers=DEFAULT_TEACHERS,    # claude-opus-4.7 + gpt-5 + deepseek-v4-pro
        max_total_usd=10.0,           # hard ceiling (spike 001 measured $0.98/trace mean)
    )
)
```

The 3 teachers are queried in parallel via OpenRouter
(`OPENROUTER_API_KEY` in env or `~/.hermes/.env`), latencies recorded,
costs tracked.

### 4b. Get `DPOPair`s from disagreement

```python
pairs = extract_dpo_pairs(
    states=states,
    teacher_actions=teacher_actions,
    agreement_threshold=2,    # at least 2/3 teachers must agree on the chosen action
)
```

Each pair is a `DPOPair` TypedDict with the exact shape the
`DJNormalizer` and downstream training expects:

```python
class DPOPair(TypedDict):
    state_id:           str
    state_messages:     list[dict]    # conversation context
    chosen:             str           # teacher-consensus action
    rejected:           str           # student action
    n_teachers_agreeing: int
```

(verified in `composer_replication/teacher_replay.py:99–105`).

### 4c. Run `DJNormalizer` with `default.yaml`

```python
from composer_replication.replaysim import DJNormalizer

normalizer = DJNormalizer()        # uses recipes/replaysim/default.yaml
normalized = normalizer.normalize(pairs)
# → list[NormalizedDPOPair]
```

`DJNormalizer` shells out to data-juicer's `DefaultExecutor` under the hood
(file-in / file-out contract). The default recipe at
`composer_replication/recipes/replaysim/default.yaml` runs four CPU-only ops
in order:

1. `text_length_filter` (8 ≤ chars ≤ 32000) on `chosen` and `rejected`
2. `words_num_filter` (2 ≤ words ≤ 4096) on both
3. `special_characters_filter` (≤50% non-alpha) on both
4. `document_deduplicator` (per-batch hashing, lowercase, ignore non-character) on `chosen`

Records carry **two parallel shapes** for `chosen`/`rejected`:
- flat strings (`chosen`, `rejected`) → consumed by data-juicer's text_key-based filters
- chat-messages lists (`chosen_messages`, `rejected_messages`) → preserved for chat-aware ops + round-trip

This dual-shape design (verified in the test
`test_dpo_pair_to_dj_record_shape`,
`composer_replication/replaysim/tests/test_replaysim.py:44`) is what
unblocked the data-juicer integration in Wave 14.

### 4d. The 3-record fixture from spike 001

The fixture lives at
`spikes/001-teacher-replay-cost/states.jsonl` (50 states) and
`spikes/001-teacher-replay-cost/results.jsonl` (the teacher responses, all
priced and timed). The first 3 states are:

```jsonl
{"id": "state-000", "task": "Fix the failing test in tests/test_auth.py::test_login_with_email", ...}
{"id": "state-001", "task": "Add rate-limiting middleware to the Flask app", ...}
{"id": "state-002", "task": "Refactor the parse_config function — it's 200 lines and has 3 responsibilities", ...}
```

For each, all 3 teachers answered (claude-opus-4.7, gpt-5, deepseek-v4-pro);
agreement on the `(c)` choice for state-000 and state-001 (read more
files / check schema first) drives a clean DPO pair where the student's
action becomes the rejected. For state-002, all 3 agreed on `(c)` (write
characterization tests first) → another clean pair. These three records
pass through the `DJNormalizer` default recipe unchanged (length, words,
special-char ratios all in bounds; no duplicates).

The full 50-state trace cost **$0.98 mean** end-to-end across all three
teachers (spike 001 verdict). The framework's cost ceiling
(`max_total_usd`) and VOI gating drop this to ~$0.30/trace projected.

### 4e. End-to-end one-liner

```python
from composer_replication.replaysim import replay_and_normalize_trace

teacher_actions, normalized_pairs = await replay_and_normalize_trace(
    states=states,
    teachers=DEFAULT_TEACHERS,
    agreement_threshold=2,
    max_total_usd=10.0,
)
```

(`async def`; for sync callers use the sibling `replay_and_normalize_trace_sync`
in `composer_replication.replaysim.normalize`.)

---

## 5. Switching DPO → SimPO: one kwarg

```python
components = compose_loss(
    model, batch,
    alpha_sdpo=0.1,
    beta_replay=0.05,
    dpo_variant="simpo",      # ← the only line that changes
    simpo_beta=2.0,           # paper default
    simpo_gamma=1.0,          # paper default
)
```

The kwarg is verified in `compose_loss`'s signature
(`composer_replication/loss.py:81`):

```python
dpo_variant: Literal["dpo", "simpo"] = "dpo",
```

### What changes in the loss curve

- **Channel 3 input requirements drop.** `compose_loss` no longer reads
  `dpo_chosen_ref_logprobs` / `dpo_rejected_ref_logprobs`. Reference-model
  VRAM cost goes to zero. (Source: `composer_replication/loss.py:111–113`
  and `composer_replication/distillation/simpo.py:23–27`.)
- **Loss scale shifts.** Standard DPO is
  `-logsigmoid(β·[(logπ(c) - logπ_ref(c)) - (logπ(r) - logπ_ref(r))])`.
  SimPO is `-logsigmoid(β·[avg_logπ(c) - avg_logπ(r)] - γ)` — average
  per-token log-prob (length-normalized) and a constant target margin γ.
- **Loss is ≤ DPO loss when chosen/rejected separation is large.** The
  unit test `test_simpo_loss_lower_for_better_separation`
  (`composer_replication/distillation/tests/test_distillation_losses.py:35`)
  verifies that a wider chosen-vs-rejected gap drives lower SimPO loss —
  meaning, in practice, SimPO curves are *steeper* than DPO when the
  preference signal is strong, and *flatter* when it's weak.
- **No KL-against-reference regularization.** This is both the upside (no
  ref-model serving) and the risk (more tendency to drift). Watch for
  reward-hacking-style degeneracies if your preference data has noise.

### When to use SimPO

- **GPU-poor.** You can't afford to keep a frozen reference policy resident
  alongside the trainer.
- **Cold-start preference data.** Length-normalization (avg_logπ vs sum)
  helps when chosen/rejected lengths are wildly imbalanced — common in
  agentic traces where the student's failed attempt is short and the
  teacher's correction is long.
- **You don't have ref logprobs precomputed.** SimPO needs nothing from
  the reference policy.

When to **stay on DPO**: when you need the explicit KL anchor against
a known-good reference (e.g., when training over a long horizon and you
want to bound the drift), or when your preference data is very noisy and
the reference acts as a regularizer.

---

## 6. Adding TAID / Entropy-Aware OPD wrappers

Channel 2 (SDPO/OPSD) can be replaced by **TAID** (Sakana AI,
arXiv:2501.16937) for capacity-gap distillation, or by
**Entropy-Aware OPD** (ICLR 2026 Spotlight) for per-token forward/reverse-KL
gating. Both are wired through `compose_loss`:

```python
sdpo_wrapper: Literal["none", "taid", "entropy_opd"] = "none",
taid_t: float | None = None,         # current TAID interpolation coeff
entropy_opd_h_max: float | None = None,
```

(verified at `composer_replication/loss.py:82–93`.)

### TAID (upstream-faithful port)

> **Wave 15 rewrite, breaking change.** The previous in-tree TAID was
> algorithmically different from the paper (it mixed in probability space
> against a frozen step-0 student snapshot and wrapped a symmetric JSD
> criterion). It has been replaced with an upstream-faithful port:
> logit-space mix, current-student-detached anchor, forward-KL criterion.
> Old kwargs `taid_schedule_step`, `taid_total_steps`, `taid_schedule`,
> `taid_alpha_min`, `taid_alpha_max`, plus `inputs["student_init_logits"]` /
> `inputs["student_init_input_ids"]` are **gone**. They have no upstream
> analogue. Use `taid_t` (and optionally `TAIDScheduler`) instead.

The TAID criterion is forward-KL against a logit-space-interpolated target:

```
p_t = softmax( (1 - t) · stop_grad(student_logits) + t · teacher_logits )
L   = - mean_token  Σ_v  p_t(v) · log_softmax(student_logits)(v)
```

where `t ∈ [0, 1]` is the interpolation coefficient. At `t=0` the target
is the (detached) student itself — the loss is the entropy of that
distribution and contributes no gradient to the student. At `t=1` it
reduces to standard forward-KL distillation against the teacher.

The schedule that produces `t` is the **trainer's** responsibility. The
package ships an optional `TAIDScheduler` that mirrors the paper's
adaptive momentum scheme:

```python
from composer_replication.distillation import TAIDScheduler

sched = TAIDScheduler(num_train_steps=10_000)   # paper defaults
for step in range(num_train_steps):
    components = compose_loss(
        model, batch,
        sdpo_wrapper="taid",
        taid_t=sched.t,
    )
    components.total.backward(); optimizer.step()
    sched.update_t(components.sdpo_jsd.detach(), global_step=step)
```

`TAIDScheduler` defaults match upstream: `t_start=0.4`, `t_end=1.0`,
`alpha=5e-4`, `beta=0.99`. Pass `disable_adaptive=True` to fall back to
the deterministic linear schedule
`t = t_start + progress · (t_end - t_start)`.

If you want a simple fixed schedule (no scheduler), just compute `t`
yourself and pass it in — `compose_loss` validates `taid_t ∈ [0, 1]`.

### Upstream-parity test

`composer_replication/distillation/tests/test_taid_parity.py` skip-imports
the upstream reference at `/tmp/taid-clone` (clone with
`git clone --depth 1 https://github.com/SakanaAI/TAID /tmp/taid-clone`)
and asserts our `taid_loss(student, teacher, mask, t)` matches upstream
`TAID.compute_loss(...)` within `atol=rtol=1e-5` across `t ∈ {0.0, 0.1, 0.4,
0.5, 0.9, 1.0}`. This is the load-bearing parity guarantee.

### Entropy-Aware OPD

Drop-in for channel 2 — gates between forward KL (mode-covering) and
reverse KL (mode-seeking) per token, weighted by the teacher's entropy:

```
L = Σ_t  w(t) · KL_fwd_t  +  (1 - w(t)) · KL_rev_t
w(t) = clamp(H_teacher(t) / h_max, 0, 1)
```

`entropy_opd_h_max=None` (the default) auto-sets `h_max = log(V)` (the
maximum-entropy bound for a vocab-V softmax).

### Boundary-condition unit test (proof of correctness)

The test `test_taid_loss_t_zero_target_matches_detached_student`
(`composer_replication/distillation/tests/test_distillation_losses.py`)
pins TAID's `t=0` invariant — the teacher is *completely* hidden from the
gradient because the target collapses to `softmax(student.detach())`:

```python
def test_taid_loss_t_zero_target_matches_detached_student():
    s1 = torch.randn(1, 2, 4, requires_grad=True)
    teacher_a = torch.zeros(1, 2, 4); teacher_a[..., 0] = 10.0
    teacher_b = torch.zeros(1, 2, 4); teacher_b[..., 3] = 10.0
    mask = torch.ones(1, 2)
    loss_a = taid_loss(s1, teacher_a, mask, t=0.0)
    loss_b = taid_loss(s1, teacher_b, mask, t=0.0)
    # Two completely different teachers must give the same loss at t=0.
    assert abs(float(loss_a) - float(loss_b)) < 1e-6
```

This is the load-bearing test for TAID: if the `t=0` endpoint ever leaks
teacher signal into the gradient, this test fires and the contract is
broken. The companion test `test_taid_loss_t_one_is_pure_forward_kl`
pins the `t=1` endpoint by hand-computing `-Σ p_teacher · log_q` and
asserting equality.
---

## 7. Going multi-replica with serverless DiLoCo

DiLoCo is the outer-loop optimizer that lets you run N replicas in
parallel, sync them every H inner steps, and tolerate slow links — see
`docs/adrs/ADR-005-serverless-diloco.md` for the design. The framework
gives you three increasingly-distant deployments:

### Step 1 — `LocalProcessExecutor` for development

```python
from composer_replication.diloco.serverless import (
    LocalProcessExecutor, ObjectStoreAllReduce,
)
import tempfile

with tempfile.TemporaryDirectory() as td:
    rendezvous = ObjectStoreAllReduce(td, rank=0, world_size=4)
    executor = LocalProcessExecutor()
    handles = executor.launch_replicas(
        n_replicas=4,
        entrypoint="composer_replication.diloco.serverless.replica_entrypoint",
        entrypoint_args={"rendezvous_uri": td, "rank_env": "REPLICA_RANK"},
    )
    results = executor.collect(handles, timeout=600)
```

`LocalProcessExecutor` (`composer_replication/diloco/serverless/executor.py:160`)
spawns N child processes via `multiprocessing.get_context("spawn")` and
sets `REPLICA_RANK={0..N-1}` in each child's env. It satisfies the
`ServerlessExecutor` Protocol (line 35) — the same Protocol the cloud
adapters implement. So the dev-loop code is byte-identical to the cloud
deploy: only the executor instance changes.

### Step 2 — `ObjectStoreAllReduce` as the rendezvous

```python
# Local file:// for tests
rendezvous = ObjectStoreAllReduce("/tmp/diloco-runs/run42/", rank=0, world_size=4)

# Real S3 (after `pip install -e .[serverless]`)
rendezvous = ObjectStoreAllReduce(
    "s3://my-bucket/diloco-runs/run42/",
    rank=0, world_size=4,
    timeout_s=1800.0,
)
```

The communication pattern is `S3 PutObject + N GetObjects` once per
inner H steps (matches DiLoCo's actual sync cadence,
arXiv:2311.08105 §3.2). For 1B-param bf16, that's ~2 GB / 30 minutes
per replica — well within S3 free-tier. On the inner side the framework
exposes a `MockManager` that drops into the `torchft.Manager` slot, so
you can validate the rendezvous logic before plugging in real torchft
(verified by `test_serverless_diloco_integration.py`).

### Step 3 — point at `ModalExecutor` / `HFJobsExecutor`

```python
# Modal (skeleton at composer_replication/diloco/serverless/modal.py)
from composer_replication.diloco.serverless.modal import ModalExecutor
executor = ModalExecutor(image="modal:python3.11", gpu="A100")

# HuggingFace Jobs (skeleton at composer_replication/diloco/serverless/hf_jobs.py)
from composer_replication.diloco.serverless.hf_jobs import HFJobsExecutor
executor = HFJobsExecutor(hardware="a10g-large")

# Same Protocol — same launch_replicas / poll / collect calls as Local
handles = executor.launch_replicas(n_replicas=4, ...)
```

Both adapters check their cloud SDK at `__init__` time (not at module
import) so they don't break the package if you don't have `modal` or
`huggingface_hub` installed. Production maturity: dev-ready for cloud
trial; per ADR-005, full HA-cluster fan-out lives in v0.2+.

---

## 8. Picking an RL backend

Four canonical recipes, each tied to an upstream framework. Source:
`docs/INTEGRATION_ARCHITECTURE.md` Recipes A–D plus
`docs/adrs/ADR-006-rl-frameworks.md`.

### Recipe A — TRL `GRPOTrainer` subclass

`ComposerReplicationTrainer` is a `trl.GRPOTrainer` subclass that
overrides `_compute_loss(model, inputs)` to compose the same 3 channels
on top of TRL's real reward + advantage machinery. Install:
`pip install -e ".[train]"`. Then:

```python
from composer_replication import ComposerReplicationTrainer
trainer = ComposerReplicationTrainer(model=..., reward_funcs=[...], ...)
trainer.train()
```

**When to use it:** This is the v0.0/v0.1 recommended path. You want
real GRPO with rewards, you have HF model + dataset + (mostly) standard
GRPO infrastructure, and you don't need >100B-param scale. TRL is
mature, the trainer is a small subclass, and the same `compose_loss`
math runs in both the verification harness and in production with no
re-coding.

→ See `docs/INTEGRATION_ARCHITECTURE.md` § "Recipe A: TRL `GRPOTrainer`
subclass" (line 205).

### Recipe B — VeRL custom `adv_estimator` + DataProto extension

VeRL replaces TRL's reward+advantage machinery with a Ray-driven actor
graph that's specifically optimized for distributed RL training of
large LMs. Composition with the framework: extend `DataProto` with the
hint-conditioned columns + DPO pair fields, register a custom
`adv_estimator` that calls the same `compose_loss` body.

**When to use it:** You're past 7B-param, you have multi-host setup
(Ray cluster), and TRL's single-process trainer is the bottleneck. VeRL
is the recommended scale path for v0.2+. Trade-off: the extension surface
is larger and the docs are sparser than TRL's.

→ See `docs/INTEGRATION_ARCHITECTURE.md` § "Recipe B: VeRL custom
`adv_estimator`" (line 289).

### Recipe C — PRIME-RL with DPPO-clip details

`composer_replication/recipes/prime_rl/composer_loss.py` ships a
`loss_fn(inputs, *, alpha_sdpo=0.0, beta_dpo=0.0, dppo_mask_high=0.2,
dppo_mask_low=0.2, adv_tau=1.0, kl_tau=1e-3)` adapter that maps
PRIME-RL's `LossInputs` struct (1-D per-sample tensors:
`trainer_logprobs`, `inference_logprobs`, `teacher_logprobs`,
`advantages`, `loss_mask`) into our 3-channel composition.

The DPPO+KL bit is what makes PRIME-RL distinctive — and we mirror
PRIME-RL's upstream `default_loss_fn` exactly (verified against
`prime_rl/trainer/rl/loss.py` lines 116-165):

```python
log_ir       = trainer_logprobs - inference_logprobs
ir           = exp(log_ir)                                  # importance ratio
probs_diff   = exp(trainer_logprobs) - exp(inference_logprobs)
invalid_high = probs_diff >  dppo_mask_high                 # for positive-advantage tokens
invalid_low  = probs_diff < -dppo_mask_low                  # for negative-advantage tokens
invalid      = where(advantages > 0, invalid_high, invalid_low)
keep         = loss_mask & ~invalid
pg_loss      = keep      * (adv_tau * advantages) * ir
kl_loss      = loss_mask * log_ir**2
loss         = (-pg_loss + kl_tau * kl_loss).sum()
```

Three things to remember: (1) the mask gate is on **probability-space**
`exp(trainer_lp) - exp(inference_lp)`, not on the log-ratio; (2) the
policy-gradient term is multiplied by the importance ratio
`exp(trainer_lp - inference_lp)`, not by `trainer_lp` directly (proper
IS-corrected gradient, not REINFORCE); (3) the mask is **conditioned on
the sign of the advantage** — positive-advantage tokens are dropped on
the upper bound, negative-advantage tokens on the lower. Defaults
`dppo_mask_high=dppo_mask_low=0.2` and `adv_tau=1.0, kl_tau=1e-3`
match PRIME-RL's `DefaultLossConfig` (all fields `Field(..., ge=0)`).
SDPO (channel 2) is gated `NotImplementedError` in v0 because PRIME-RL
exposes log-probs, not full vocab logits; channel 3 (trace-replay DPO)
emits a warning if `beta_dpo != 0`.

**When to use it:** You're already in the PRIME-Intellect / decentralized
training universe, you want INTELLECT-style scaling on a long-horizon
training run, and DPPO masking is part of your existing reward+advantage
recipe. Install: `pip install -e ".[prime-rl]"`.

→ See `composer_replication/recipes/prime_rl/prime_rl_recipe.md` and
`docs/INTEGRATION_ARCHITECTURE.md` § "Recipe C: TorchForge + Monarch"
(line 356).

### Recipe C+D — Monarch as actor mesh

Monarch (the actor framework underpinning TorchForge) hosts the
trainer/generator/manager actors in a topology-aware mesh. The framework
ships *skeleton* actor definitions at
`composer_replication/recipes/monarch/actors.py` (TrainerActor,
GeneratorActor) and a layout doc at `monarch_actor_layout.md`. v0
intentionally *fails fast* if you try to instantiate them
(`raise NotImplementedError("v0 skeleton; deferred to v0.2 per ADR-006")`)
because the upstream Monarch API is still moving.

**When to use it:** Reference-pattern reading only in v0. Decision point
is v0.2 once the upstream actor API stabilizes. Treat the skeleton as
shape-of-the-final-answer documentation, not as a production target.
Install: `pip install -e ".[prime-rl,monarch]"` for the full surface.

→ See `composer_replication/recipes/monarch/monarch_actor_layout.md`
and `docs/adrs/ADR-006-rl-frameworks.md`.

---

## Common pitfalls + what tests catch them

The framework's 115-test suite (post-Wave-15) is structured so each pitfall has a
specific test-file home. If you hit one of these in production, the
corresponding test is your fastest reproducer.

| Pitfall                                                                                       | Symptom                                              | Test file (catches it)                                                                                              |
|-----------------------------------------------------------------------------------------------|------------------------------------------------------|----------------------------------------------------------------------------------------------------------------------|
| Forgetting `taid_schedule_step` when `sdpo_wrapper="taid"`                                    | `ValueError` at first step                           | `composer_replication/tests/test_compose_loss_integration.py` (kwarg validation)                                    |
| TAID α=0 endpoint leaks teacher signal                                                        | Teacher swap changes the loss when α should be 0     | `test_taid_loss_alpha_zero_ignores_teacher` in `composer_replication/distillation/tests/test_distillation_losses.py:153` |
| TAID α=1 endpoint differs from plain SDPO                                                     | Bit-difference vs reference SDPO at the schedule end | `test_taid_blended_logits_endpoints` in `composer_replication/distillation/tests/test_distillation_losses.py:115`   |
| SimPO loss not differentiable through the loss-of-sigmoid path                                | `chosen.grad is None` after backward                 | `test_simpo_loss_differentiable` in `composer_replication/distillation/tests/test_distillation_losses.py:50`        |
| SimPO shape-mismatch slips through silently                                                   | Broadcasting bug, NaN downstream                     | `test_simpo_loss_shape_mismatch_raises` in `composer_replication/distillation/tests/test_distillation_losses.py:61` |
| Entropy-OPD failing to zero out when distributions match                                      | Loss > 0 when student≡teacher                        | `test_entropy_aware_opd_zero_when_distributions_match` in `composer_replication/distillation/tests/test_distillation_losses.py:217` |
| Entropy of one-hot ≠ 0 / entropy of uniform ≠ log(V)                                          | Wrong gating weights w(t)                            | `test_teacher_entropy_one_hot_is_zero` and `test_teacher_entropy_uniform_is_log_v` in `composer_replication/distillation/tests/test_distillation_losses.py:175,183` |
| `DJNormalizer` records missing the chat-messages shape                                        | Filters silently no-op or crash                      | `test_dpo_pair_to_dj_record_shape` in `composer_replication/replaysim/tests/test_replaysim.py:44`                   |
| `DJNormalizer` round-trip drops `state_messages` / metadata                                   | Lost provenance                                      | `test_dj_record_to_normalized_roundtrip` and `test_dj_record_to_normalized_preserves_state_messages` in `composer_replication/replaysim/tests/test_replaysim.py` |
| `ObjectStoreAllReduce` accepts an out-of-bounds rank                                          | Silent corruption of the all-reduce average          | `test_object_store_allreduce_init_validates_rank` in `composer_replication/diloco/serverless/tests/test_serverless_local.py:31` |
| `ObjectStoreAllReduce(world_size=1)` doesn't passthrough cleanly                              | False all-reduce on single replica                   | `test_object_store_allreduce_world_size_1_passthrough` in `composer_replication/diloco/serverless/tests/test_serverless_local.py:46` |
| `LocalProcessExecutor` doesn't propagate child failures to `collect()`                        | Silent test pass when a replica crashed              | `test_serverless_diloco_integration.py` in `composer_replication/diloco/serverless/tests/`                          |
| PRIME-RL adapter accidentally uses `(B, T)` shape instead of per-sample `(seq,)`              | Shape mismatch / wrong reduction                     | `composer_replication/recipes/prime_rl/tests/test_composer_loss.py` (10 tests covering shape and DPPO mask edges)   |
| Channel 2/3 fails to auto-disable when its inputs are absent                                  | Crash on missing key, not graceful skip              | `composer_replication/tests/test_compose_loss_integration.py` (`(a) defaults reproduce existing compose_loss output bit-exact`) |

Run the full suite with `pytest` from the repo root.

---

**File path:** `/mnt/e/CS/HF/composer-replication-framework/docs/USER_GUIDE.md`