examples/gsm8k_grpo_with_sdpo/run.py

"""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())