feat(trainer): ADR-008 Dr.GRPO config + SDPO strict-alignment guard

Phase-6 execution of ADR-008 (gates 1, 2, 5 green; gate 3 = the live
GRPOTrainer smoke runs detached separately).

- make_dr_grpo_config(): builds trl.GRPOConfig to the Dr. GRPO recipe
(loss_type=dr_grpo => no length-standardization bias; scale_rewards=none
=> no std-dev advantage normalization; num_iterations=1 => single-epoch),
with a drift-guard assertion. TRL 1.5.0 natively supports these knobs and
its own help text cites the Dr. GRPO paper for both.
- ComposerReplicationTrainer.strict_sdpo_alignment (default True): the
student/teacher logit shape-mismatch path now RAISES instead of silently
zeroing the SDPO channel (closes the ADR-008 trust-gap, composer_trainer.py
lines 158-160). Opt-out to warn-and-skip for production resilience.
- prime_rl/composer_loss.py: NotImplementedError(alpha_sdpo>0) message now
points at the TRL host as the SDPO home (documents the ADR-006 amendment).
- Tests: 7 new (config knobs asserted; strict raises / non-strict skips /
aligned fires / no-error-site no-ops). Full trainer + prime_rl suites:
28 passed, 1 skipped, no regressions.
- examples/composer_grpo_sdpo_smoke/: gate-3 live-loop smoke (detached-run).

Constraint note: cross-family ADR review + subagent execution blocked by
OpenRouter 402 (out of credits); executing inline. TRL import is ~140s
(heavy transitive deps) so the live GRPO smoke runs as a detached
systemd-scope job.

composer_replication/recipes/prime_rl/composer_loss.py

@@ -263,8 +263,10 @@ def loss_fn(
             "distribution. Set alpha_sdpo=0.0 to silence this and use "
             "channel 1 (DPPO+KL) only. teacher_logprobs is "
             f"{'present' if teacher_lp is not None else 'absent'} in this "
-            "call but unused. See docs/research/WAVE_13_FINAL_REVIEW.md "
-            "Finding 1."
+            "call but unused. For the SDPO channel, use the TRL host "
+            "(composer_replication.trainer.ComposerReplicationTrainer with "
+            "alpha_sdpo>0), which has full logits — see ADR-008. "
+            "See docs/research/WAVE_13_FINAL_REVIEW.md Finding 1."
         )
 
     # --- Channel 3: not supported in PRIME-RL recipe v0 -------------------

composer_replication/trainer/composer_trainer.py

@@ -74,6 +74,7 @@ class ComposerReplicationTrainer(GRPOTrainer):  # type: ignore[misc, valid-type]
         sdpo_temperature: float = 1.0,
         sdpo_token_clip: float | None = None,
         replay_dpo_beta: float = 0.1,
+        strict_sdpo_alignment: bool = True,
         **kwargs: Any,
     ):
         if not _TRL_AVAILABLE:
@@ -88,6 +89,12 @@ class ComposerReplicationTrainer(GRPOTrainer):  # type: ignore[misc, valid-type]
         self.sdpo_temperature = sdpo_temperature
         self.sdpo_token_clip = sdpo_token_clip
         self.replay_dpo_beta = replay_dpo_beta
+        # When True (default), an SDPO student/teacher shape mismatch is a hard
+        # error — it means the data collator failed to align the post-hint
+        # section, which silently zeroes the distillation signal (the exact
+        # trust-gap flagged in ADR-008). Set False only for production runs
+        # where a single malformed batch should warn-and-skip rather than abort.
+        self.strict_sdpo_alignment = strict_sdpo_alignment
 
     # ----------------------------------------------------------------------
     # Loss override (the integration core)
@@ -159,12 +166,21 @@ class ComposerReplicationTrainer(GRPOTrainer):  # type: ignore[misc, valid-type]
         # SAME LENGTH at the post-hint section so logits align position-by-position.
         # The data collator pads/aligns. The skeleton trusts that's done correctly.
         if student_logits.shape != teacher_logits.shape:
-            logger.warning(
-                "SDPO logit shape mismatch: student=%s vs teacher=%s. "
-                "Skipping SDPO loss for this step. Check the data collator's "
-                "alignment — the post-hint section must have identical token-counts.",
-                student_logits.shape, teacher_logits.shape,
+            msg = (
+                f"SDPO logit shape mismatch: student={tuple(student_logits.shape)} "
+                f"vs teacher={tuple(teacher_logits.shape)}. The data collator must "
+                "pad/align the post-hint section so student and teacher have "
+                "identical token-counts; otherwise the distillation signal is "
+                "silently zeroed."
             )
+            if self.strict_sdpo_alignment:
+                # Hard error (default): a mismatch means the collator failed to
+                # align — the ADR-008 trust-gap. Do not silently zero the channel.
+                raise ValueError(
+                    msg + " (strict_sdpo_alignment=True; pass False to warn-and-skip "
+                    "instead for production resilience.)"
+                )
+            logger.warning("%s Skipping SDPO loss for this step.", msg)
             return torch.tensor(0.0, device=_device_of(model), requires_grad=True)
 
         return generalized_jsd_loss(
@@ -246,4 +262,45 @@ def _device_of(model: torch.nn.Module) -> torch.device:
     return next(model.parameters()).device
 
 
-__all__ = ["ComposerReplicationTrainer"]
+def make_dr_grpo_config(**overrides: Any):
+    """Build a `trl.GRPOConfig` configured to the **Dr. GRPO** recipe.
+
+    Per the Composer 2 technical report (arXiv:2603.24477,
+    research/10-composer2-techreport-mining.md) the RL base is Dr. GRPO
+    (Liu et al., arXiv:2503.20783):
+
+      - ``loss_type="dr_grpo"``  — removes GRPO's length-standardization term
+        (which injects a length bias). TRL's own help text cites the Dr. GRPO
+        paper for this.
+      - ``scale_rewards="none"`` — NO std-dev advantage normalization. TRL docs:
+        "The Dr. GRPO paper recommends not scaling rewards, as scaling by the
+        standard deviation introduces a question-level difficulty bias."
+      - ``num_iterations=1``     — single-epoch regime (a prompt is never
+        trained on twice), matching the tech report.
+      - ``beta`` (KL-to-ref coef) kept; TRL uses the k1 (−log r)-family
+        estimator the report selects.
+
+    Any field can be overridden via kwargs (e.g. ``learning_rate=...``,
+    ``output_dir=...``). The three Dr. GRPO-defining knobs are forced unless
+    explicitly overridden, and a sanity assertion guards against silent drift.
+    """
+    from trl import GRPOConfig  # local import: only when actually building a config
+
+    dr_grpo_defaults: dict[str, Any] = {
+        "loss_type": "dr_grpo",
+        "scale_rewards": "none",
+        "num_iterations": 1,
+    }
+    merged = {**dr_grpo_defaults, **overrides}
+    cfg = GRPOConfig(**merged)
+    # Guard: fail loudly if a future TRL renames/repurposes these knobs.
+    assert cfg.loss_type == merged["loss_type"], "GRPOConfig dropped loss_type"
+    assert str(cfg.scale_rewards) in ("none", "False", "False"), (
+        f"Dr. GRPO requires scale_rewards='none' (no std-norm); got {cfg.scale_rewards!r}. "
+        "TRL knob may have drifted — re-verify against trl version."
+    )
+    assert cfg.num_iterations == merged["num_iterations"], "GRPOConfig dropped num_iterations"
+    return cfg
+
+
+__all__ = ["ComposerReplicationTrainer", "make_dr_grpo_config"]

composer_replication/trainer/tests/test_dr_grpo_config_and_alignment.py

@@ -0,0 +1,147 @@
+"""Tests for the Dr. GRPO config helper + SDPO strict-alignment guard (ADR-008).
+
+Covers ADR-008 acceptance gates:
+  - gate 1: make_dr_grpo_config sets loss_type=dr_grpo, scale_rewards=none,
+    num_iterations=1 (asserted against the resulting GRPOConfig).
+  - gate 2: strict_sdpo_alignment=True raises on a student/teacher shape
+    mismatch rather than silently zeroing the channel.
+  - gate 5: the PRIME-RL recipe raises NotImplementedError for alpha_sdpo>0
+    with a message pointing at the TRL host (documents the ADR-006 amendment).
+
+These are CPU-only and fast. The full end-to-end smoke (gate 3) that
+instantiates a real GRPOTrainer on Qwen2.5-0.5B lives in
+examples/composer_grpo_sdpo_smoke/ (gated on model cache + memory).
+"""
+from __future__ import annotations
+
+import pytest
+import torch
+
+trl = pytest.importorskip("trl")  # whole module is meaningless without TRL
+
+
+# ---------------------------------------------------------------------------
+# Gate 1 — Dr. GRPO config
+# ---------------------------------------------------------------------------
+
+def test_make_dr_grpo_config_sets_dr_grpo_knobs(tmp_path):
+    from composer_replication.trainer.composer_trainer import make_dr_grpo_config
+
+    cfg = make_dr_grpo_config(output_dir=str(tmp_path))
+    # loss_type=dr_grpo removes GRPO's length-standardization bias.
+    assert cfg.loss_type == "dr_grpo"
+    # scale_rewards=none => NO std-dev advantage normalization (Dr. GRPO).
+    assert str(cfg.scale_rewards) in ("none", "False")
+    # single-epoch: a prompt is never trained on twice.
+    assert cfg.num_iterations == 1
+
+
+def test_make_dr_grpo_config_allows_overrides(tmp_path):
+    from composer_replication.trainer.composer_trainer import make_dr_grpo_config
+
+    cfg = make_dr_grpo_config(output_dir=str(tmp_path), beta=0.05, num_generations=2)
+    assert cfg.beta == 0.05
+    assert cfg.num_generations == 2
+    # Dr. GRPO knobs still forced when not overridden.
+    assert cfg.loss_type == "dr_grpo"
+    assert str(cfg.scale_rewards) in ("none", "False")
+
+
+def test_make_dr_grpo_config_override_does_not_silently_break_guard(tmp_path):
+    """Overriding loss_type away from dr_grpo is allowed (caller's choice) and
+    the returned config reflects it — the guard only protects the *defaults*
+    from silent TRL drift, it does not forbid explicit opt-out."""
+    from composer_replication.trainer.composer_trainer import make_dr_grpo_config
+
+    cfg = make_dr_grpo_config(output_dir=str(tmp_path), loss_type="grpo")
+    assert cfg.loss_type == "grpo"
+
+
+# ---------------------------------------------------------------------------
+# Gate 2 — SDPO strict-alignment guard (no real GRPOTrainer needed)
+# ---------------------------------------------------------------------------
+
+class _TinyLM(torch.nn.Module):
+    """Minimal HF-style model: model(input_ids=...).logits, controllable vocab."""
+
+    def __init__(self, vocab: int = 16, hidden: int = 8):
+        super().__init__()
+        self.embed = torch.nn.Embedding(vocab, hidden)
+        self.head = torch.nn.Linear(hidden, vocab)
+
+    def forward(self, input_ids: torch.Tensor):
+        logits = self.head(self.embed(input_ids))
+
+        class _Out:
+            pass
+        out = _Out()
+        out.logits = logits
+        return out
+
+
+def _make_trainer_without_init(strict: bool):
+    """Build a ComposerReplicationTrainer instance WITHOUT running GRPOTrainer's
+    heavy __init__ (which needs a real model+dataset). We only exercise
+    _compute_sdpo_loss, so we set the attributes it reads directly."""
+    from composer_replication.trainer.composer_trainer import ComposerReplicationTrainer
+
+    obj = ComposerReplicationTrainer.__new__(ComposerReplicationTrainer)
+    obj.alpha_sdpo = 1.0
+    obj.sdpo_jsd_beta = 0.5
+    obj.sdpo_temperature = 1.0
+    obj.sdpo_token_clip = None
+    obj.strict_sdpo_alignment = strict
+    return obj
+
+
+def test_strict_alignment_raises_on_shape_mismatch():
+    obj = _make_trainer_without_init(strict=True)
+    model = _TinyLM(vocab=16)
+    # student seq len 5, teacher seq len 7 -> logit shape mismatch
+    inputs = {
+        "input_ids": torch.randint(0, 16, (1, 5)),
+        "ctx_teacher_input_ids": torch.randint(0, 16, (1, 7)),
+        "sdpo_loss_mask": torch.ones(1, 5),
+    }
+    with pytest.raises(ValueError, match="shape mismatch"):
+        obj._compute_sdpo_loss(model, inputs)
+
+
+def test_nonstrict_alignment_warns_and_skips_on_mismatch():
+    obj = _make_trainer_without_init(strict=False)
+    model = _TinyLM(vocab=16)
+    inputs = {
+        "input_ids": torch.randint(0, 16, (1, 5)),
+        "ctx_teacher_input_ids": torch.randint(0, 16, (1, 7)),
+        "sdpo_loss_mask": torch.ones(1, 5),
+    }
+    loss = obj._compute_sdpo_loss(model, inputs)
+    assert float(loss.detach()) == 0.0  # skipped, not crashed
+
+
+def test_aligned_shapes_produce_finite_sdpo_loss():
+    obj = _make_trainer_without_init(strict=True)
+    model = _TinyLM(vocab=16)
+    # student and teacher SAME length -> channel fires
+    inputs = {
+        "input_ids": torch.randint(0, 16, (1, 6)),
+        "ctx_teacher_input_ids": torch.randint(0, 16, (1, 6)),
+        "sdpo_loss_mask": torch.ones(1, 6),
+    }
+    loss = obj._compute_sdpo_loss(model, inputs)
+    val = float(loss.detach())
+    assert val == val  # not NaN
+    assert val not in (float("inf"), float("-inf"))
+
+
+def test_no_error_sites_returns_zero_without_forward():
+    """Empty ctx_teacher_input_ids => no SDPO signal, returns 0 (no crash)."""
+    obj = _make_trainer_without_init(strict=True)
+    model = _TinyLM(vocab=16)
+    inputs = {
+        "input_ids": torch.randint(0, 16, (1, 6)),
+        "ctx_teacher_input_ids": torch.empty(0, dtype=torch.long),
+        "sdpo_loss_mask": torch.zeros(1, 6),
+    }
+    loss = obj._compute_sdpo_loss(model, inputs)
+    assert float(loss.detach()) == 0.0

examples/composer_grpo_sdpo_smoke/run.py

@@ -0,0 +1,133 @@
+"""ComposerGRPOTrainer ⊕ SDPO live smoke (ADR-008 gate 3).
+
+Instantiates a REAL `trl.GRPOTrainer` via `ComposerReplicationTrainer`, configured
+to the Dr. GRPO recipe (`make_dr_grpo_config`), on a tiny model, and runs a
+short training run with `alpha_sdpo>0` so the SDPO channel is live on top of the
+Dr. GRPO policy-gradient loss.
+
+This is the wrapper-level proof. The loss-composition CORE (compose_loss forward
++ backward + optimizer.step with the SDPO JSD firing on real traces) is already
+proven CPU-only by `examples/sdpo_real_trace_train_smoke/run.py`. This script
+proves the same SDPO channel survives inside a live TRL GRPO rollout→update loop.
+
+Heavy + slow on CPU (TRL import alone is ~140s; GRPO generation on CPU is slow).
+RUN DETACHED so a gateway restart can't reap it:
+
+    systemd-run --user --scope -p MemoryMax=28G -- \
+      bash -lc 'cd <repo> && source .venv/bin/activate && \
+        python examples/composer_grpo_sdpo_smoke/run.py > /tmp/grpo_smoke.log 2>&1; \
+        echo EXIT=$? >> /tmp/grpo_smoke.log; touch /tmp/grpo_smoke.done'
+
+Gates asserted:
+  - trainer instantiates with the Dr. GRPO config (loss_type=dr_grpo,
+    scale_rewards=none, num_iterations=1) and alpha_sdpo>0;
+  - a training step runs without crashing;
+  - total loss is finite;
+  - the SDPO channel is wired (loss/sdpo_kl logged) — value may be 0.0 if the
+    tiny synthetic rollouts happen to produce no error-aligned batch, which is
+    acceptable for the WRAPPER smoke (signal-firing is proven elsewhere).
+
+Exit 0 = PASS, 1 = FAIL, 2 = SKIP (model/TRL unavailable).
+"""
+from __future__ import annotations
+
+import os
+import sys
+
+
+def main() -> int:
+    os.environ.setdefault("HF_HUB_OFFLINE", "1")
+    os.environ.setdefault("TRANSFORMERS_OFFLINE", "1")
+    os.environ.setdefault("TRL_USE_VLLM", "0")
+    os.environ.setdefault("OMP_NUM_THREADS", "8")
+
+    model_id = os.environ.get("SMOKE_MODEL", "Qwen/Qwen2.5-0.5B-Instruct")
+
+    try:
+        import torch  # noqa: F401
+        from datasets import Dataset
+        from transformers import AutoModelForCausalLM, AutoTokenizer
+
+        from composer_replication.trainer.composer_trainer import (
+            ComposerReplicationTrainer,
+            make_dr_grpo_config,
+        )
+    except Exception as e:  # noqa: BLE001
+        print(f"SKIP: import failed: {e!r}")
+        return 2
+
+    print(f"[grpo-smoke] loading {model_id} (CPU) — slow ...")
+    try:
+        tok = AutoTokenizer.from_pretrained(model_id)
+        if tok.pad_token is None:
+            tok.pad_token = tok.eos_token
+        model = AutoModelForCausalLM.from_pretrained(model_id)
+    except Exception as e:  # noqa: BLE001
+        print(f"SKIP: model/tokenizer load failed: {e!r}")
+        return 2
+
+    # Trivial verifiable reward: reward length-1 presence of a digit (toy).
+    def reward_has_digit(completions, **kwargs):
+        return [1.0 if any(c.isdigit() for c in (t or "")) else 0.0 for t in completions]
+
+    # Tiny prompt dataset.
+    prompts = [{"prompt": "Reply with a number:"}, {"prompt": "Count to three:"}]
+    ds = Dataset.from_list(prompts)
+
+    cfg = make_dr_grpo_config(
+        output_dir="/tmp/grpo_smoke_out",
+        per_device_train_batch_size=2,
+        num_generations=2,
+        max_completion_length=8,
+        max_prompt_length=32,
+        max_steps=1,
+        logging_steps=1,
+        report_to=[],
+        beta=0.0,            # drop KL-to-ref for the smoke (no ref model load)
+        use_vllm=False,
+    )
+    print(f"[grpo-smoke] Dr.GRPO config: loss_type={cfg.loss_type} "
+          f"scale_rewards={cfg.scale_rewards} num_iterations={cfg.num_iterations}")
+
+    try:
+        trainer = ComposerReplicationTrainer(
+            model=model,
+            reward_funcs=reward_has_digit,
+            args=cfg,
+            train_dataset=ds,
+            processing_class=tok,
+            # SDPO channel ON. The toy rollouts won't carry collator-built
+            # ctx_teacher_input_ids, so _compute_sdpo_loss returns 0 (no error
+            # sites) — but the channel is WIRED and logged. strict=False so the
+            # absence of error sites is a clean no-op, not an abort.
+            alpha_sdpo=1.0,
+            strict_sdpo_alignment=False,
+        )
+    except Exception as e:  # noqa: BLE001
+        print(f"FAIL: trainer instantiation failed: {e!r}")
+        import traceback
+        traceback.print_exc()
+        return 1
+
+    print("[grpo-smoke] trainer instantiated; running 1 Dr. GRPO step "
+          "with alpha_sdpo=1.0 ...")
+    try:
+        trainer.train()
+    except Exception as e:  # noqa: BLE001
+        print(f"FAIL: train() crashed: {e!r}")
+        import traceback
+        traceback.print_exc()
+        return 1
+
+    # If we got here, the live loop ran with the SDPO channel wired in.
+    log_history = getattr(trainer.state, "log_history", [])
+    sdpo_logged = any("loss/sdpo_kl" in row for row in log_history)
+    print("=" * 60)
+    print(f"  trainer ran 1 Dr. GRPO step:        OK")
+    print(f"  loss/sdpo_kl present in log_history: {sdpo_logged}")
+    print(f"  RESULT: PASS ✅ (SDPO channel wired into live Dr. GRPO loop)")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())