Codeseys/composer-replication-framework

1"""composer_replication — Composer 2.5 Replication Framework.2 3A research-grade, open replication of Cursor Composer 2.5's training recipe:4take any HuggingFace model, further-RL-train it using a 3-channel loss combining5 6    1. RLVR / GRPO (channel 1, via TRL)7    2. SDPO hint-distillation (channel 2, OPSD-based)8    3. Multi-teacher trace-replay DPO (channel 3, this framework's contribution)9 10with optional DiLoCo / Streaming DiLoCo outer-loop sync for distributed runs.11 12See https://huggingface.co/Codeseys/composer-replication-framework for the13full project README, design docs, ADRs, and verification spikes.14 15## Two API surfaces, on purpose16 17This package exposes BOTH a verification-harness API and a production-trainer18API. Use the right one for your purpose:19 20### Verification harness (small, easy to call, NOT for real training)21 22`compose_loss(model, batch, alpha_sdpo, beta_replay)` is a free function23that returns `LossComponents(lm_ce, sdpo_jsd, trace_replay_dpo, total)`.24It stubs the GRPO channel with LM cross-entropy on response tokens (the25limit GRPO converges to under deterministic rewards) so you can verify26the 3-channel composition wires together WITHOUT spinning up TRL's full27reward + advantage machinery.28 29`build_batch(tokenizer)` produces a real chat-template-formatted batch30with all keys `compose_loss` may consume.31 32Use these for:33- CPU smokes on real HF models (Spike 006 / Spike 002a-mini-gpu)34- Unit testing custom loss-composition variants35- Debugging gradient flow through one of the three channels36- Anything where you want to call backward() on a real model without37  spinning up TRL38 39### Production trainer (use for actual training runs)40 41`ComposerReplicationTrainer` is a `trl.GRPOTrainer` subclass that42overrides `_compute_loss(model, inputs)` to compose the same 3 channels43on top of TRL's real GRPO machinery. This is what you train models with.44 45Use this for:46- Real training runs on HF models with real rollouts + rewards47- Anything where the GRPO channel's policy-gradient signal matters48  (i.e., not a memorization smoke)49 50The verification harness's `compose_loss` is intentionally NOT a51drop-in replacement for `_compute_loss` — they target different52phases of the framework's lifecycle.53 54## Quickstart (verification-harness API)55 56    >>> from composer_replication import compose_loss, build_batch57    >>> from transformers import AutoModelForCausalLM, AutoTokenizer58    >>> model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct")59    >>> tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct")60    >>> batch = build_batch(tokenizer)61    >>> components = compose_loss(model, batch, alpha_sdpo=0.1, beta_replay=0.05)62    >>> components.total.backward()63 64See `examples/qwen_05b_quickstart/run.py` in the repo for a complete CPU65smoke (verification harness) and `spikes/002a-mini-gpu-smoke/run_gpu_smoke.py`66for a GPU smoke (verification harness, bf16, 50 steps).67 68For production-trainer usage, see `docs/INTEGRATION_ARCHITECTURE.md` Recipe A.69"""70from __future__ import annotations71 72# Loss composition (Spike 006)73from composer_replication.loss import LossComponents, compose_loss74from composer_replication.batch import build_batch75 76# Trace ingestion (Spike 007)77from composer_replication.ingestion.claude_code import (78    SYSTEM_PROMPT,79    ClaudeCodeIngester,80    IngestionStats,81)82 83# OPSD / SDPO loss (verified extension from siyan-zhao/OPSD, MIT)84from composer_replication.opsd import generalized_jsd_loss85 86# Teacher replay (Spike 001 → trainer)87from composer_replication.teacher_replay import (88    DEFAULT_TEACHERS,89    DPOPair,90    TeacherCallResult,91    TeacherSpec,92    TraceState,93    extract_dpo_pairs,94    replay_trace,95)96 97# Trainer (Spike 005) + policy-optimization config factories (ADR-008/ADR-014)98from composer_replication.trainer import (99    PO_OBJECTIVES,100    ComposerReplicationTrainer,101    make_dr_grpo_config,102    make_po_config,103)104 105# DiLoCo (Spike 008) — optional, requires torchft106try:107    from composer_replication.diloco import make_diloco_outer_loop108    _DILOCO_AVAILABLE = True109except ImportError:110    _DILOCO_AVAILABLE = False111    make_diloco_outer_loop = None  # type: ignore[assignment]112 113__version__ = "0.1.0"114 115__all__ = [116    # Core loss117    "compose_loss",118    "LossComponents",119    "build_batch",120    "generalized_jsd_loss",121    # Trace ingestion122    "ClaudeCodeIngester",123    "IngestionStats",124    "SYSTEM_PROMPT",125    "TraceState",126    # Teacher replay127    "DEFAULT_TEACHERS",128    "DPOPair",129    "TeacherCallResult",130    "TeacherSpec",131    "extract_dpo_pairs",132    "replay_trace",133    # Trainer134    "ComposerReplicationTrainer",135    "make_dr_grpo_config",136    "make_po_config",137    "PO_OBJECTIVES",138    # DiLoCo (optional)139    "make_diloco_outer_loop",140    # Meta141    "_DILOCO_AVAILABLE",142    "__version__",143]144