"""Production-grade SDPO end-to-end on real Claude Code traces (CPU, ~2min). This is the FIFTH example in the SDPO progression — the production-grade sibling to `examples/sdpo_with_real_traces/`: examples/gsm8k_grpo/ -- plain GRPO baseline examples/gsm8k_grpo_with_sdpo/ -- SDPO on hand-crafted prompts examples/sdpo_with_real_traces/ -- SDPO WIRING smoke (misaligned) examples/sdpo_with_real_traces_production/ -- SDPO PRODUCTION-GRADE (this) Where `sdpo_with_real_traces` was a wiring-only smoke (HINT appended to messages → student/teacher right-edge tokens diverge → JSD measured on different content), THIS example uses the production path: ClaudeCodeIngester → claude_states_to_trace_examples() [Wave 19 NEW adapter] → ComposerDataCollator(hint_generator=...) → batch with PROPERLY-ALIGNED ctx_teacher_input_ids + sdpo_loss_mask → compose_loss The data collator's `_build_hint_injected_trace` walks the turns, detects error sites via `tool_error` markers, injects the hint as a system turn BEFORE the assistant recovery turn, and builds an `sdpo_loss_mask` that's 1 only at the post-hint assistant tokens (positions where student and teacher are predicting the SAME content). This example demonstrates: ✅ The full production data path: ingester → adapter → collator ✅ SDPO column firing on PROPERLY-ALIGNED student/teacher contexts ✅ Real tool error detection via the [TOOL_RESULT (ERROR)] tag flow ✅ A deterministic hint generator wired into CollatorConfig ✅ Gradient flow through Qwen2.5-0.5B-Instruct's params Closes the V5 gap end-to-end (the path is production-grade and content-honest, with a detailed hint at the actual error site of the trace), within the constraint that the trace fixture is hand-authored (PII reasons; users can point at their own JSONL). Usage: pip install -e ".[train]" python examples/sdpo_with_real_traces_production/run.py Cross-references: - composer_replication.ingestion.trace_examples.claude_states_to_trace_examples - composer_replication.trainer.data_collator.ComposerDataCollator - composer_replication.trainer.data_collator._build_hint_injected_trace - examples/sdpo_with_real_traces/ (the wiring-only sibling for comparison) """ from __future__ import annotations import logging import math import sys import time from pathlib import Path import torch from transformers import AutoModelForCausalLM, AutoTokenizer from composer_replication import compose_loss from composer_replication.ingestion import ( ClaudeCodeIngester, claude_states_to_trace_examples, ) from composer_replication.trainer.data_collator import ( CollatorConfig, ComposerDataCollator, ) # --------------------------------------------------------------------------- # Config # --------------------------------------------------------------------------- MODEL_REPO = "Qwen/Qwen2.5-0.5B-Instruct" N_STEPS = 5 LR = 1e-5 ALPHA_SDPO = 0.5 BETA_REPLAY = 0.0 MAX_SEQ_LEN = 1024 # generous; the with-error fixture is short OUTPUT_DIR = Path(__file__).resolve().parent / "output" OUTPUT_DIR.mkdir(parents=True, exist_ok=True) # This fixture is the WITH-ERROR variant — it has an `is_error: true` # tool_result that the adapter detects and the collator injects a hint # before. The clean Spike 007 fixture has no errors and would produce # a no-op SDPO batch. FIXTURE_PATH = ( Path(__file__).resolve().parents[2] / "spikes" / "007-real-trace-ingestion" / "fixtures" / "synthetic_session_with_error.jsonl" ) # --------------------------------------------------------------------------- # Hint generator — deterministic, error-kind-aware # --------------------------------------------------------------------------- def hint_for_error(error_kind: str, error_meta: dict) -> str | None: """Return a hint text given the classified error kind. A real production hint generator would pull from a curated hint library or an LLM-as-teacher; here we use a small static map for determinism. Returning None for an error kind tells the collator to skip the SDPO injection for that turn. """ library = { "file_not_found": ( "Hint: when reading a file fails with 'does not exist', " "first verify the path with `ls` on the parent directory " "or use a glob to find similar names before retrying." ), "permission_denied": ( "Hint: when 'permission denied', check ownership with `ls -l` " "before retrying. Don't blindly add `sudo`; read the situation." ), "command_not_found": ( "Hint: when a command isn't found, check `which ` " "and `echo $PATH`; the binary may need to be installed first." ), "tool_error": ( "Hint: this tool call failed. Read the error carefully and " "consider whether to retry, change inputs, or pivot to a " "different tool before continuing." ), } return library.get(error_kind, library["tool_error"]) # --------------------------------------------------------------------------- # Build batch via production path # --------------------------------------------------------------------------- def build_production_batch( tokenizer, fixture_path: Path, ) -> tuple[dict[str, torch.Tensor], int, int]: """Run the full production pipeline. Returns: (batch, n_states, n_error_sites) """ ingester = ClaudeCodeIngester(skip_sidechain=True, strip_thinking=True) states = list(ingester.ingest(fixture_path)) if not states: raise RuntimeError(f"No TraceState yielded from {fixture_path}") examples = claude_states_to_trace_examples(states) n_error_sites = sum( 1 for ex in examples for t in ex["turns"] if t.get("tool_error") ) config = CollatorConfig( hint_generator=hint_for_error, enable_replay_dpo=False, # this example focuses on SDPO pad_token_id=tokenizer.pad_token_id or 0, max_seq_len=MAX_SEQ_LEN, ) collator = ComposerDataCollator(tokenizer=tokenizer, config=config) batch = collator(examples) return batch, len(states), n_error_sites # --------------------------------------------------------------------------- # Main # --------------------------------------------------------------------------- def main() -> int: 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("sdpo_production") log.info("=" * 64) log.info("PRODUCTION-GRADE SDPO + ClaudeCodeIngester + ComposerDataCollator") log.info("Model: %s (CPU)", MODEL_REPO) log.info("=" * 64) if not FIXTURE_PATH.is_file(): log.error("Fixture not found at %s", FIXTURE_PATH) return 1 log.info("[1/5] Fixture: %s (size=%d bytes)", FIXTURE_PATH.name, FIXTURE_PATH.stat().st_size) log.info("[2/5] 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("[3/5] Building batch via production pipeline ...") log.info(" ClaudeCodeIngester → claude_states_to_trace_examples → ComposerDataCollator") batch, n_states, n_error_sites = build_production_batch(tokenizer, FIXTURE_PATH) log.info(" ingested %d states; adapter detected %d error site(s)", n_states, n_error_sites) if n_error_sites == 0: log.error(" No error sites detected — SDPO will be a no-op. " "Use the with-error fixture or extend the adapter.") return 1 for k, v in batch.items(): log.info(" %s: shape=%s dtype=%s", k, tuple(v.shape), v.dtype) if "ctx_teacher_input_ids" not in batch: log.error(" Collator did not produce ctx_teacher_input_ids — " "no error sites survived hint generator. Aborting.") return 1 sdpo_in_loss = (batch["sdpo_loss_mask"] == 1).sum().item() log.info(" sdpo_loss_mask: %d positions in loss (per-row: %s)", sdpo_in_loss, (batch["sdpo_loss_mask"] == 1).sum(dim=-1).tolist()) s_shape = batch["input_ids"].shape t_shape = batch["ctx_teacher_input_ids"].shape log.info(" shape reconciliation: student %s vs teacher %s — %s", tuple(s_shape), tuple(t_shape), "ALIGNED" if s_shape == t_shape else "MISMATCH (collator bug?)") assert s_shape == t_shape, ( f"Shape mismatch after collator: student {s_shape} vs teacher {t_shape}. " f"compose_loss requires student_logits.shape == teacher_logits.shape; " f"the collator's __call__ must reconcile them." ) log.info("[4/5] 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, batch, alpha_sdpo=ALPHA_SDPO, beta_replay=BETA_REPLAY, ) out.total.backward() 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) log.info("[5/5] Verifying production-grade SDPO behavior ...") sdpo_values = [h["sdpo_jsd"] for h in history] # Production-grade SDPO MUST produce a non-zero JSD signal because # the collator put the hint in a position where it actually changes # the teacher's prediction at the masked positions. assert all(abs(s) > 1e-7 for s in sdpo_values), ( f"Production-grade SDPO column produced negligible JSD: {sdpo_values}. " f"The hint isn't perturbing teacher logits at masked positions — " f"check the collator's hint injection or the loss mask." ) log.info(" ✓ sdpo_jsd > 1e-7 at every step (min=%.6f max=%.6f)", min(sdpo_values), max(sdpo_values)) # The composed total must differ from lm_ce alone — confirms SDPO contributes diffs = [abs(h["total"] - h["lm_ce"]) for h in history] assert all(d > 1e-6 for d in diffs), ( f"total ≈ lm_ce — SDPO contribution negligible. abs(total-lm_ce)={diffs}" ) log.info(" ✓ total != lm_ce at every step (min |diff|=%.4f)", min(diffs)) gnorms = [h["grad_norm"] for h in history] assert all(g > 0.0 and math.isfinite(g) for g in gnorms), ( f"Some grads non-finite or zero: {gnorms}" ) log.info(" ✓ |grad| finite at every step (min=%.2e max=%.2e)", min(gnorms), max(gnorms)) # ---------------------------------------------------------------- # Alignment audit (Wave 19 honesty: documents the residual drift) # ---------------------------------------------------------------- s_in = batch["input_ids"] t_in = batch["ctx_teacher_input_ids"] m_in = batch["sdpo_loss_mask"] n_aligned = 0 n_total_in_loss = 0 for row in range(s_in.shape[0]): in_loss = (m_in[row] == 1) n_pos = in_loss.sum().item() if n_pos == 0: continue s_at = s_in[row][in_loss] t_at = t_in[row][in_loss] n_aligned += int((s_at == t_at).sum().item()) n_total_in_loss += n_pos if n_total_in_loss: ratio = n_aligned / n_total_in_loss log.info(" alignment audit: %d / %d in-loss positions match student==teacher (%.1f%%)", n_aligned, n_total_in_loss, 100 * ratio) if ratio < 0.95: log.warning( " NOTE: %d positions (%.1f%%) of the SDPO mask cover non-aligned " "tokens. As of Wave 20 the chat-template drift was fixed via " "ComposerDataCollator._build_chat_aligned_mask (per-message " "apply_chat_template prefix deltas). A ratio below ~100%% now " "indicates a NEW regression — investigate the collator, not a " "known-residual bug.", n_total_in_loss - n_aligned, 100 * (1 - ratio), ) else: log.info( " ✓ Wave 20 chat-template alignment holding (%.1f%% — was ~67%% " "before the _build_chat_aligned_mask fix).", 100 * ratio, ) log.info("=" * 64) log.info("Summary") log.info("=" * 64) log.info(" trace fixture: %s", FIXTURE_PATH.name) log.info(" states: %d", n_states) log.info(" error sites: %d", n_error_sites) log.info(" sdpo_loss_mask: %d positions in loss", sdpo_in_loss) log.info(" alpha_sdpo: %.2f", ALPHA_SDPO) 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("=" * 64) log.info("✅ Production-grade SDPO verified end-to-end via ComposerDataCollator.") return 0 if __name__ == "__main__": sys.exit(main())