Codeseys/composer-replication-framework
0
1# INTEGRATION_RECIPES.md — Wiring the 3-channel composer loss into your RL stack2 3> **Status:** Wave 14 release reference. Supersedes the historical4> [`docs/INTEGRATION_ARCHITECTURE.md`](INTEGRATION_ARCHITECTURE.md) (Recipes5> A–D), which is retained as background reading for the original6> mechanism-level diagrams.7>8> **Companion docs:**9> - [`docs/USER_GUIDE.md`](USER_GUIDE.md) — narrative walk-through, sections 1–810> - [`docs/API_REFERENCE.md`](API_REFERENCE.md) — exact kwarg signatures11> - [`docs/TROUBLESHOOTING.md`](TROUBLESHOOTING.md) — error → fix index12> - [`docs/V3_SUBSTRATE_COVERAGE.md`](V3_SUBSTRATE_COVERAGE.md) — what each13> substrate covers14> - [`docs/adrs/ADR-006-rl-frameworks.md`](adrs/ADR-006-rl-frameworks.md) —15> why these five recipes and not others16 17This document is the canonical answer to **"how do I plug the 3-channel18composer loss into framework X?"** for the five frameworks the project19supports as of Wave 14:20 211. [TRL `GRPOTrainer` subclass](#recipe-1--trl-grpotrainer-subclass)222. [VeRL custom `adv_estimator` + DataProto extension](#recipe-2--verl-custom-adv_estimator--dataproto-extension)233. [PRIME-RL custom-loss config](#recipe-3--prime-rl-customlossconfig)244. [Serverless Decoupled DiLoCo (Modal / HF Jobs / SageMaker)](#recipe-4--serverless-decoupled-diloco)255. [Monarch actor mesh (TorchForge-style topology)](#recipe-5--monarch-actor-mesh)26 27Each recipe follows the same seven-part template:28 291. **When to use it** — decision criteria.302. **Install command** — which optional extras of `composer-replication`.313. **Minimum-viable Python script** — copy-pasteable, ≤ 60 lines.324. **Decoupled DiLoCo wiring** — how `ServerlessExecutor` +33 `ObjectStoreAllReduce` + `MockManager` layer on top.345. **Distillation-loss wiring** — how to switch DPO → SimPO and add TAID35 via `compose_loss(..., dpo_variant=..., sdpo_wrapper=...)` or the36 recipe's own loss-config field.376. **Cost ballpark** — GPU $/hr + API spend, sourced from38 [`docs/research/DILOCO_SERVERLESS_RECONNAISSANCE.md`](research/DILOCO_SERVERLESS_RECONNAISSANCE.md).397. **Known limitations as of Wave 14**.40 41A cross-recipe [comparison matrix](#comparison-matrix) closes the doc.42 43## TL;DR — the unified loss44 45For any of the five recipes, the v0.1 trainer step computes:46 47```48total_loss = grpo_loss49 + α * sdpo_kl_loss (channel 2 — Composer hint-distill;50 optional TAID or Entropy-OPD wrapper)51 + β * trace_replay_loss (channel 3 — N-teacher DPO;52 switchable to SimPO)53```54 55This is implemented once, in56[`composer_replication/loss.py::compose_loss`](../composer_replication/loss.py),57and re-used by every recipe via the kwargs documented in58[`API_REFERENCE.md`](API_REFERENCE.md). The full signature — including59all ADR-007 channel-2/3 knobs (`dpo_variant`, `sdpo_wrapper`, `taid_t`,60`simpo_beta`/`simpo_gamma`, `entropy_opd_h_max`, …) — is the61single source of truth in62[API_REFERENCE.md § `compose_loss`](API_REFERENCE.md#compose_loss).63The conceptual call shape is just:64 65```python66compose_loss(model, inputs, **kwargs) # see API_REFERENCE.md#compose_loss for full signature67```68 69All five recipes below either call `compose_loss` directly or call a70thin per-framework adapter that forwards these kwargs unchanged. Each71recipe's **§5 Distillation-loss wiring** documents the kwargs *that72recipe* uses by default and why; refer back to API_REFERENCE.md for73defaults, types, and which kwargs are mutually exclusive.74 75---76 77## Recipe 1 — TRL `GRPOTrainer` subclass78 79### 1. When to use it80 81This is the **default v0.0/v0.1 path** and the one we recommend for82~99% of users today. Pick TRL when:83 84- Your model fits on ≤ 32 GPUs (typically ≤ 70B-param FSDP).85- You already have a HuggingFace `model` + `tokenizer` + `datasets` flow.86- You want minimum integration cost — `ComposerReplicationTrainer` is a87 single subclass override of `_compute_loss` over `trl.GRPOTrainer`,88 no Ray, no actor mesh.89- You're doing single-host (one node, possibly multi-GPU FSDP) training.90 91Don't pick TRL when you need >100 B-param scale, when you must async-decouple92tool calls from the GPU loop, or when a Ray cluster is already in your stack93(in which case Recipe 2 is cheaper).94 95### 2. Install command96 97```bash98pip install -e ".[train,replaysim]"99```100 101The `train` extra pulls `trl>=0.12`, `peft`, `accelerate`, and `datasets`.102The `replaysim` extra pulls `data-juicer` for CPU-side DPO normalization103(channel 3 cleaning step). Add `[serverless]` if you also want Decoupled104DiLoCo (see step 4).105 106### 3. Minimum-viable Python script107 108```python109# train_trl.py — minimum viable Recipe 1110from datasets import load_dataset111from transformers import AutoModelForCausalLM, AutoTokenizer112from composer_replication import ComposerReplicationTrainer113 114MODEL_ID = "Qwen/Qwen2.5-0.5B-Instruct" # swap for 7B once it works115model = AutoModelForCausalLM.from_pretrained(MODEL_ID)116tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)117dataset = load_dataset("trl-lib/tldr", split="train[:512]")118 119def reward_length(completions, **_):120 return [-abs(len(c) - 64) for c in completions]121 122trainer = ComposerReplicationTrainer(123 model = model,124 processing_class = tokenizer,125 reward_funcs = [reward_length],126 train_dataset = dataset,127 # Composer extras (defaults shown):128 alpha_sdpo = 0.1,129 beta_replay = 0.05,130 sdpo_jsd_beta = 0.5,131 sdpo_temperature = 1.0,132 sdpo_token_clip = None,133 replay_dpo_beta = 0.1,134)135trainer.train()136```137 138Channels 2 and 3 **auto-disable per step** when their inputs aren't139present in the batch (e.g. batches with no error sites get140`sdpo_kl=0`). Set `alpha_sdpo=0` / `beta_replay=0` to disable globally141for ablations.142 143### 4. Decoupled DiLoCo wiring144 145`ComposerReplicationTrainer` is a single-process trainer. To run N146replicas of it under Decoupled DiLoCo, layer the serverless stack on the147outside: each replica runs the script above; `MockManager` stands in for148`torchft.Manager` on the inner loop and `ObjectStoreAllReduce` runs the149outer-loop pseudo-gradient exchange:150 151```python152# diloco_replica.py — what each of the N replicas runs153import os154from composer_replication.diloco import make_diloco_outer_loop155from composer_replication.diloco.serverless import (156 LocalProcessExecutor, ObjectStoreAllReduce, MockManager,157)158 159rendezvous = ObjectStoreAllReduce(160 uri = "s3://my-bucket/diloco-runs/run42/",161 world_size = 4,162 rank = int(os.environ["REPLICA_RANK"]),163)164manager = MockManager(allreduce=rendezvous)165# trainer.optimizer is the *inner* optimizer; the outer is built here:166outer = make_diloco_outer_loop(167 inner_optimizer = trainer.optimizer,168 manager = manager,169 sync_every_h = 500,170)171trainer.add_callback(outer.callback()) # syncs every H inner steps172trainer.train()173```174 175The driver process spins these up with any `ServerlessExecutor`:176 177```python178# Wave 14: ModalExecutor / HFJobsExecutor are skeletons (raise NotImplementedError);179# use LocalProcessExecutor for testing. Swap once the cloud backends land.180executor = LocalProcessExecutor()181handles = executor.launch_replicas(182 n_replicas = 4,183 entrypoint = "diloco_replica.py",184 entrypoint_args = {"rendezvous": rendezvous.uri,185 "rank_env": "REPLICA_RANK"},186)187result = executor.collect(handles, timeout=3600)188```189 190### 5. Distillation-loss wiring191 192`ComposerReplicationTrainer` exposes the new ADR-007 channels via the193shared `compose_loss` kwargs — pass them through `**kwargs` on the194trainer and they're forwarded to `compose_loss`:195 196```python197trainer = ComposerReplicationTrainer(198 model = model, processing_class = tokenizer,199 reward_funcs = [reward_length], train_dataset = dataset,200 # SimPO instead of DPO for channel 3:201 dpo_variant = "simpo",202 simpo_beta = 2.0,203 simpo_gamma = 1.0,204 # TAID for channel 2 (SakanaAI port; logit-space mix + forward-KL):205 sdpo_wrapper = "taid",206 taid_t = 0.4, # current TAID coeff in [0, 1];207 # drive from TAIDScheduler if you want208 # the paper's adaptive scheme209)210```211 212Or, equivalently, drop `entropy_opd` in for `taid` if you want213per-token entropy-gated forward/reverse KL instead of the214linear-blend interpolation. SimPO does **not** require reference215log-probs (channel 3 batches with `dpo_chosen_ref_logprobs` /216`dpo_rejected_ref_logprobs` set are silently ignored).217 218### 6. Cost ballpark219 220- **GPU**: single host, `g5.12xlarge` ($5.67/hr) or RunPod 4×A100-80GB221 (~$5–9/hr) gets you Qwen2.5-7B at moderate throughput. For Qwen2.5-72B222 you'll want 2–4× H100 — `p5.48xlarge` (~$98/hr on AWS, ~$25–30/hr on223 Lambda Cloud / RunPod community).224- **API**: channel 3 teacher replay via OpenRouter — verified225 ~$0.98/trace at 50 steps × 3 teachers (spike 001). For a 100-trace226 curriculum that's ~$100 in teacher tokens.227- **Storage**: negligible until you turn on DiLoCo (then see Recipe 4).228 229### 7. Known limitations as of Wave 14230 231- **Tool calls block the GPU.** TRL's rollout is synchronous; long232 tool-call latency idles the trainer. Async-decouple via Recipe 2/3/5233 if this matters.234- **No native multi-node.** TRL is single-process; multi-host scaling is235 via Decoupled DiLoCo (Recipe 4) on top, not via TRL itself.236- **vLLM weight sync is co-located** — no resharding between FSDP and TP.237 At 70B+ this becomes the bottleneck and you should move to Recipe 2.238- **`reward_funcs` must be Python callables** that return `list[float]`;239 shell-out reward graders need a wrapper.240 241---242 243## Recipe 2 — VeRL custom `adv_estimator` + DataProto extension244 245### 1. When to use it246 247Pick VeRL when:248 249- You need >70B-param scale or >32-GPU multi-host, *and* a Ray cluster250 is acceptable in your stack.251- You're already using or willing to adopt **3D-HybridEngine** for252 efficient FSDP↔TP weight resharding (verified ~5× weight-sync speed-up253 vs co-located vLLM at 70B+).254- You need async multi-turn rollouts where tool-call latency must not255 block the GPU loop. VeRL's `AsyncServer` + `AgentLoop` is the256 best-in-class option here.257- You want extension points the framework's authors *expect* third258 parties to use — the `@register_adv_est("...")` decorator and the259 `DataProto` extension contract are first-class APIs.260 261Don't pick VeRL if you're <7B-param or single-host (overkill —262Recipe 1's Trainer subclass is one file, not a Ray cluster).263 264### 2. Install command265 266```bash267pip install -e ".[replaysim]"268pip install verl # not packaged as an extra; pinned at >=0.3269# Optional, for the Composer adapter:270pip install -e ".[serverless]" # for Decoupled DiLoCo on top271```272 273The framework's verl adapter lives at274`composer_replication.recipes.verl` (currently shape-only — see275[Limitations](#7-known-limitations-as-of-wave-14-2) below).276 277### 3. Minimum-viable Python script278 279VeRL's actual entry point is a Hydra/YAML config + `verl.trainer.main_ppo`280CLI; the pythonic surface looks like this:281 282```python283# train_verl.py — minimum viable Recipe 2 sketch284from verl.trainer.ppo import core_algos285from verl.trainer.ppo.ray_trainer import RayPPOTrainer286from composer_replication.loss import compose_loss287 288@core_algos.register_adv_est("grpo_composer")289def composer_advantage(data, **kwargs):290 """Custom adv-estimator that adds SDPO + DPO channels to GRPO.291 292 Reads three extra DataProto keys (populated by the data prep step):293 - data.batch["sdpo_teacher_logits"] (channel 2)294 - data.non_tensor_batch["teacher_actions"] (channel 3)295 and returns the standard (advantages, returns) tuple plus a stashed296 composer-loss term consumed by the critic worker.297 """298 advantages, returns = core_algos.compute_grpo_outcome_advantage(data, **kwargs)299 composer_term = compose_loss(300 model = kwargs["actor_module"],301 inputs = data.batch,302 alpha_sdpo = 0.1,303 beta_replay = 0.05,304 dpo_variant = "dpo",305 sdpo_wrapper = "none",306 )307 data.meta_info["composer_loss"] = composer_term308 return advantages, returns309 310# Then in your YAML:311# algorithm:312# adv_estimator: grpo_composer313# and run: python -m verl.trainer.main_ppo --config-name composer_grpo314```315 316The full driver wires `RayPPOTrainer` against your config; consult VeRL's317own quickstart for the Ray-cluster boilerplate. The composer-specific318piece is just the registered estimator above.319 320### 4. Decoupled DiLoCo wiring321 322VeRL's actor workers run in Ray; DiLoCo replicates the **whole VeRL job**.323Each "replica" is one Ray cluster running Recipe 2 end-to-end; the outer324loop is independent of Ray and just exchanges pseudo-gradients via the325object store between Ray-job invocations:326 327```python328from composer_replication.diloco.serverless import (329 LocalProcessExecutor, ObjectStoreAllReduce,330)331 332rendezvous = ObjectStoreAllReduce(333 uri = "s3://verl-diloco/run/",334 world_size = 4,335)336executor = LocalProcessExecutor() # Wave 14: ModalExecutor is a skeleton (raises NotImplementedError) — keep LocalProcessExecutor for now337handles = executor.launch_replicas(338 n_replicas = 4,339 entrypoint = "verl.trainer.main_ppo",340 entrypoint_args = {341 "+algorithm.adv_estimator": "grpo_composer",342 "+algorithm.diloco.rendezvous": rendezvous.uri,343 "+algorithm.diloco.sync_every_h": 500,344 },345)346executor.collect(handles, timeout=24 * 3600)347```348 349The Ray cluster inside each replica handles intra-replica scaling350(FSDP / TP / vLLM); the object-store exchange handles cross-replica351sync. Bandwidth is identical to Recipe 1 (~2 GB / 30 min per replica352for a 7B-param model in bf16) and well within S3 free-tier.353 354### 5. Distillation-loss wiring355 356The custom `adv_estimator` from step 3 already calls `compose_loss`;357flip the kwargs there to switch DPO → SimPO or add TAID:358 359```python360composer_term = compose_loss(361 model = kwargs["actor_module"],362 inputs = data.batch,363 alpha_sdpo = 0.1,364 beta_replay = 0.05,365 dpo_variant = "simpo", # ← SimPO swap366 simpo_beta = 2.0,367 simpo_gamma = 1.0,368 sdpo_wrapper = "taid", # ← TAID wrap369 taid_schedule_step = data.meta_info.get("global_step", 0),370 taid_total_steps = 10_000,371)372```373 374VeRL's `data.meta_info` carries the global step automatically, which is375exactly what TAID's interpolation schedule needs. Channel 2 batches376without `student_init_logits` / `student_init_input_ids` are auto-skipped377(returns 0 for that step).378 379### 6. Cost ballpark380 381- **GPU**: 8× H100 (`p5.48xlarge` ~$98/hr on AWS, ~$25/hr on Lambda or382 RunPod community) is the entry point for 70B-class. Expect 32–256383 H100 for full 671B (matches DeepSeek's reported VeRL config).384- **API**: same ~$0.98/trace as Recipe 1 (channel 3 is a Python helper,385 not a VeRL primitive — costs are framework-independent).386- **Ray cluster overhead**: head node + redis + dashboard adds ~1387 CPU-instance ($0.10–0.50/hr) per cluster, negligible at GPU scale.388 389### 7. Known limitations as of Wave 14390 391- **`composer_replication.recipes.verl` is shape-only.** The decorator392 registration and DataProto extension are documented but not yet shipped393 as a runnable adapter — Wave 14 release exposes the *contract*, not the394 glue. Expect this to land in a v0.2 follow-up spike.395- **Ray dependency.** Adds a heavyweight runtime; debugging396 cross-actor crashes can be painful. Use VeRL's `--debug` mode early.397- **Custom-`adv_estimator` LOC**: writing your own takes ~50–150 LOC398 including DataProto plumbing. Not a one-liner.399- **No first-class TAID hook in VeRL itself** — we route TAID through400 the meta_info channel; this works but means you can't use VeRL's401 built-in checkpoint-replay tooling without re-stamping `taid_schedule_step`402 on each replay.403 404---405 406## Recipe 3 — PRIME-RL `CustomLossConfig`407 408### 1. When to use it409 410Pick PRIME-RL when:411 412- You're operating in the **PRIME-Intellect / decentralized training**413 universe and want INTELLECT-style scaling on a long-horizon training414 run.415- You need **DPPO importance-ratio masking** (the rationale most users416 arrive with) — PRIME-RL's headline contribution is the417 out-of-band-token *mask* (not clip) on `log_ratio = trainer_lp -418 inference_lp`, with defaults `low=-4.0, high=4.0`.419- You want a **first-class custom-loss surface**: PRIME-RL ships420 `CustomLossConfig` that takes an importable Python function and a421 `LossInputs` struct exposing exactly the tensors we need422 (`trainer_logprobs`, `inference_logprobs`, `teacher_logprobs`,423 `advantages`, `loss_mask`). No fork, no Trainer subclass, no monkey-patch.424- You have access to multi-node infrastructure that PRIME-RL's425 trainer/inference/orchestrator split is designed for.426 427Don't pick PRIME-RL if you need full vocab logits (channel 2 SDPO428requires logits not log-probs — see Limitations).429 430### 2. Install command431 432```bash433pip install -e ".[prime-rl,replaysim]"434# pulls prime-rl>=0.5435```436 437### 3. Minimum-viable Python script438 439PRIME-RL drives via YAML config; the only Python you write is the440custom-loss function (already shipped at441`composer_replication/recipes/prime_rl/composer_loss.py`). Wire it in:442 443```yaml444# prime_rl_config.yaml — point at the framework's adapter445loss:446 custom:447 import_path: composer_replication.recipes.prime_rl.composer_loss:loss_fn448 kwargs:449 alpha_sdpo: 0.0 # channel 2 deferred in v0 (see below)450 beta_dpo: 0.0 # channel 3 emits a warning if non-zero451 dppo_mask_high: 4.0 # PRIME-RL DPPO mask bounds452 dppo_mask_low: -4.0453 epsilon: 1.0e-6454 455trainer:456 model: Qwen/Qwen2.5-7B-Instruct457 ... # standard PRIME-RL fields458```459 460The shipped `loss_fn` signature is fixed by PRIME-RL's contract:461 462```python463def loss_fn(464 inputs: LossInputs,465 *,466 alpha_sdpo: float = 0.0,467 beta_dpo: float = 0.0,468 dppo_mask_high: float = 4.0,469 dppo_mask_low: float = -4.0,470 epsilon: float = 1e-6,471) -> torch.Tensor:472 log_ratio = inputs.trainer_logprobs - inputs.inference_logprobs473 dppo_invalid = (log_ratio > dppo_mask_high) | (log_ratio < dppo_mask_low)474 keep_mask = inputs.loss_mask & ~dppo_invalid475 grpo = -(inputs.advantages * inputs.trainer_logprobs * keep_mask).sum() \476 / keep_mask.sum().clamp_min(epsilon)477 if alpha_sdpo != 0.0:478 raise NotImplementedError(479 "Channel 2 SDPO requires full-vocab logits; PRIME-RL v0.5 "480 "exposes only log-probs. Deferred to v0.2."481 )482 if beta_dpo != 0.0:483 import warnings; warnings.warn(484 "Channel 3 trace-replay DPO is out-of-scope for PRIME-RL recipe v0",485 stacklevel=2,486 )487 return grpo488```489 490**Shape note** (caught in the Wave 13 cross-model review): PRIME-RL491calls the loss function **once per sample**; tensors are 1-D `(seq,)`,492*not* batched `(B, T)`. The 10 unit tests in493`composer_replication/recipes/prime_rl/tests/test_composer_loss.py`494cover this plus DPPO mask edges.495 496### 4. Decoupled DiLoCo wiring497 498PRIME-RL was designed for decentralized training and ships its own499weight-sync primitives. Stack DiLoCo on top via the500`ServerlessExecutor` Protocol — each replica runs an independent501PRIME-RL job pointing at the same `composer_loss:loss_fn`:502 503```python504from composer_replication.diloco.serverless import (505 LocalProcessExecutor, ObjectStoreAllReduce,506)507 508rendezvous = ObjectStoreAllReduce(509 uri = "s3://prime-rl-diloco/run/",510 world_size = 4,511)512# Wave 14: ModalExecutor is a skeleton (raises NotImplementedError until v0.x).513# Use LocalProcessExecutor for the inner-replica wiring; swap to the cloud514# executor once it lands. The DiLoCo + rendezvous code below is identical.515executor = LocalProcessExecutor()516handles = executor.launch_replicas(517 n_replicas = 4,518 entrypoint = "prime_rl.cli:main",519 entrypoint_args = {520 "config": "prime_rl_config.yaml",521 "+diloco.rendezvous": rendezvous.uri,522 "+diloco.sync_every_h": 500,523 },524)525executor.collect(handles, timeout=24 * 3600)526```527 528Note PRIME-RL's own multi-node story (the trainer / inference /529orchestrator split) is **orthogonal** to Decoupled DiLoCo: PRIME-RL530multi-node = single replica scaled across many GPUs; DiLoCo = N531independent replicas synchronizing via object store. Combine both for532"big PRIME-RL job × N replicas".533 534### 5. Distillation-loss wiring535 536Channel 2 (SDPO + TAID + Entropy-OPD) is **deferred** in v0 because537PRIME-RL's `LossInputs` exposes log-probs not full vocab logits. The538SimPO swap on channel 3 is also gated by the same shape constraint, but539DPPO-clip itself doesn't change. To get TAID/SimPO into a PRIME-RL job540today you must:541 5421. Switch to Recipe 1 or 2 for the SFT/distill phase.5432. Use PRIME-RL only for the on-policy GRPO+DPPO phase.544 545The v0.2 plan (per ADR-007) is to extend `LossInputs` with a546`teacher_logits` field; the loss adapter is already shape-ready.547 548### 6. Cost ballpark549 550- **GPU**: similar profile to Recipe 2 — 8–32 H100 typical, scales to551 hundreds for INTELLECT-class runs. Lambda Cloud or RunPod community552 H100 community pricing (~$2–4/hr per H100) is most cost-effective.553- **API**: channel 3 is gated, so the only OpenRouter spend is from the554 *offline data-prep* spike (using the verifier harness in Recipe 1 to555 pre-bake DPO pairs), not from the training loop itself. Order of556 magnitude: $50–500 for a curriculum-bake one-time, then $0/run.557- **Network**: PRIME-RL's own decentralized weight sync uses substantial558 bandwidth between training replicas (one of its design constraints);559 this is *separate* from the Decoupled DiLoCo bandwidth and shows up560 as a ceiling on cross-region replica placement.561 562### 7. Known limitations as of Wave 14563 564- **Channel 2 deferred** — see step 5. `alpha_sdpo > 0` raises565 `NotImplementedError`.566- **Channel 3 emits a warning** if `beta_dpo != 0`; trace-replay DPO567 pairs must be folded into the *training data* (offline) rather than568 the *loss* (online) until v0.2.569- **PRIME-RL ≥ 0.5 required.** Earlier versions don't ship570 `CustomLossConfig`.571- **Smoke test deferred.** Per `prime_rl_recipe.md`, the runtime smoke572 test requires a CUDA box + `prime-rl >= 0.5` install and is gated573 to a follow-up spike. The 10 unit tests run cleanly without GPU.574- **DPPO defaults are PRIME-RL's, not ours.** We pin `low=-4.0,575 high=4.0` to match. If you change them, you're now diverging from576 PRIME-RL's example configs.577 578---579 580## Recipe 4 — Serverless Decoupled DiLoCo581 582### 1. When to use it583 584Pick Decoupled DiLoCo when:585 586- You have **N independent training replicas** that should sync587 occasionally but can't (or shouldn't) cross-talk on every step.588- The cost or operational burden of an always-on multi-node cluster is589 unacceptable, but you're happy paying for 4× independent **serverless590 jobs**.591- Your inner trainer is one of Recipes 1–3 — DiLoCo wraps any inner592 optimizer; it's *purely outer-loop*.593- You need **failure isolation**: if one replica crashes, the others594 keep training; on restart it picks up from the last outer round.595 596DiLoCo's design rests on two abstractions (per ADR-005):597 5981. **`ServerlessExecutor` Protocol** — uniform interface for spinning up599 N replicas across cloud backends (Modal / HF Jobs / SageMaker / k8s).6002. **`ObjectStoreAllReduce`** — fsspec-backed pseudo-gradient exchange601 that replaces the in-process `torchft.Manager.allreduce` call.602 603The communication pattern is `S3 PutObject + N GetObjects` once per604inner-H steps, matching DiLoCo paper §3.2 (arXiv:2311.08105). For6051B-param bf16 that's ~2 GB / 30 min per replica — well within S3606free-tier.607 608### 2. Install command609 610```bash611pip install -e ".[diloco,serverless]"612# also one of the inner-trainer extras:613pip install -e ".[train]" # if the inner trainer is Recipe 1614# OR pip install verl # if the inner trainer is Recipe 2615# OR pip install -e ".[prime-rl]" # if the inner trainer is Recipe 3616```617 618### 3. Minimum-viable Python script619 620This pattern is independent of the inner trainer — pick any of Recipes6211/2/3 and wrap it with a `ServerlessExecutor`. The replica entrypoint622runs the inner trainer; the driver launches N of them and waits.623 624```python625# diloco_driver.py — driver that launches N replicas626from composer_replication.diloco.serverless import (627 LocalProcessExecutor, # for dev — runs replicas as local subprocesses628 ObjectStoreAllReduce,629)630 631rendezvous = ObjectStoreAllReduce(632 uri = "s3://my-bucket/diloco-runs/run42/", # or file:// for local633 world_size = 4,634)635executor = LocalProcessExecutor() # Wave 14: ModalExecutor skeleton raises NotImplementedError; swap once cloud backend lands636handles = executor.launch_replicas(637 n_replicas = 4,638 entrypoint = "diloco_replica.py", # (script below)639 entrypoint_args = {640 "rendezvous": rendezvous.uri,641 "rank_env": "REPLICA_RANK",642 },643)644result = executor.collect(handles, timeout=3600)645print({h.replica_id: h.exit_code for h in result})646```647 648```python649# diloco_replica.py — runs inside each replica650import os651from composer_replication.diloco import make_diloco_outer_loop652from composer_replication.diloco.serverless import (653 ObjectStoreAllReduce, MockManager,654)655 656# Build inner trainer (Recipe 1 example):657from train_trl import trainer658 659rendezvous = ObjectStoreAllReduce(660 uri = os.environ["DILOCO_RENDEZVOUS"],661 world_size = 4,662 rank = int(os.environ["REPLICA_RANK"]),663)664manager = MockManager(allreduce=rendezvous)665outer = make_diloco_outer_loop(666 inner_optimizer = trainer.optimizer,667 manager = manager,668 sync_every_h = 500,669)670trainer.add_callback(outer.callback())671trainer.train()672```673 674### 4. Decoupled DiLoCo wiring675 676This recipe **is** the DiLoCo wiring — see step 3. The available677executor adapters are:678 679| Executor | Status | Use case |680|---------------------------|-------------------------------|--------------------------------------|681| `LocalProcessExecutor` | Production-ready | Dev loop — N subprocesses on one box |682| `ModalExecutor` | Skeleton (modal-client gated) | Modal cloud, $/sec billing |683| `HFJobsExecutor` | Skeleton (hf-hub gated) | HuggingFace Jobs, transformer-shop |684| `SageMakerExecutor` | Roadmap (post-v0.2) | AWS, warm-pool ~10s cold start |685| `K8sExecutor` | Roadmap | KubeRay / Volcano gang scheduling |686 687Cross-cloud replica placement (e.g. 2× Modal + 2× HF Jobs) is supported688in principle — they all read/write the same S3 / GCS / HF rendezvous —689but treat as experimental.690 691### 5. Distillation-loss wiring692 693DiLoCo is loss-agnostic — it operates purely on inner-optimizer state.694Whichever inner trainer you're running (Recipe 1, 2, or 3) handles695distillation kwargs as documented in that recipe's step 5. The only696DiLoCo-specific knob worth knowing: TAID's `taid_schedule_step` is a697*global* counter, but each replica increments it independently. If you698care about replicas all reading the same α at outer-sync time, set699`taid_schedule_step = trainer.state.global_step + replica_offset` and700let the outer-loop sync average them out.701 702### 6. Cost ballpark703 704Pulled from705[`docs/research/DILOCO_SERVERLESS_RECONNAISSANCE.md`](research/DILOCO_SERVERLESS_RECONNAISSANCE.md):706 707| Backend | A100-80GB $/hr | H100 $/hr | Cold-start | Notes |708|---------------|----------------|-----------|------------|------------------------------------------|709| Modal | $1.39/sec → 4× ≈ $20/hr per A100 | ~$8/hr per H100 | 1–60s warm, 60–120s first-run | $/sec billing; no minimum |710| AWS SageMaker | $4.10/A100·hr | $12.29/hr | 2–5 min cold, ~10s warm pool | Min 60min on warm pool |711| GCP Vertex | $3.67/A100·hr | $11/hr | 2–6 min cold | 30–50% premium over raw GPU |712| Azure ML | ~$3.67/A100·hr | ~$12.25/hr | 3–8 min cold | Use curated env to cut cold-start |713| RunPod | $1.19/hr (community), $2.17 (secure) | $1.99/hr (community), $4.18 (secure) | seconds | No federation; same-DC only |714| HF Jobs | comparable to Modal | ~$8–12/hr | 30–90s | Best DX for HF-shop |715 716**Object-store cost.** ~$0.02/GB-month for S3 standard, ~$0/free-tier.717Pseudo-gradients are ~2 GB per replica per outer round; for a 24-hour7184-replica run at H=500 that's ~50 outer rounds × 2 GB × 4 replicas = ~400719GB written. Free-tier blows through fast — budget $10–20 in storage.720 721### 7. Known limitations as of Wave 14722 723- **`ModalExecutor` and `HFJobsExecutor` are skeletons.** They check724 `import modal` / `import huggingface_hub` at *adapter init* time and725 raise; the actual `launch_replicas` is shape-only until the relevant726 spike lands. Use `LocalProcessExecutor` for dev.727- **`ObjectStoreAllReduce(world_size=1)`** must passthrough cleanly —728 the unit test `test_object_store_allreduce_world_size_1_passthrough`729 is the regression guard. Don't override unless you've read it.730- **Rank validation is mandatory.** Tests assert731 `ObjectStoreAllReduce(rank=N, world_size=N)` raises (rank must be732 `< world_size`); silent corruption otherwise.733- **`MockManager` is *not* feature-complete.** It implements the734 `Manager.allreduce` surface that DiLoCo's outer-loop needs, but735 not the full `torchft.Manager` API (no fault-tolerance, no736 membership protocol). Don't use it as a drop-in for live torchft.737- **No native heterogeneous compute** — all replicas are assumed to738 have the same compute shape. Mixed A100+H100 placements work but739 the slow replica gates outer-loop progress.740 741---742 743## Recipe 5 — Monarch actor mesh744 745### 1. When to use it746 747Pick Monarch when:748 749- You're at **TorchForge-style topology scale**: trainer / generator /750 rewarder / N-teachers all want to be independent, asynchronously751 scheduled, fault-tolerant actors on a typed mesh.752- You want **heterogeneous executor support** — different actors run753 in different clouds (e.g. `TrainerActor` on Modal A100s,754 `GeneratorActor` on dedicated H100s, `TeacherPoolActor` as 0-GPU CPU755 pods on k8s).756- You need **hot-swap of actor implementations** — replace757 "OpenRouter teachers" with "local vLLM teachers" by changing one758 Monarch binding, no trainer code change.759- You're prepared to track **upstream Monarch** (v0.4.1 stable, v0.5760 dev daily); the API is moving and v0 of this recipe is intentionally761 deferred per ADR-006.762 763Don't pick Monarch in Wave 14 unless you're explicitly scoping a764v0.2+ pilot. The framework ships *skeleton* actors that fail-fast on765instantiation; this is a reference-pattern reading exercise, not a766production target.767 768### 2. Install command769 770```bash771pip install -e ".[prime-rl,monarch]"772# pulls monarch>=0.4.1 plus the PRIME-RL trainer used inside actors773```774 775### 3. Minimum-viable Python script776 777The framework ships skeleton actor definitions at778`composer_replication/recipes/monarch/actors.py`; they raise779`NotImplementedError` on instantiation in Wave 14. The shape of the780final answer:781 782```python783# monarch_train.py — what v0.2+ usage will look like784from monarch import Actor, mesh, endpoint785from composer_replication.recipes.monarch.actors import (786 TrainerActor, GeneratorActor, RewarderActor, TeacherPoolActor,787)788 789# Topology790trainers = mesh.spawn(TrainerActor, n=4, gpu="A100")791generator = mesh.spawn(GeneratorActor, n=1, gpu="A100")792rewarder = mesh.spawn(RewarderActor, n=1, gpu=None)793teachers = mesh.spawn(TeacherPoolActor, n=1, gpu=None)794 795# Wire endpoints796async def outer_step(batch_id: int):797 prompts = await trainers[0].sample_prompts.call(batch_id)798 rollouts = await generator.rollout.call(prompts)799 rewards = await rewarder.score.call(rollouts)800 teacher_acts = await teachers.replay.call([801 {"state": r["state"]} for r in rollouts802 ])803 await trainers.train_outer_step.call(804 batch_id, rollouts=rollouts, rewards=rewards,805 teacher_actions=teacher_acts,806 )807 808# Run809import asyncio810for batch_id in range(1000):811 asyncio.run(outer_step(batch_id))812```813 814The Composer 3-channel loss lives inside `TrainerActor.train_outer_step`,815which calls `compose_loss(...)` exactly as Recipe 1 does. The816*orchestration* changes; the *loss math* doesn't.817 818### 4. Decoupled DiLoCo wiring819 820Monarch + Decoupled DiLoCo compose naturally: each `TrainerActor` is a821DiLoCo replica, and Monarch's supervision tree handles the failure822recovery that ADR-005 lists as a DiLoCo design constraint. The wire-up823is identical to Recipe 4's `LocalProcessExecutor` pattern, just running824inside Monarch instead of `subprocess`:825 826```python827from composer_replication.diloco.serverless import (828 ObjectStoreAllReduce, MockManager,829)830 831class TrainerActor(Actor):832 def __init__(self, rendezvous_uri: str, rank: int, world_size: int):833 self.rendezvous = ObjectStoreAllReduce(834 uri=rendezvous_uri, rank=rank, world_size=world_size,835 )836 self.manager = MockManager(allreduce=self.rendezvous)837 # ... build inner ComposerReplicationTrainer ...838 839 @endpoint840 async def train_outer_step(self, batch_id: int, **kw):841 # Inner H steps locally, then sync via self.rendezvous842 ...843```844 845The "object store" is the cross-actor synchronization point that846*doesn't* go through Monarch's RDMA data plane — by design, slow847syncs (S3) and fast syncs (RDMA for in-actor weight broadcast) live on848different planes.849 850### 5. Distillation-loss wiring851 852Monarch sees the loss as opaque: it lives inside `TrainerActor` and853takes the same `compose_loss` kwargs as Recipe 1. The mesh-level854benefit is **swap-by-binding**: you can replace `TeacherPoolActor`855("OpenRouter") with a `LocalVLLMTeacherActor` to switch the856*supplier* of teacher log-probs without touching the loss config.857 858```python859# Original binding — channel 3 via OpenRouter860teachers = mesh.spawn(TeacherPoolActor, n=1, gpu=None)861 862# Swap binding — channel 3 via local vLLM863teachers = mesh.spawn(LocalVLLMTeacherActor, n=1, gpu="A100",864 model_id="Qwen/Qwen2.5-72B-Instruct")865 866# Trainer config unchanged:867trainer.compose_loss_kwargs = dict(868 dpo_variant = "simpo", # same as before869 sdpo_wrapper = "taid",870 taid_schedule_step = batch_id,871 taid_total_steps = 10_000,872)873```874 875### 6. Cost ballpark876 877In Wave 14: $0 (skeleton fails fast; no compute used). Projected for v0.2+:878 879- **Mesh overhead**: Monarch's coordination plane is light — typically880 <1% of total compute even at 4-actor scale. The dominant cost is881 whatever the actors run.882- **Heterogeneous placement** is the cost lever: e.g. a 4-trainer mesh883 with `TeacherPoolActor` on 0-GPU CPU pods can cut total $/hr by884 ~10–20% vs forcing all actors onto GPU nodes.885- **Cluster bring-up**: Monarch v0.5's Slurm backend is stable; k8s886 backend is dev-track; bare-metal SSH backend is documented.887 888### 7. Known limitations as of Wave 14889 890- **Skeleton only, fails fast.** Importing `actors.py` is fine;891 instantiating `TrainerActor(...)` raises `NotImplementedError("v0892 skeleton; deferred to v0.2 per ADR-006")`. By design.893- **Upstream Monarch API is moving.** v0.4.1 stable + v0.5 dev daily894 means breaking changes are expected. Pin to a Monarch hash if you895 prototype.896- **TorchForge is paused.** Per its own repo banner — don't take897 TorchForge's recipes as production patterns. Monarch alone is898 active; Forge as a layered framework is reference reading.899- **Open question (deferred):** does Monarch v0.5's Slurm backend900 hand-shake cleanly with HF Jobs lifecycle? See901 `monarch_actor_layout.md` for the open-questions list.902- **Open question (deferred):** can `TrainerActor` host903 `ComposerReplicationTrainer` unmodified, or does it need a904 `step_init` / `step_compute` split for Monarch's async actor model?905 906---907 908## Comparison matrix909 910| Dimension | Recipe 1 — TRL | Recipe 2 — VeRL | Recipe 3 — PRIME-RL | Recipe 4 — Serverless DiLoCo | Recipe 5 — Monarch |911|------------------------------------|-----------------------------|----------------------------------|-----------------------------------|------------------------------------|-------------------------------------|912| **Maturity (Wave 14)** | Production-ready | Production-ready (adapter shape-only) | Recipe ready, runtime smoke deferred | `LocalProcessExecutor` ready; cloud adapters skeleton | Skeleton only; v0.2+ scope |913| **Supports DAPO / GRPO** | GRPO ✅; DAPO via TRL master | GRPO ✅; DAPO ✅ (built-in) | GRPO+DPPO ✅ (DAPO mask is the headline) | Inherits from inner trainer | Inherits from inner trainer |914| **Custom-loss extension cost (LOC)** | ~30 LOC (subclass override) | ~50–150 LOC (registered estimator) | ~20 LOC (single Python fn) | 0 (transparent wrapper) | ~30 LOC (loss inside actor) |915| **OpenEnv-compatible** | ✅ (HF datasets layer) | ✅ (DataProto extension) | ✅ (rollout JSONL contract) | ✅ (orthogonal) | ✅ (RewarderActor binding) |916| **Native multi-node** | ❌ (single-host FSDP only) | ✅ (Ray cluster + 3D-HybridEngine) | ✅ (trainer/inference/orchestrator split) | ✅ (the *whole point*) | ✅ (mesh of actors) |917| **Native Decoupled DiLoCo** | ❌ — wrap with Recipe 4 | ❌ — wrap with Recipe 4 | ❌ — wrap with Recipe 4 | ✅ (this *is* it) | ✅ (compose with Recipe 4 inside actor) |918| **License** | Apache 2.0 (TRL) | Apache 2.0 (VeRL) | Apache 2.0 (PRIME-RL) | Apache 2.0 (this repo) | BSD-3 (Monarch) |919| **Our recommendation (Wave 14)** | **Default for ≤ 70B / single-host** | Pick at >70B *if* Ray is acceptable | Pick if PRIME-Intellect / DPPO mask is required | Stack on top of 1/2/3 for N replicas | Reference pattern only — revisit v0.2 |920 921---922 923## Cross-recipe checklist924 925Regardless of which recipe you pick, these invariants are tested across926the test suite (266 passing / 62 skipped; canonical count in docs/V1_V8_COVERAGE.md) and should be true of your wired-up system:927 928- **`alpha_sdpo=0`** must reproduce the channel-1-only baseline929 bit-exact (`test_compose_loss_integration.py`).930- **`beta_replay=0`** must reproduce the no-channel-3 baseline931 bit-exact.932- **`sdpo_wrapper="taid"` without `taid_schedule_step`** must `ValueError`933 at first step (`test_compose_loss_integration.py`).934- **`sdpo_wrapper="taid"` at `taid_schedule_step / taid_total_steps = 0`**935 must ignore the teacher signal (`test_taid_loss_alpha_zero_ignores_teacher`).936- **`sdpo_wrapper="taid"` at `taid_schedule_step / taid_total_steps = 1`**937 must equal plain SDPO (`test_taid_blended_logits_endpoints`).938- **`dpo_variant="simpo"`** must be differentiable through the939 `loss-of-sigmoid` path (`test_simpo_loss_differentiable`).940- **`sdpo_wrapper="entropy_opd"`** must zero out when student ≡ teacher941 (`test_entropy_aware_opd_zero_when_distributions_match`).942- **`ObjectStoreAllReduce(world_size=1)`** must passthrough cleanly943 (`test_object_store_allreduce_world_size_1_passthrough`).944 945If any of these fail in your wired-up system, run the corresponding946unit test to localize: most break because a kwarg got dropped at the947adapter boundary, not because the loss math is wrong.948 949---950 951## Picking a recipe — decision flow952 9531. **Piloting Monarch (v0.2+)?** → Recipe 5.9542. **Else, need >70B / multi-host?** → Recipe 2 (VeRL) if Ray is OK,955 Recipe 3 (PRIME-RL) if you're in the PRIME-Intellect / DPPO universe,956 otherwise wait for Recipe 5.9573. **Else** → Recipe 1 (TRL) is the v0.0/v0.1 default.9584. **At any of 1–3, need N independent replicas / failure isolation?**959 → Stack Recipe 4 (Decoupled DiLoCo) on top.960 961---962 963## Pointers to source964 965- Loss core: [`composer_replication/loss.py`](../composer_replication/loss.py)966- TRL trainer: [`composer_replication/trainer/composer_trainer.py`](../composer_replication/trainer/composer_trainer.py)967- PRIME-RL adapter:968 [`composer_replication/recipes/prime_rl/composer_loss.py`](../composer_replication/recipes/prime_rl/composer_loss.py),969 recipe doc:970 [`composer_replication/recipes/prime_rl/prime_rl_recipe.md`](../composer_replication/recipes/prime_rl/prime_rl_recipe.md)971- Monarch skeleton:972 [`composer_replication/recipes/monarch/actors.py`](../composer_replication/recipes/monarch/actors.py),973 layout doc:974 [`composer_replication/recipes/monarch/monarch_actor_layout.md`](../composer_replication/recipes/monarch/monarch_actor_layout.md)975- Serverless DiLoCo:976 [`composer_replication/diloco/serverless/`](../composer_replication/diloco/serverless/)977- VeRL adapter (shape-only): `composer_replication/recipes/verl/`978- ADRs:979 [`docs/adrs/ADR-005-serverless-diloco.md`](adrs/ADR-005-serverless-diloco.md),980 [`docs/adrs/ADR-006-rl-frameworks.md`](adrs/ADR-006-rl-frameworks.md),981 [`docs/adrs/ADR-007-self-distillation-losses.md`](adrs/ADR-007-self-distillation-losses.md)982 983---984 985**File path:** `docs/INTEGRATION_RECIPES.md` (repo-relative)986