Codeseys/composer-replication-framework
0
1# API Reference — composer-replication-framework2 3Complete reference for every public symbol in `composer_replication`. Source-of-truth is the `.py` files in `composer_replication/`; docstrings have been pulled verbatim where they exist and supplemented where missing.4 5**Legend**6 7- ⚠️ **UNTESTED-CONTRACT** — symbol exists and is callable, but its behaviour is not pinned by an automated test in `composer_replication/**/tests/` or `spikes/**/tests/`.8- 🟡 **SKELETON** — class/method body raises `NotImplementedError`; ships as design-of-record per ADR-005 / ADR-006.9 10**Module groups (in this document)**11 121. `composer_replication` (top-level re-exports)132. `composer_replication.loss`143. `composer_replication.batch`154. `composer_replication.opsd`165. `composer_replication.distillation`176. `composer_replication.teacher_replay`187. `composer_replication.replaysim`198. `composer_replication.ingestion` (+ `.claude_code`)209. `composer_replication.hint_generator`2110. `composer_replication.trainer` (+ `.composer_trainer`, `.data_collator`)2211. `composer_replication.diloco`2312. `composer_replication.diloco.serverless` (+ `.executor`, `.allreduce`, `.modal`, `.hf_jobs`, `.replica_entrypoint`)2413. `composer_replication.recipes.prime_rl.composer_loss`2514. `composer_replication.recipes.monarch.actors`26 27---28 29## 1. `composer_replication` — top-level package30 31The package re-exports the most common entry points from sub-modules. `__all__` is the canonical list of public top-level names.32 33### `composer_replication.__version__: str`34 35Package version string. Currently `"0.1.0"`.36 37```python38import composer_replication39print(composer_replication.__version__) # "0.1.0"40```41 42### `composer_replication._DILOCO_AVAILABLE: bool`43 44`True` iff `torchft` is importable in the running Python environment (gates `make_diloco_outer_loop`). Set to `False` and `make_diloco_outer_loop` is set to `None` when `torchft` is missing.45 46```python47from composer_replication import _DILOCO_AVAILABLE48if _DILOCO_AVAILABLE:49 from composer_replication import make_diloco_outer_loop50```51 52### Re-exports53 54| Name | Source module |55|---|---|56| `compose_loss` | `composer_replication.loss` |57| `LossComponents` | `composer_replication.loss` |58| `build_batch` | `composer_replication.batch` |59| `generalized_jsd_loss` | `composer_replication.opsd` |60| `ClaudeCodeIngester` | `composer_replication.ingestion.claude_code` |61| `IngestionStats` | `composer_replication.ingestion.claude_code` |62| `SYSTEM_PROMPT` | `composer_replication.ingestion.claude_code` |63| `DEFAULT_TEACHERS` | `composer_replication.teacher_replay` |64| `DPOPair` | `composer_replication.teacher_replay` |65| `TeacherCallResult` | `composer_replication.teacher_replay` |66| `TeacherSpec` | `composer_replication.teacher_replay` |67| `TraceState` | `composer_replication.teacher_replay` |68| `extract_dpo_pairs` | `composer_replication.teacher_replay` |69| `replay_trace` | `composer_replication.teacher_replay` |70| `ComposerReplicationTrainer` | `composer_replication.trainer` |71| `make_diloco_outer_loop` | `composer_replication.diloco` (or `None` if `torchft` missing) |72 73See each source module below for full signatures.74 75---76 77## 2. `composer_replication.loss`78 79Verification-harness 3-channel loss. Free function, does not depend on `trl`.80 81### `class LossComponents`82 83```python84@dataclass85class LossComponents:86 lm_ce: torch.Tensor87 sdpo_jsd: torch.Tensor88 trace_replay_dpo: torch.Tensor89 total: torch.Tensor90 91 def detached(self) -> dict[str, float]: ...92```93 94Per-channel breakdown of the total loss for logging and ablation. All four fields are scalar `torch.Tensor`s (`shape=()`); `total = lm_ce + alpha_sdpo * sdpo_jsd + beta_replay * trace_replay_dpo`.95 96**`detached() -> dict[str, float]`** — returns Python-float copies of all four fields with no grad. Useful for W&B logging.97 98```python99from composer_replication import compose_loss, build_batch100components = compose_loss(model, build_batch(tokenizer))101print(components.detached()) # {'lm_ce': 2.34, 'sdpo_jsd': 0.12, ...}102components.total.backward()103```104 105### `compose_loss(model, inputs, *, ...) -> LossComponents`106<a id="compose_loss"></a>107 108```python109def compose_loss(110 model: torch.nn.Module,111 inputs: dict[str, torch.Tensor],112 *,113 alpha_sdpo: float = 0.1,114 beta_replay: float = 0.05,115 sdpo_jsd_beta: float = 0.5,116 sdpo_temperature: float = 1.0,117 sdpo_token_clip: float | None = None,118 replay_dpo_beta: float = 0.1,119 lm_ce_label_smoothing: float = 0.0,120 dpo_variant: Literal["dpo", "simpo"] = "dpo",121 sdpo_wrapper: Literal["none", "taid", "entropy_opd"] = "none",122 taid_t: float | None = None,123 simpo_beta: float = 2.0,124 simpo_gamma: float = 1.0,125 entropy_opd_h_max: float | None = None,126) -> LossComponents127```128 129Compute `total = lm_ce + alpha_sdpo * sdpo_jsd + beta_replay * trace_replay_dpo`.130 131**Required keys in `inputs`**132 133- `input_ids`: `(B, T_s)` student rollout token ids.134- `response_mask`: `(B, T_s)` 1 on assistant-response tokens, 0 elsewhere.135 136**Optional keys** (channel auto-disables if missing OR if its weight = 0):137 138- SDPO: `ctx_teacher_input_ids` `(B, T_t)`, `sdpo_loss_mask` `(B, T_t)`.139- DPO (`dpo_variant="dpo"`): `dpo_chosen_input_ids`, `dpo_chosen_response_mask`, `dpo_rejected_input_ids`, `dpo_rejected_response_mask`, `dpo_chosen_ref_logprobs`, `dpo_rejected_ref_logprobs` (precomputed).140- SimPO (`dpo_variant="simpo"`): same DPO ids/masks; reference logprobs are silently ignored.141- TAID (`sdpo_wrapper="taid"`): no extra `inputs` keys needed; the optional `sdpo_loss_mask` is reused as the per-token TAID mask. Pass `taid_t` directly (or drive it from `TAIDScheduler`).142 143**Parameters**144 145| Name | Type | Default | Meaning |146|---|---|---|---|147| `model` | `torch.nn.Module` | — | HF causal-LM. Must accept `input_ids=` and return an object with `.logits`. |148| `inputs` | `dict[str, torch.Tensor]` | — | Batch dict (see required/optional keys above). |149| `alpha_sdpo` | `float` | `0.1` | Weight on SDPO/JSD channel. `0.0` disables. |150| `beta_replay` | `float` | `0.05` | Weight on trace-replay DPO channel. `0.0` disables. |151| `sdpo_jsd_beta` | `float` | `0.5` | β param for `generalized_jsd_loss` (0=fwd KL, 0.5=JSD, 1=rev KL). Unused when `sdpo_wrapper="taid"`. |152| `sdpo_temperature` | `float` | `1.0` | Softmax temperature in SDPO. Unused when `sdpo_wrapper="taid"`. |153| `sdpo_token_clip` | `float \| None` | `None` | Per-token JSD clamp. |154| `replay_dpo_beta` | `float` | `0.1` | β in standard DPO logit. |155| `lm_ce_label_smoothing` | `float` | `0.0` | `F.cross_entropy(label_smoothing=)`. |156| `dpo_variant` | `Literal["dpo","simpo"]` | `"dpo"` | Channel-3 algorithm. |157| `sdpo_wrapper` | `Literal["none","taid","entropy_opd"]` | `"none"` | Channel-2 wrapper. |158| `taid_t` | `float \| None` | `None` | Current TAID interpolation coefficient in `[0, 1]`. Required when `sdpo_wrapper="taid"`. Drive from `TAIDScheduler` or pass a fixed value. |159| `simpo_beta` | `float` | `2.0` | SimPO β (paper default). |160| `simpo_gamma` | `float` | `1.0` | SimPO target margin γ (paper default). |161| `entropy_opd_h_max` | `float \| None` | `None` | Max-entropy normalizer; `None` ⇒ `log(V)`. |162 163**Returns** `LossComponents` (see above).164 165**Raises** `ValueError` if `dpo_variant` or `sdpo_wrapper` is unknown, if `sdpo_wrapper="taid"` is requested without `taid_t`, or if `taid_t` is outside `[0, 1]`.166 167```python168from composer_replication import compose_loss, build_batch169batch = build_batch(tokenizer)170out = compose_loss(model, batch, alpha_sdpo=0.1, beta_replay=0.05)171out.total.backward()172print(out.detached())173```174 175---176 177## 3. `composer_replication.batch`178 179Verification-harness batch builder.180 181### `build_batch(tokenizer, *, ...) -> dict[str, torch.Tensor]`182 183```python184def build_batch(185 tokenizer: Any,186 *,187 device: torch.device | str = "cpu",188 seed: int = 42,189 variant: str = "factorial",190 align_sdpo_shapes: bool = False,191) -> dict[str, torch.Tensor]192```193 194Construct a full 3-channel batch from a real HF tokenizer. The DPO ref-logprobs are dummy tensors (the smoke verifies loss composition wires together, not the reference-policy precompute).195 196**Returned keys**: `input_ids`, `response_mask`, `ctx_teacher_input_ids`, `sdpo_loss_mask`, `dpo_chosen_input_ids`, `dpo_chosen_response_mask`, `dpo_rejected_input_ids`, `dpo_rejected_response_mask`, `dpo_chosen_ref_logprobs`, `dpo_rejected_ref_logprobs`.197 198**Parameters**199 200| Name | Type | Default | Meaning |201|---|---|---|---|202| `tokenizer` | HF `AutoTokenizer` (duck-typed) | — | Must support `apply_chat_template` and `__call__`. |203| `device` | `torch.device \| str` | `"cpu"` | Target device for all returned tensors. |204| `seed` | `int` | `42` | Fixes `torch.manual_seed`. |205| `variant` | `str` | `"factorial"` | One of `"factorial"`, `"binary_search"`. |206| `align_sdpo_shapes` | `bool` | `False` | If True, truncate/pad `ctx_teacher_input_ids` to `input_ids` length so the SDPO channel actually fires. |207 208**Raises** `ValueError` if `variant` is unknown.209 210```python211from transformers import AutoTokenizer212from composer_replication import build_batch213tok = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct")214batch = build_batch(tok, variant="factorial", align_sdpo_shapes=True)215print({k: v.shape for k, v in batch.items()})216```217 218---219 220## 4. `composer_replication.opsd`221 222Self-distillation generalized-JSD loss, lifted verbatim from `siyan-zhao/OPSD` (MIT) per ADR-006.223 224### `generalized_jsd_loss(student_logits, teacher_logits, labels=None, beta=0.5, ...) -> torch.Tensor`225 226```python227def generalized_jsd_loss(228 student_logits: torch.Tensor,229 teacher_logits: torch.Tensor,230 labels: torch.Tensor | None = None,231 beta: float = 0.5,232 temperature: float = 1.0,233 reduction: str = "batchmean",234 logits_are_probs: bool = False,235 top_k: int | None = None,236 token_clip: float | None = None,237) -> torch.Tensor238```239 240Generalized JSD between student and teacher distributions. Same model on different contexts in the SDPO recipe; student and teacher params come from the SAME model.241 242**Parameters**243 244| Name | Type | Default | Meaning |245|---|---|---|---|246| `student_logits` | `Tensor (B, T, V)` | — | Student logits with grad. |247| `teacher_logits` | `Tensor (B, T, V)` | — | Teacher logits (no grad in SDPO). |248| `labels` | `Tensor (B, T) \| None` | `None` | Per-token mask. `-100` positions are ignored (HF convention). |249| `beta` | `float` in [0, 1] | `0.5` | 0=fwd KL, 1=rev KL, 0.5=symmetric JSD. |250| `temperature` | `float` | `1.0` | Softmax temperature. |251| `reduction` | `str` | `"batchmean"` | `"batchmean"`, `"sum"`, `"mean"`, `"none"`. |252| `logits_are_probs` | `bool` | `False` | Skip softmax if inputs are already probabilities. |253| `top_k` | `int \| None` | `None` | Restrict KL to teacher's top-k tokens. |254| `token_clip` | `float \| None` | `None` | Clip per-token JSD for stability. |255 256**Returns** scalar tensor (or `(B, T)` if `reduction="none"`).257 258**Raises** `ValueError` for unknown `reduction`.259 260```python261import torch262from composer_replication.opsd import generalized_jsd_loss263s = torch.randn(2, 8, 32, requires_grad=True)264t = torch.randn(2, 8, 32)265loss = generalized_jsd_loss(s, t, beta=0.5, reduction="batchmean")266loss.backward()267```268 269---270 271## 5. `composer_replication.distillation`272 273Pluggable self-distillation losses (ADR-007). All pure PyTorch.274 275### `simpo_loss(chosen_avg_logprobs, rejected_avg_logprobs, *, beta=2.0, gamma=1.0) -> torch.Tensor`276 277```python278def simpo_loss(279 chosen_avg_logprobs: torch.Tensor,280 rejected_avg_logprobs: torch.Tensor,281 *,282 beta: float = 2.0,283 gamma: float = 1.0,284) -> torch.Tensor285```286 287Reference-free DPO with target margin γ (Meng et al., NeurIPS 2024). `L = -log σ(β · (avg_logπ(c) − avg_logπ(r)) − γ)`.288 289**Parameters**290 291| Name | Type | Default | Meaning |292|---|---|---|---|293| `chosen_avg_logprobs` | `Tensor (B,)` | — | Per-sequence avg logprob over chosen response tokens. |294| `rejected_avg_logprobs` | `Tensor (B,)` | — | Same for rejected. |295| `beta` | `float` | `2.0` | Scaling factor (paper default). |296| `gamma` | `float` | `1.0` | Target margin (paper default). |297 298**Returns** scalar; **Raises** `ValueError` if shapes mismatch.299 300```python301import torch302from composer_replication.distillation import simpo_loss303loss = simpo_loss(torch.tensor([-2.1, -1.8]), torch.tensor([-3.0, -2.5]),304 beta=2.0, gamma=1.0)305```306 307### `avg_sequence_logprob(model_logprobs, response_mask) -> torch.Tensor`308 309⚠️ UNTESTED-CONTRACT (helper exported from `simpo.py` but not asserted by a test).310 311```python312def avg_sequence_logprob(313 model_logprobs: torch.Tensor,314 response_mask: torch.Tensor,315) -> torch.Tensor316```317 318Convert `(B, T)` per-token logprobs + `(B, T)` response mask into `(B,)` per-sequence average over response tokens.319 320```python321from composer_replication.distillation.simpo import avg_sequence_logprob322import torch323lp = torch.randn(2, 8); m = torch.tensor([[0,0,1,1,1,0,0,0],[0,1,1,1,1,1,0,0]])324out = avg_sequence_logprob(lp, m) # shape (2,)325```326 327### `taid_loss(student_logits, teacher_logits, mask=None, *, t) -> torch.Tensor`328 329```python330def taid_loss(331 student_logits: torch.Tensor,332 teacher_logits: torch.Tensor,333 mask: torch.Tensor | None = None,334 *,335 t: float | torch.Tensor,336) -> torch.Tensor337```338 339Faithful port of `SakanaAI/TAID` (arXiv:2501.16937). Forward-KL distillation against a logit-space-interpolated target whose anchor is the **current student detached**:340 341```342p_t = softmax( (1 - t) · stop_grad(student_logits) + t · teacher_logits )343L = - mean_token Σ_v p_t(v) · log_softmax(student_logits)(v)344```345 346At `t=0` the target collapses to the detached student (no teacher signal in the gradient). At `t=1` it reduces to standard forward-KL distillation against the teacher.347 348**Wave 15 breaking change.** The previous signature `taid_loss(student, teacher, student_init, *, schedule_step, total_steps, schedule, alpha_min, alpha_max, jsd_beta, temperature, reduction)` was algorithmically wrong (probability-space mix, frozen step-0 anchor, JSD criterion). All those kwargs are removed; the schedule is now the caller's responsibility (see `TAIDScheduler` below for the upstream adaptive scheme).349 350**Parameters**351 352| Name | Type | Default | Meaning |353|---|---|---|---|354| `student_logits` | `Tensor (B, T, V)` | — | Current student (with grad). |355| `teacher_logits` | `Tensor (B, T, V)` | — | Teacher logits. |356| `mask` | `Tensor (B, T) \| None` | `None` | Token mask. `None` ⇒ all-ones. |357| `t` | `float \| Tensor` | — | Interpolation coefficient in `[0, 1]`. |358 359**Raises** `ValueError` for shape mismatch.360 361```python362from composer_replication.distillation import taid_loss363loss = taid_loss(s_logits, t_logits, mask, t=0.4)364```365 366### `TAIDScheduler(num_train_steps, *, t_start=0.4, t_end=1.0, alpha=5e-4, beta=0.99, disable_adaptive=False)`367 368Stateful schedule that mirrors upstream `TAID.update_t`. Monotone non-decreasing, bumped above the linear floor by an EMA on the relative loss change. Use as:369 370```python371from composer_replication.distillation import TAIDScheduler372 373sched = TAIDScheduler(num_train_steps=10_000) # paper defaults374for step in range(num_train_steps):375 loss = taid_loss(s, t, mask, t=sched.t)376 loss.backward(); optimizer.step()377 sched.update_t(loss.detach(), global_step=step)378```379 380**Parameters**381 382| Name | Type | Default | Meaning |383|---|---|---|---|384| `num_train_steps` | `int` | — | Total planned training steps; sets the linear floor. |385| `t_start` | `float` | `0.4` | Initial `t` (paper default). |386| `t_end` | `float` | `1.0` | Terminal `t`; hard ceiling at every step. |387| `alpha` | `float` | `5e-4` | Adaptive bump magnitude. |388| `beta` | `float` | `0.99` | EMA decay on relative-loss-change momentum. |389| `disable_adaptive` | `bool` | `False` | If True, fall back to deterministic linear schedule. |390| `device` | `torch.device \| str` | `"cpu"` | Where to allocate state buffers. |391 392**Properties / methods**393 394- `sched.t -> float` — current `t` as a Python float (zero-arg property).395- `sched.update_t(loss, global_step) -> Tensor | None` — update internal state. First finite-loss call only seeds `prev_loss` and returns `None`; subsequent calls return the (positive) `delta_t` added on top of the linear floor.396 397### `entropy_aware_opd_loss(student_logits, teacher_logits, *, labels=None, h_max=None, temperature=1.0, reduction="batchmean") -> torch.Tensor`398 399```python400def entropy_aware_opd_loss(401 student_logits: torch.Tensor,402 teacher_logits: torch.Tensor,403 *,404 labels: torch.Tensor | None = None,405 h_max: float | None = None,406 temperature: float = 1.0,407 reduction: str = "batchmean",408) -> torch.Tensor409```410 411Per-token mixture of forward and reverse KL gated by teacher entropy: `w(t) = clamp(H_teacher(t)/h_max, 0, 1)`. High-entropy tokens use forward KL (mode-covering), low-entropy tokens use reverse KL (mode-seeking).412 413**Parameters**414 415| Name | Type | Default | Meaning |416|---|---|---|---|417| `student_logits` | `Tensor (B,T,V)` | — | Student logits (grad). |418| `teacher_logits` | `Tensor (B,T,V)` | — | Teacher logits (no grad). |419| `labels` | `Tensor (B,T) \| None` | `None` | 0/1 mask, applied multiplicatively after the per-token mix. |420| `h_max` | `float \| None` | `None` ⇒ `log(V)` | Max-entropy normalizer. |421| `temperature` | `float` | `1.0` | Softmax temperature on both. |422| `reduction` | `str` | `"batchmean"` | `"batchmean"`, `"sum"`, `"mean"`, `"none"`. |423 424**Raises** `ValueError` on shape mismatch (student vs teacher; labels vs per-token loss) or unknown `reduction`.425 426```python427from composer_replication.distillation import entropy_aware_opd_loss428loss = entropy_aware_opd_loss(s_logits, t_logits, temperature=1.0)429loss.backward()430```431 432### `teacher_entropy(teacher_logits) -> torch.Tensor`433 434⚠️ UNTESTED-CONTRACT (helper exposed from `entropy_aware_opd.py`'s `__all__` but not directly asserted).435 436Per-token entropy in nats. Input `(B,T,V)`, output `(B,T)`.437 438```python439from composer_replication.distillation.entropy_aware_opd import teacher_entropy440H = teacher_entropy(teacher_logits) # (B, T)441```442 443---444 445## 6. `composer_replication.teacher_replay`446 447N-teacher OpenRouter parallel client + DPO-pair extractor. `httpx` is lazy-imported inside `replay_trace`; the deterministic local logic is testable without it.448 449### `DEFAULT_TEACHERS: list[TeacherSpec]`450 451Three-teacher default set: `anthropic/claude-opus-4.7`, `openai/gpt-5`, `deepseek/deepseek-v4-pro` with paper-baseline OpenRouter pricing.452 453```python454from composer_replication.teacher_replay import DEFAULT_TEACHERS455print([t["slug"] for t in DEFAULT_TEACHERS])456```457 458### `class TeacherSpec(TypedDict)`459 460```python461class TeacherSpec(TypedDict):462 slug: str463 input_per_mtok: float464 output_per_mtok: float465```466 467OpenRouter model slug + per-million-token pricing.468 469```python470spec: TeacherSpec = {"slug": "openai/gpt-5",471 "input_per_mtok": 1.25, "output_per_mtok": 10.0}472```473 474### `class TraceState(TypedDict)`475 476```python477class TraceState(TypedDict):478 state_id: str # unique within the trace479 messages: list[dict] # OpenAI-style chat history up to (and incl.) this user prompt480 student_action: str # what the student actually did at this step481```482 483One step of a frozen agentic trace. `student_action` is the raw text emitted by the student; teachers are queried with `messages` and asked to predict the assistant's next action.484 485```python486state: TraceState = {"state_id": "ex001::0042",487 "messages": [{"role": "user", "content": "..."}],488 "student_action": "[TOOL_USE] name=Read input={...}"}489```490 491### `class TeacherCallResult(TypedDict)`492 493```python494class TeacherCallResult(TypedDict):495 state_id: str496 teacher_slug: str497 response_text: str | None # None on error498 latency_s: float499 prompt_tokens: int500 completion_tokens: int501 cost_usd: float502 error: str | None # None on success503```504 505One row of N×T results from `replay_trace`.506 507```python508r: TeacherCallResult = {"state_id": "x", "teacher_slug": "openai/gpt-5",509 "response_text": "ok", "latency_s": 1.2, "prompt_tokens": 100,510 "completion_tokens": 5, "cost_usd": 0.001, "error": None}511```512 513### `class DPOPair(TypedDict)`514 515```python516class DPOPair(TypedDict):517 state_id: str518 state_messages: list[dict]519 chosen: str # teacher-consensus action520 rejected: str # student action521 n_teachers_agreeing: int522```523 524One preference pair extracted from teacher-vs-student disagreement.525 526```python527p: DPOPair = {"state_id": "x", "state_messages": [...], "chosen": "...",528 "rejected": "...", "n_teachers_agreeing": 2}529```530 531### `async replay_trace(states, teachers=DEFAULT_TEACHERS, max_total_usd=5.0, api_key=None) -> list[TeacherCallResult]`532 533```python534async def replay_trace(535 states: Sequence[TraceState],536 teachers: Sequence[TeacherSpec] = tuple(DEFAULT_TEACHERS),537 max_total_usd: float = 5.0,538 api_key: str | None = None,539) -> list[TeacherCallResult]540```541 542For each state, fan-out one parallel call per teacher via OpenRouter. Hard-caps cumulative spend at `max_total_usd` (stops after the offending state completes).543 544**Parameters**545 546| Name | Type | Default | Meaning |547|---|---|---|---|548| `states` | `Sequence[TraceState]` | — | Frozen trace, one entry per assistant turn. |549| `teachers` | `Sequence[TeacherSpec]` | `DEFAULT_TEACHERS` | Models to query in parallel. |550| `max_total_usd` | `float` | `5.0` | Cumulative spend cap. |551| `api_key` | `str \| None` | `None` | OpenRouter key; defaults to `OPENROUTER_API_KEY` env or `~/.hermes/.env`. |552 553**Returns** flat list of `TeacherCallResult`s (length `len(states) * len(teachers)` modulo budget cutoff).554 555**Raises** `RuntimeError` if `OPENROUTER_API_KEY` is not findable; `ImportError` if `httpx` is missing at call time.556 557```python558import asyncio559from composer_replication import replay_trace560results = asyncio.run(replay_trace(states=my_trace, max_total_usd=1.0))561```562 563### `extract_dpo_pairs(states, teacher_actions, agreement_threshold=2) -> list[DPOPair]`564 565```python566def extract_dpo_pairs(567 states: Sequence[TraceState],568 teacher_actions: Sequence[TeacherCallResult],569 agreement_threshold: int = 2,570) -> list[DPOPair]571```572 573Group teacher_actions by `state_id`, normalize whitespace, and emit one `DPOPair` per state where ≥`agreement_threshold` teachers agreed on an action that differs from the student's. `chosen` is the original (un-normalized) teacher response text.574 575**Parameters**576 577| Name | Type | Default | Meaning |578|---|---|---|---|579| `states` | `Sequence[TraceState]` | — | Same as passed to `replay_trace`. |580| `teacher_actions` | `Sequence[TeacherCallResult]` | — | Output of `replay_trace`. |581| `agreement_threshold` | `int` | `2` | Min teachers that must agree for a pair to fire. |582 583**Returns** list of `DPOPair`. At most one pair per state (the most-agreed-upon action wins).584 585```python586from composer_replication import extract_dpo_pairs587pairs = extract_dpo_pairs(my_states, results, agreement_threshold=2)588```589 590### `save_pairs(pairs, path) -> None`591 592⚠️ UNTESTED-CONTRACT.593 594```python595def save_pairs(pairs: Sequence[DPOPair], path: str | Path) -> None596```597 598Write pairs to JSONL (one dict per line). Creates parent dirs.599 600```python601from composer_replication.teacher_replay import save_pairs602save_pairs(pairs, "/tmp/dpo_pairs.jsonl")603```604 605---606 607## 7. `composer_replication.replaysim`608 609ADR-004 normalization layer over `teacher_replay`. Re-exports `DPOPair`, `TeacherCallResult`, `extract_dpo_pairs`, `replay_trace` from `teacher_replay`.610 611### `class NormalizedDPOPair`612 613```python614@dataclass615class NormalizedDPOPair:616 state_id: str617 state_messages: list[dict[str, Any]]618 chosen_messages: list[dict[str, Any]]619 rejected_messages: list[dict[str, Any]]620 n_teachers_agreeing: int621 metadata: dict[str, Any]622```623 624Post-normalization shape. `chosen_messages`/`rejected_messages` are chat-format (`[{"role": "assistant", "content": ...}]`). `metadata` carries op-graph provenance, including `{"skipped": True}` when the normalizer was bypassed (`skip_dj=True`).625 626```python627from composer_replication.replaysim import NormalizedDPOPair628n = NormalizedDPOPair(state_id="x", state_messages=[],629 chosen_messages=[{"role": "assistant", "content": "ok"}],630 rejected_messages=[{"role": "assistant", "content": "no"}],631 n_teachers_agreeing=2, metadata={})632```633 634### `class DJNormalizer`635 636```python637class DJNormalizer:638 DEFAULT_RECIPE: ClassVar[Path] # composer_replication/recipes/replaysim/default.yaml639 640 def __init__(641 self,642 recipe_path: str | os.PathLike[str] | None = None,643 *,644 skip_dj: bool = False,645 ) -> None: ...646 647 def normalize(648 self,649 pairs: Iterable[DPOPair | dict[str, Any]],650 ) -> list[NormalizedDPOPair]: ...651```652 653`data-juicer`-backed normalizer. Pipeline: each `DPOPair` → JSONL record → `data_juicer.core.DefaultExecutor.run()` against the recipe → JSONL → `NormalizedDPOPair`.654 655**Constructor parameters**656 657| Name | Type | Default | Meaning |658|---|---|---|---|659| `recipe_path` | `str \| PathLike \| None` | `None` ⇒ default recipe | data-juicer YAML recipe path. |660| `skip_dj` | `bool` (kw-only) | `False` | If True: passthrough; records get `metadata={"skipped": True}` and no ops run. |661 662**`normalize(pairs) -> list[NormalizedDPOPair]`** runs the op-graph. Output may be shorter than input if filter ops drop records.663 664**Raises** `RuntimeError` at construction time if `skip_dj=False` and `data_juicer` is not importable. `FileNotFoundError` if `recipe_path` (default or explicit) is missing and `skip_dj=False`.665 666```python667from composer_replication.replaysim import DJNormalizer668norm = DJNormalizer(skip_dj=True)669out = norm.normalize(my_pairs)670```671 672### `async replay_and_normalize_trace(*, states, teachers=None, agreement_threshold=2, max_total_usd=5.0, normalizer=None, **replay_kwargs) -> tuple[list[TeacherCallResult], list[NormalizedDPOPair]]`673 674```python675async def replay_and_normalize_trace(676 *,677 states: Any,678 teachers: Any = None,679 agreement_threshold: int = 2,680 max_total_usd: float = 5.0,681 normalizer: DJNormalizer | None = None,682 **replay_kwargs: Any,683) -> tuple[list[TeacherCallResult], list[NormalizedDPOPair]]684```685 686End-to-end async: replay → extract pairs → normalize.687 688**Parameters**689 690| Name | Type | Default | Meaning |691|---|---|---|---|692| `states` | `Sequence[TraceState]` | — | Frozen trace. |693| `teachers` | `Sequence[TeacherSpec] \| None` | `None` ⇒ defaults | Forwarded to `replay_trace`. |694| `agreement_threshold` | `int` | `2` | Forwarded to `extract_dpo_pairs`. |695| `max_total_usd` | `float` | `5.0` | Spend cap. |696| `normalizer` | `DJNormalizer \| None` | `None` ⇒ `DJNormalizer()` | Pass `DJNormalizer(skip_dj=True)` to bypass. |697| `**replay_kwargs` | `Any` | — | Forwarded to `replay_trace` (e.g. `api_key`). |698 699**Returns** `(raw_teacher_actions, normalized_pairs)`.700 701```python702import asyncio703from composer_replication.replaysim import replay_and_normalize_trace, DJNormalizer704raw, norm = asyncio.run(replay_and_normalize_trace(705 states=my_states, normalizer=DJNormalizer(skip_dj=True)))706```707 708### `replay_and_normalize_trace_sync(*args, **kwargs) -> tuple[list[TeacherCallResult], list[NormalizedDPOPair]]`709 710⚠️ UNTESTED-CONTRACT (sync wrapper around the async function; tests call the async form via `asyncio.run`).711 712```python713def replay_and_normalize_trace_sync(*args, **kwargs) -> ...714```715 716Sync convenience wrapping `asyncio.run(replay_and_normalize_trace(...))`.717 718```python719from composer_replication.replaysim.normalize import replay_and_normalize_trace_sync720raw, norm = replay_and_normalize_trace_sync(states=my_states)721```722 723---724 725## 8. `composer_replication.ingestion` & `composer_replication.ingestion.claude_code`726 727Trace-source adapters (ADR-002). v0.1 supports Claude Code session JSONL.728 729### `SYSTEM_PROMPT: str`730 731Default synthetic system prompt injected at `messages[0]` for ingested traces (most Claude Code sessions don't write one). Truncated head: `"You are a senior software engineer working as a coding agent in a terminal environment..."`.732 733```python734from composer_replication import SYSTEM_PROMPT735print(SYSTEM_PROMPT[:60])736```737 738### `class IngestionStats`739 740```python741@dataclass742class IngestionStats:743 n_records_total: int = 0744 n_records_skipped: int = 0745 n_states_emitted: int = 0746 n_assistant_turns: int = 0747 n_tool_use_blocks: int = 0748 n_text_blocks: int = 0749 skipped_subagent: int = 0750 skipped_summary: int = 0751 skipped_truncated_lines: int = 0752 version_warnings: list[str] | None = None # initialized to [] in __post_init__753```754 755Counters populated by `ClaudeCodeIngester.ingest()` and exposed as `ingester.last_stats`.756 757```python758from composer_replication import IngestionStats759s = IngestionStats(n_records_total=5)760print(s.version_warnings) # []761```762 763### `class ClaudeCodeIngester`764 765```python766class ClaudeCodeIngester:767 def __init__(768 self,769 *,770 system_prompt: str = SYSTEM_PROMPT,771 skip_sidechain: bool = True,772 strip_thinking: bool = True,773 max_history_tokens: int | None = None,774 ) -> None: ...775 776 def ingest(self, path: Path) -> Iterator[TraceState]: ...777```778 779Convert a Claude Code session JSONL to a stream of `TraceState`s — one per assistant TURN (not per `tool_use` block).780 781**Constructor parameters**782 783| Name | Type | Default | Meaning |784|---|---|---|---|785| `system_prompt` | `str` | `SYSTEM_PROMPT` | Synthetic system message injected at history[0]. |786| `skip_sidechain` | `bool` | `True` | Skip subagent files (`agent-*.jsonl`) and records with `isSidechain=True`. |787| `strip_thinking` | `bool` | `True` | Remove `[THINKING]` blocks from history handed to teachers (kept inside `student_action`). |788| `max_history_tokens` | `int \| None` | `None` | ⚠️ UNTESTED-CONTRACT — accepted but currently not used to truncate. |789 790**`ingest(path) -> Iterator[TraceState]`**: generator over `TraceState` objects. Each turn's `state_id` is `f"{path.stem}::{idx:04d}"`. Side effect: replaces `self.last_stats` with a fresh `IngestionStats` and updates it as records stream.791 792```python793from pathlib import Path794from composer_replication import ClaudeCodeIngester795ing = ClaudeCodeIngester()796for state in ing.ingest(Path("session.jsonl")):797 print(state["state_id"])798print(ing.last_stats.n_states_emitted)799```800 801---802 803## 9. `composer_replication.hint_generator`804 805⚠️ UNTESTED-CONTRACT (entire module — used by the data collator config but not pinned by a test).806 807Template-based hint registry for SDPO error-site injection.808 809### `class HintContext(TypedDict, total=False)`810 811```python812class HintContext(TypedDict, total=False):813 error_kind: str814 error_message: str815 available_tools: list[str]816 tool_name: str817 tool_schema: dict818 intent: str819```820 821Per-error context dict consumed by hint templates.822 823### `HINT_TEMPLATES: dict[str, Callable[[HintContext], str]]`824 825Default registry keys: `"tool_not_found"`, `"json_decode"`, `"type_error"`, `"runtime_error"`, `"repeated_failure"`.826 827### `dispatch(error_kind, ctx=None) -> str | None`828 829```python830def dispatch(error_kind: str, ctx: HintContext | None = None) -> str | None831```832 833Look up `error_kind` in `HINT_TEMPLATES`. Returns the template's hint text, or `None` if the kind is unknown.834 835```python836from composer_replication.hint_generator import dispatch837hint = dispatch("json_decode") # "Reminder: tool arguments must be valid JSON. ..."838```839 840### `register(error_kind, fn) -> None`841 842```python843def register(error_kind: str, fn: Callable[[HintContext], str]) -> None844```845 846Add or override a custom hint template.847 848```python849from composer_replication.hint_generator import register850register("my_error", lambda ctx: "Reminder: try X.")851```852 853### Individual template functions854 855⚠️ UNTESTED-CONTRACT — exported only via `HINT_TEMPLATES`, useful as building blocks:856 857- `hint_tool_not_found(ctx) -> str`858- `hint_json_decode(ctx) -> str`859- `hint_type_error(ctx) -> str`860- `hint_runtime_error(ctx) -> str`861- `hint_repeated_failure(ctx) -> str`862 863Each accepts a `HintContext` and returns hint text. Signatures are uniform: `Callable[[HintContext], str]`.864 865```python866from composer_replication.hint_generator import hint_tool_not_found867text = hint_tool_not_found({"available_tools": ["Read", "Write"]})868```869 870---871 872## 10. `composer_replication.trainer` & sub-modules873 874Production trainer (TRL `GRPOTrainer` subclass) plus data collator.875 876### `class ComposerReplicationTrainer`877 878```python879class ComposerReplicationTrainer(GRPOTrainer):880 def __init__(881 self,882 *args: Any,883 alpha_sdpo: float = 0.1,884 beta_replay: float = 0.05,885 sdpo_jsd_beta: float = 0.5,886 sdpo_temperature: float = 1.0,887 sdpo_token_clip: float | None = None,888 replay_dpo_beta: float = 0.1,889 **kwargs: Any,890 ) -> None: ...891 892 def _compute_loss(893 self,894 model: torch.nn.Module,895 inputs: dict[str, torch.Tensor],896 ) -> torch.Tensor: ...897```898 899`trl.GRPOTrainer` subclass that overrides `_compute_loss(model, inputs)` to compose `total = grpo + α·sdpo + β·trace_replay_dpo`. When `trl` is not installed, the parent class falls back to `object` so the module imports — but instantiation will fail because the parent's GRPO machinery is missing.900 901**Constructor (kw-only beyond GRPOTrainer's own `*args, **kwargs`)**902 903| Name | Type | Default | Meaning |904|---|---|---|---|905| `alpha_sdpo` | `float` | `0.1` | Channel-2 weight. |906| `beta_replay` | `float` | `0.05` | Channel-3 weight. |907| `sdpo_jsd_beta` | `float` | `0.5` | β for `generalized_jsd_loss`. |908| `sdpo_temperature` | `float` | `1.0` | SDPO softmax temperature. |909| `sdpo_token_clip` | `float \| None` | `None` | Per-token JSD clip. |910| `replay_dpo_beta` | `float` | `0.1` | DPO β. |911 912**`_compute_loss(model, inputs) -> torch.Tensor`** — overrides `GRPOTrainer._compute_loss`. Calls `super()._compute_loss` for channel 1, then `_compute_sdpo_loss` and `_compute_trace_replay_loss`, then composes. Logs per-channel components every `args.logging_steps` (default 50). **Raises** whatever `super()` raises (TRL-shaped errors).913 914**Internal methods (publicly accessible, exercised by spike tests)**915 916- ⚠️ UNTESTED-CONTRACT `_compute_sdpo_loss(model, inputs) -> torch.Tensor` — generalized-JSD between student forward and `ctx_teacher_input_ids` forward. Returns `0.0` (with grad) when `alpha_sdpo == 0`, the key is missing, or shapes mismatch. Logs a warning on shape mismatch.917- ⚠️ UNTESTED-CONTRACT `_compute_trace_replay_loss(model, inputs) -> torch.Tensor` — standard DPO over `dpo_chosen_*` and `dpo_rejected_*`, using precomputed `dpo_chosen_ref_logprobs` / `dpo_rejected_ref_logprobs`.918- ⚠️ UNTESTED-CONTRACT `@staticmethod _sequence_logprobs(model, input_ids, response_mask) -> torch.Tensor` — sum logprobs over response tokens; standard DPO accounting.919 920```python921from composer_replication import ComposerReplicationTrainer922trainer = ComposerReplicationTrainer(923 model=my_model, args=my_grpo_args, train_dataset=ds,924 data_collator=my_collator, alpha_sdpo=0.1, beta_replay=0.05,925)926# trainer.train() # uses overridden _compute_loss927```928 929### `class TraceTurn(TypedDict, total=False)` — `trainer.data_collator`930 931```python932class TraceTurn(TypedDict, total=False):933 role: str # "user" | "assistant" | "tool"934 content: str935 tool_call: dict | None936 tool_error: str | None937 error_meta: dict938```939 940One turn of an agentic trace as consumed by `ComposerDataCollator`.941 942### `class TraceExample(TypedDict, total=False)` — `trainer.data_collator`943 944```python945class TraceExample(TypedDict, total=False):946 trace_id: str947 turns: list[TraceTurn]948 final_reward: float949 dpo_pairs: list[dict] | None950```951 952One training example: `(turns, optional dpo_pairs)`. `dpo_pairs` shape matches `DPOPair`.953 954### `class TokenizerLike` — `trainer.data_collator`955 956⚠️ UNTESTED-CONTRACT (duck-typed protocol; used as a type hint).957 958```python959class TokenizerLike:960 pad_token_id: int961 def __call__(self, text: str | list[str], **kwargs: Any) -> dict[str, list]: ...962 def apply_chat_template(self, messages: list[dict], **kwargs: Any) -> str | list[int]: ...963```964 965Minimal protocol the collator needs. Compatible with HF `AutoTokenizer`.966 967### `class CollatorConfig` — `trainer.data_collator`968 969```python970@dataclass971class CollatorConfig:972 max_seq_len: int = 4096973 max_dpo_seq_len: int = 2048974 pad_token_id: int = 0975 ignore_index: int = -100976 enable_sdpo: bool = True977 hint_generator: Callable[[str, dict], str | None] | None = None978 enable_replay_dpo: bool = True979 rlvr_reward_key: str = "final_reward"980```981 982Tunables for `ComposerDataCollator`.983 984| Field | Default | Meaning |985|---|---|---|986| `max_seq_len` | `4096` | Truncation cap for student/teacher sequences. |987| `max_dpo_seq_len` | `2048` | Truncation cap for DPO chosen/rejected sequences. |988| `pad_token_id` | `0` | Padding token id. |989| `ignore_index` | `-100` | HF "ignore in loss" sentinel for SDPO mask. |990| `enable_sdpo` | `True` | Toggle channel-2 fields. |991| `hint_generator` | `Callable[[str, dict], str \| None] \| None` (`None`) | `(error_kind, error_meta) -> hint_text`. SDPO is no-op without this. |992| `enable_replay_dpo` | `True` | Toggle channel-3 fields. |993| `rlvr_reward_key` | `"final_reward"` | Key in `TraceExample` to read scalar reward. |994 995```python996from composer_replication.trainer.data_collator import CollatorConfig997cfg = CollatorConfig(max_seq_len=2048, hint_generator=my_dispatch)998```999 1000### `class ComposerDataCollator` — `trainer.data_collator`1001 1002```python1003@dataclass1004class ComposerDataCollator:1005 tokenizer: TokenizerLike1006 config: CollatorConfig = field(default_factory=CollatorConfig)1007 1008 def __call__(1009 self, batch: Sequence[TraceExample]1010 ) -> dict[str, torch.Tensor]: ...1011```1012 1013Build trainer-ready batches from raw traces + optional DPO pairs.1014 1015**Output dict keys** (tested in `spikes/005-integrated-trainer-skeleton/tests/test_data_collator.py`):1016 1017- Channel 1 (always): `input_ids`, `attention_mask`, `response_mask`, `rewards`.1018- Channel 2 (when `enable_sdpo=True` AND batch has at least one error site AND `hint_generator` is set): `ctx_teacher_input_ids`, `sdpo_loss_mask`.1019- Channel 3 (when `enable_replay_dpo=True` AND batch has at least one `dpo_pair`): `dpo_chosen_input_ids`, `dpo_chosen_response_mask`, `dpo_rejected_input_ids`, `dpo_rejected_response_mask`. (Reference logprobs are NOT computed here — the trainer does that pass.)1020 1021```python1022from composer_replication.trainer.data_collator import (1023 ComposerDataCollator, CollatorConfig)1024collator = ComposerDataCollator(tokenizer=tok, config=CollatorConfig())1025batch = collator([{"trace_id": "x", "turns": [...], "final_reward": 1.0}])1026```1027 1028---1029 1030## 11. `composer_replication.diloco`1031 1032DiLoCo outer-loop wrapper around `torchft.local_sgd.DiLoCo`. Optional dep — when `torchft` is missing the package re-export `composer_replication.make_diloco_outer_loop` is `None`.1033 1034### Module-level attributes1035 1036- `DiLoCo: Any` — `torchft.local_sgd.DiLoCo` if importable else `None`.1037- `Manager: Any` — `torchft.manager.Manager` if importable else `None`.1038- `_DummyWork: Any` — `torchft.work._DummyWork` if importable else `None`.1039- `_TORCHFT_AVAILABLE: bool` — whether the imports succeeded.1040 1041```python1042from composer_replication.diloco import _TORCHFT_AVAILABLE, DiLoCo1043```1044 1045### `make_diloco_outer_loop(manager, model_fragments, inner_optimizer, *, ...) -> torchft.local_sgd.DiLoCo`1046 1047```python1048def make_diloco_outer_loop(1049 manager: Any,1050 model_fragments: list[torch.nn.Module],1051 inner_optimizer: torch.optim.Optimizer,1052 *,1053 outer_lr: float = 0.7,1054 outer_momentum: float = 0.9,1055 nesterov: bool = True,1056 sync_every: int = 100,1057 fragment_sync_delay: int = 0,1058 fragment_update_alpha: float = 0.0,1059) -> Any1060```1061 1062Construct a `torchft.DiLoCo` configured with framework-default hyperparams (DiLoCo paper §3.2: `lr=0.7, momentum=0.9, Nesterov`).1063 1064**Parameters**1065 1066| Name | Type | Default | Meaning |1067|---|---|---|---|1068| `manager` | `torchft.Manager` (or duck-typed `MockManager`) | — | Provides `allreduce`, `should_commit`, `current_step`, `start_quorum`, etc. |1069| `model_fragments` | `list[torch.nn.Module]` | — | One module for vanilla DiLoCo; N modules for Streaming DiLoCo. |1070| `inner_optimizer` | `torch.optim.Optimizer` | — | Inner-step optimizer (steps every batch). |1071| `outer_lr` | `float` | `0.7` | Outer SGD lr. |1072| `outer_momentum` | `float` | `0.9` | Outer SGD momentum. |1073| `nesterov` | `bool` | `True` | Nesterov momentum on outer SGD. |1074| `sync_every` | `int` | `100` | Inner steps per outer round. |1075| `fragment_sync_delay` | `int` | `0` | 0 = vanilla; >0 = Streaming DiLoCo (requires CUDA streams). |1076| `fragment_update_alpha` | `float` | `0.0` | 0 = full replacement on sync; >0 = exponential mix. |1077 1078**Returns** a `torchft.local_sgd.DiLoCo` instance — usable as a context manager.1079 1080**Raises** `RuntimeError` if `torchft` is not installed.1081 1082```python1083import torch1084from composer_replication.diloco import make_diloco_outer_loop1085opt = torch.optim.AdamW(model.parameters(), lr=1e-5)1086outer = make_diloco_outer_loop(manager=mgr, model_fragments=[model],1087 inner_optimizer=opt, sync_every=100)1088with outer:1089 for _ in range(N):1090 opt.zero_grad(); loss.backward(); opt.step()1091```1092 1093---1094 1095## 12. `composer_replication.diloco.serverless`1096 1097ADR-005 serverless DiLoCo executors + object-store all-reduce.1098 1099### `class ReplicaHandle` — `serverless.executor`1100 1101```python1102@dataclass1103class ReplicaHandle:1104 rank: int1105 backend_name: str1106 metadata: dict[str, Any] = field(default_factory=dict)1107```1108 1109Opaque handle returned by `ServerlessExecutor.launch_replicas`. `metadata` is backend-specific.1110 1111```python1112from composer_replication.diloco.serverless import ReplicaHandle1113h = ReplicaHandle(rank=0, backend_name="local_process",1114 metadata={"pid": 12345})1115```1116 1117### `class ServerlessExecutor` (Protocol) — `serverless.executor`1118 1119```python1120@runtime_checkable1121class ServerlessExecutor(Protocol):1122 backend_name: str1123 supports_inter_replica_network: bool1124 1125 def launch_replicas(1126 self,1127 n_replicas: int,1128 entrypoint: str | Callable[..., Any],1129 entrypoint_args: Mapping[str, Any],1130 *,1131 gpu: str | None = None,1132 timeout: int = 3600,1133 ) -> list[ReplicaHandle]: ...1134 1135 def poll(self, handle: ReplicaHandle) -> str: ...1136 def stream_logs(self, handle: ReplicaHandle, *, n_lines: int = 200) -> str: ...1137 def cancel(self, handle: ReplicaHandle) -> None: ...1138 def collect(1139 self, handles: list[ReplicaHandle], *, timeout: int | None = None,1140 ) -> list[dict[str, Any]]: ...1141```1142 1143Structural protocol for serverless backends.1144 1145- `launch_replicas(...)` returns `list[ReplicaHandle]` of length `n_replicas` in rank order. `entrypoint` is either an importable module path (uses `main()`) or a `module.function` path or a `Callable` (Local executor only). `entrypoint_args` may include `rank_env` (default `"REPLICA_RANK"`).1146- `poll(handle) -> str`: one of `"pending"`, `"running"`, `"succeeded"`, `"failed"`, `"cancelled"`.1147- `stream_logs(handle, n_lines=200) -> str`: best-effort recent stdout/stderr.1148- `cancel(handle) -> None`: best-effort.1149- `collect(handles, timeout=None) -> list[dict]`: blocks; each result dict has `rank`, `status`, `exit_code`, `error` (and `result` from `LocalProcessExecutor`).1150 1151```python1152from composer_replication.diloco.serverless import ServerlessExecutor1153def supports(x: ServerlessExecutor) -> bool:1154 return isinstance(x, ServerlessExecutor) # runtime_checkable1155```1156 1157### `class LocalProcessExecutor` — `serverless.executor`1158 1159```python1160class LocalProcessExecutor:1161 backend_name = "local_process"1162 supports_inter_replica_network = True1163 1164 def __init__(self) -> None: ...1165 # implements ServerlessExecutor protocol1166```1167 1168Reference implementation using Python `multiprocessing` (`spawn` context). Used for tests, CI smokes, and local development with `file://` rendezvous.1169 1170`launch_replicas(...)`: emits a soft warning on `gpu != None` (local processes share whatever GPUs are visible). `metadata = {"pid": ..., "start_ts": ...}`.1171 1172```python1173from composer_replication.diloco.serverless import LocalProcessExecutor1174ex = LocalProcessExecutor()1175handles = ex.launch_replicas(1176 n_replicas=2,1177 entrypoint="composer_replication.diloco.serverless.replica_entrypoint",1178 entrypoint_args={"rendezvous_uri": "/tmp/run/", "world_size": 2,1179 "trainer_module": "my.trainer"},1180)1181results = ex.collect(handles, timeout=60)1182```1183 1184### `class ObjectStoreAllReduce` — `serverless.allreduce`1185 1186```python1187class ObjectStoreAllReduce:1188 def __init__(1189 self,1190 uri: str,1191 rank: int,1192 world_size: int,1193 *,1194 round_id: int | None = None,1195 timeout_s: float = 1800.0,1196 poll_interval_s: float = 1.0,1197 ) -> None: ...1198 1199 @property1200 def round_id(self) -> int: ...