"""GRPO + SDPO column wiring on Qwen2.5-0.5B-Instruct (CPU end-to-end). This is the sibling to `examples/gsm8k_grpo/` that demonstrates the **SDPO hint-distillation column** firing end-to-end on a real HuggingFace model, on CPU, without needing TRL's full GRPO training loop. Where `gsm8k_grpo/run.py` runs plain GRPO with `alpha_sdpo=0`, this script loads the same model and shows that: 1. `compose_loss(model, inputs, alpha_sdpo=0.5, ...)` produces a non-zero `sdpo_jsd` channel on a real HF causal LM, and 2. backward through that channel reaches model parameters with finite gradients, and 3. running 5 SGD steps with the SDPO column enabled reduces the channel-decomposed total loss. This is the smallest possible "real-model" SDPO-wiring proof — the hand-crafted hint contexts here are not realistic training data, they just exercise the SDPO code path. For production SDPO, use `ComposerReplicationTrainer` with a `ComposerDataCollator` that emits `ctx_teacher_input_ids` / `sdpo_loss_mask` columns from your real trace data (see `composer_replication.trainer.data_collator`). Usage: pip install -e ".[train]" python examples/gsm8k_grpo_with_sdpo/run.py Cross-references: - `composer_replication.compose_loss` — the loss-composition entrypoint - `docs/COMPOSER_RECIPE_MAPPING.md` — how SDPO maps to Cursor's Composer-2.5 hint-distillation - `docs/adrs/ADR-008-drgrpo-sdpo-live-channel.md` — SDPO design - `examples/gsm8k_grpo/run.py` — plain GRPO (no SDPO) sibling """ from __future__ import annotations import logging import random import sys import time from pathlib import Path import torch from transformers import AutoModelForCausalLM, AutoTokenizer from composer_replication import compose_loss # --------------------------------------------------------------------------- # Config # --------------------------------------------------------------------------- MODEL_REPO = "Qwen/Qwen2.5-0.5B-Instruct" N_STEPS = 5 B = 2 # batch size T = 32 # sequence length (small to keep CPU fast) LR = 1e-5 ALPHA_SDPO = 0.5 # SDPO column weight; large enough to dominate the signal BETA_REPLAY = 0.0 # DPO column off — this example focuses on SDPO OUTPUT_DIR = Path(__file__).resolve().parent / "output" OUTPUT_DIR.mkdir(parents=True, exist_ok=True) # --------------------------------------------------------------------------- # Tiny GSM8K-shaped fixture — we fabricate the chat strings so the model # sees realistic prose. The "hint" is the same prompt with a # "remember to verify your arithmetic" line inserted; that's what makes # the teacher context differ from the student context. # --------------------------------------------------------------------------- PROBLEMS = [ { "question": "Janet has 3 boxes with 4 apples each. How many apples total?", "gold": "12", }, { "question": "A train travels 60 miles in 2 hours. What's its average speed?", "gold": "30", }, ] SYS = "You are a math tutor. End with `#### N` where N is the answer." HINT = "Hint: re-check your arithmetic before giving the final answer." def _build_chat_messages(question: str, *, with_hint: bool) -> list[dict]: """Format a single example as chat messages. with_hint=True is the teacher context (hint inserted as an extra system turn). Returns the OpenAI-style messages list, ready for tokenizer.apply_chat_template. Verified 2026-05-26: Qwen2.5 uses ChatML markers (`<|im_start|>` / `<|im_end|>`), NOT `<|system|>` / `<|end|>`. Using `apply_chat_template` is the only safe way to format input — raw marker strings get tokenized as 11 punctuation tokens and the model sees nonsense. """ messages = [{"role": "system", "content": SYS}] if with_hint: # Two system turns: Qwen's chat template will format both with # <|im_start|>system / <|im_end|> markers. messages.append({"role": "system", "content": HINT}) messages.append({"role": "user", "content": question}) return messages def build_inputs(tokenizer) -> dict[str, torch.Tensor]: """Tokenize PROBLEMS into a compose_loss-shaped batch. Returns a dict with: - input_ids: (B, T) student rollouts (no hint), left-padded - response_mask: (B, T) 1 on the assistant-response area - ctx_teacher_input_ids: (B, T) hint-conditioned context, left-padded - sdpo_loss_mask: (B, T) 1 at the aligned post-prompt area SDPO requires student and teacher logits to align position-by-position over the loss mask. The student and teacher prompts have different prefix lengths (teacher is longer because of the inserted hint system turn), so we LEFT-pad both to T tokens — the right edge (the assistant generation marker) lines up across the batch and across student vs teacher. The SDPO mask covers the right-most ALIGN_LEN positions, all of which correspond to identical "post-prompt / assistant-response area" tokens in both contexts. This matches the alignment discipline the production `ComposerDataCollator` (composer_replication/trainer/data_collator.py) must enforce: the post-hint section must have identical token positions in student vs teacher, or `_compute_sdpo_loss` will detect a shape mismatch and skip the channel for that step. """ # ALIGN_LEN: how many right-most positions to use for the SDPO loss. # These positions correspond to the assistant-generation area, which # is identical (token-for-token) across student and teacher because # apply_chat_template appends the same `<|im_start|>assistant\n` # marker regardless of how many system turns came before. ALIGN_LEN = T // 2 # 16 of 32; same as response_mask back-half student_msg_lists = [_build_chat_messages(p["question"], with_hint=False) for p in PROBLEMS[:B]] teacher_msg_lists = [_build_chat_messages(p["question"], with_hint=True) for p in PROBLEMS[:B]] student_strs = [ tokenizer.apply_chat_template(m, tokenize=False, add_generation_prompt=True) for m in student_msg_lists ] teacher_strs = [ tokenizer.apply_chat_template(m, tokenize=False, add_generation_prompt=True) for m in teacher_msg_lists ] # LEFT-pad AND LEFT-truncate to T: temporarily flip both the # tokenizer's padding side and truncation side. This ensures the # right edge (the assistant generation marker) is preserved at # position T-1 regardless of whether the input is shorter than T # (gets left-padded) or longer than T (gets left-truncated, dropping # the leading system turns first). Without this, the default # right-truncation discards the assistant marker — which means the # SDPO mask covers tokens from the system prompt instead of the # assistant response area, and the channel computes JSD over # nonsense alignment. original_pad = tokenizer.padding_side original_trunc = tokenizer.truncation_side tokenizer.padding_side = "left" tokenizer.truncation_side = "left" try: s_tok = tokenizer( student_strs, max_length=T, truncation=True, padding="max_length", return_tensors="pt", ) t_tok = tokenizer( teacher_strs, max_length=T, truncation=True, padding="max_length", return_tensors="pt", ) finally: tokenizer.padding_side = original_pad tokenizer.truncation_side = original_trunc # response_mask: 1 on the right-most ALIGN_LEN tokens, 0 elsewhere # (left padding + prompt area). For both student and teacher these # positions cover the assistant-generation marker + any padding # that happens to fall there. Same indices apply to both because # of left-padding alignment. response_mask = torch.zeros(B, T, dtype=torch.long) response_mask[:, -ALIGN_LEN:] = 1 sdpo_loss_mask = response_mask.clone() return { "input_ids": s_tok["input_ids"], "response_mask": response_mask, "ctx_teacher_input_ids": t_tok["input_ids"], "sdpo_loss_mask": sdpo_loss_mask, } # --------------------------------------------------------------------------- # Main # --------------------------------------------------------------------------- def main() -> int: random.seed(42) torch.manual_seed(42) log_path = OUTPUT_DIR.parent / "run.log" logging.basicConfig( level=logging.INFO, format="%(asctime)s [%(levelname)s] %(name)s: %(message)s", handlers=[ logging.StreamHandler(sys.stdout), logging.FileHandler(log_path, mode="w"), ], ) log = logging.getLogger("gsm8k_grpo_with_sdpo") log.info("=" * 64) log.info("GRPO + SDPO + GSM8K + Qwen2.5-0.5B-Instruct (CPU)") log.info("=" * 64) log.info("[1/4] Loading model + tokenizer ...") t0 = time.time() tokenizer = AutoTokenizer.from_pretrained(MODEL_REPO) if tokenizer.pad_token_id is None: tokenizer.pad_token = tokenizer.eos_token model = AutoModelForCausalLM.from_pretrained(MODEL_REPO, torch_dtype=torch.float32) model.to("cpu") n_params = sum(p.numel() for p in model.parameters()) log.info(" loaded in %.1fs (%.3fB params)", time.time() - t0, n_params / 1e9) log.info("[2/4] Building hint-conditioned batch (B=%d, T=%d) ...", B, T) inputs = build_inputs(tokenizer) for k, v in inputs.items(): log.info(" %s: shape=%s, dtype=%s", k, tuple(v.shape), v.dtype) log.info("[3/4] Running %d SGD steps with alpha_sdpo=%.2f ...", N_STEPS, ALPHA_SDPO) optim = torch.optim.SGD(model.parameters(), lr=LR) history: list[dict[str, float]] = [] model.train() t0 = time.time() for step in range(N_STEPS): optim.zero_grad() out = compose_loss( model, inputs, alpha_sdpo=ALPHA_SDPO, beta_replay=BETA_REPLAY, ) out.total.backward() # Sanity: gradients are finite + non-zero gnorm = sum( p.grad.abs().sum().item() for p in model.parameters() if p.grad is not None ) optim.step() components = out.detached() components["grad_norm"] = gnorm history.append(components) log.info( " step %d/%d: total=%.4f lm_ce=%.4f sdpo_jsd=%.4f trace_replay_dpo=%.4f |grad|=%.2e", step + 1, N_STEPS, components["total"], components["lm_ce"], components["sdpo_jsd"], components["trace_replay_dpo"], gnorm, ) dt = time.time() - t0 log.info("Training complete in %.1fs (avg %.1fs/step)", dt, dt / N_STEPS) # ------------------------------------------------------------------ # Acceptance assertions — SDPO column actually fired # ------------------------------------------------------------------ log.info("[4/4] Verifying SDPO column wiring ...") # 1. SDPO channel was non-zero at every step (channel actually fired) sdpo_values = [h["sdpo_jsd"] for h in history] assert all(s > 0.0 for s in sdpo_values), ( f"SDPO column is identically zero — channel did not fire. " f"sdpo_jsd values: {sdpo_values}" ) log.info(" ✓ sdpo_jsd > 0 at every step (min=%.4f, max=%.4f)", min(sdpo_values), max(sdpo_values)) # 2. total != lm_ce at every step (SDPO actually contributed to total) diffs = [abs(h["total"] - h["lm_ce"]) for h in history] assert all(d > 1e-6 for d in diffs), ( f"total ≈ lm_ce at every step — SDPO contribution is negligible. " f"abs(total - lm_ce): {diffs}" ) log.info(" ✓ total != lm_ce at every step (min |diff|=%.4f, max=%.4f)", min(diffs), max(diffs)) # 3. Gradients were finite + non-zero throughout gnorms = [h["grad_norm"] for h in history] assert all(g > 0.0 for g in gnorms), ( f"Some steps had zero gradient norm: {gnorms}" ) import math assert all(math.isfinite(g) for g in gnorms), ( f"Some steps had non-finite gradient norm: {gnorms}" ) log.info(" ✓ |grad| > 0 and finite at every step (min=%.2e, max=%.2e)", min(gnorms), max(gnorms)) # ------------------------------------------------------------------ # Summary # ------------------------------------------------------------------ log.info("=" * 64) log.info("Summary") log.info("=" * 64) log.info(" steps: %d", N_STEPS) log.info(" alpha_sdpo: %.2f", ALPHA_SDPO) log.info(" beta_replay: %.2f", BETA_REPLAY) log.info(" model params: %.3fB", n_params / 1e9) log.info(" total step 1: %.4f", history[0]["total"]) log.info(" total step %d: %.4f", N_STEPS, history[-1]["total"]) log.info(" wall-clock: %.1fs", dt) log.info(" log file: %s", log_path) log.info("=" * 64) log.info("✅ SDPO column wiring verified end-to-end.") return 0 if __name__ == "__main__": sys.exit(main())