examples/altered_minds_channel_ladder/run.py

"""B4 end-to-end CPU proof: the SDPO channel actually FIRES (NONZERO) on a real
collator-built batch with genuine alignment indices (ADR-013).

The existing examples/composer_grpo_sdpo_smoke proves the SDPO channel is *wired*
into a live TRL Dr.GRPO loop, but its toy synthetic rollouts carry no error
sites, so _compute_sdpo_loss returns 0 (the channel never actually fires). This
script closes that gap: it builds a REAL ComposerDataCollator batch from a trace
that HAS an error turn — so ctx_teacher_input_ids + student/teacher_response_idx
are emitted by the shipped collator — and proves the SDPO JSD is NONZERO over
>=1 step, in the A2 ladder config (alpha_sdpo=0.02).

PROOF ACHIEVED: stub-with-differing-tokens (NOT a real Qwen checkpoint).

  - Alignment indices: REAL (production ComposerDataCollator, real error turn).
  - Model: a deterministic position-dependent TinyLM stub (CPU, no download),
    the same pattern used by trainer/tests/test_sdpo_alignment_indices.py.
  - Why perturb student tokens: the collator's placeholder-alignment trick makes
    student & teacher carry identical tokens at identical positions at the valid
    aligned indices, so a deterministic stub yields JSD≈0 there (correct for a
    perfectly-aligned identical model). To prove the channel GATHERS the aligned
    positions and computes a real divergence, the student's input_ids are made to
    DIFFER from the teacher's at exactly those aligned positions — mimicking the
    hint actually changing the recovery tokens (the real-world case where SDPO
    has signal to distill). Different aligned tokens => different logits =>
    provably NONZERO JSD, on a differentiable grad path.

To run the SAME assertion against a real Qwen2.5-0.5B-Instruct (if cached
offline), set ALTERED_MINDS_REAL_MODEL=1 — note that even with a real model the
NONZERO signal still requires the aligned student/teacher tokens to differ, so
this script keeps the same token-perturbation; the real-model path only swaps
the stub for the HF model and is much slower on CPU.

Exit 0 = PASS (SDPO fired nonzero), 1 = FAIL, 2 = SKIP (deps unavailable).
"""
from __future__ import annotations

import os
import sys


def _build_tiny_lm(vocab: int):
    import torch

    class _TinyLM(torch.nn.Module):
        def __init__(self, vocab: int = 64, hidden: int = 8, max_pos: int = 512):
            super().__init__()
            torch.manual_seed(0)
            self.embed = torch.nn.Embedding(vocab, hidden)
            self.pos = torch.nn.Embedding(max_pos, hidden)
            self.head = torch.nn.Linear(hidden, vocab)

        def forward(self, input_ids):
            T = input_ids.size(1)
            positions = torch.arange(T, device=input_ids.device).unsqueeze(0)
            h = self.embed(input_ids) + self.pos(positions)

            class _Out:
                pass

            out = _Out()
            out.logits = self.head(h)
            return out

    return _TinyLM(vocab=max(vocab, 8))


class _StubTok:
    pad_token_id = 0

    def __init__(self) -> None:
        self._v = {"<pad>": 0, "<bos>": 1, "<eos>": 2}

    def _id(self, w: str) -> int:
        if w not in self._v:
            self._v[w] = len(self._v)
        return self._v[w]

    def __call__(self, text, **_k):
        return {"input_ids": [self._id(w) for w in text.split()] if text else []}

    def apply_chat_template(self, messages, tokenize=True, **_k):  # noqa: ARG002
        return [
            self._id(w)
            for w in " ".join(m.get("content", "") for m in messages).split()
        ]


def _hint_gen(_kind, _meta):
    return "HINT search before reading"


def _error_trace():
    return {
        "trace_id": "b4-channel-ladder",
        "turns": [
            {"role": "user", "content": "do the task now"},
            {"role": "user", "content": "tool not found error occurred"},
            {
                "role": "assistant",
                "content": "let me use a real working tool instead now",
                "tool_error": "tool_not_found",
                "error_meta": {},
            },
        ],
        "final_reward": 0.0,
    }


def main() -> int:
    os.environ.setdefault("HF_HUB_OFFLINE", "1")
    os.environ.setdefault("TRANSFORMERS_OFFLINE", "1")

    try:
        import torch  # noqa: F401

        from composer_replication.integrations.altered_minds import (
            channel_ladder_configs,
        )
        from composer_replication.trainer.composer_trainer import (
            ComposerReplicationTrainer,
            make_dr_grpo_config,
        )
        from composer_replication.trainer.data_collator import (
            CollatorConfig,
            ComposerDataCollator,
        )
    except Exception as e:  # noqa: BLE001
        print(f"SKIP: import failed: {e!r}")
        return 2

    # A2 arm = +SDPO small (alpha_sdpo=0.02), the amplification probe.
    a2 = next(a for a in channel_ladder_configs() if a["arm"] == "A2")
    print(f"[b4] ladder arm A2: alpha_sdpo={a2['alpha_sdpo']} "
          f"beta_replay={a2['beta_replay']} kl_beta={a2['kl_beta']}")

    # make_dr_grpo_config is exercised to prove the config wiring is intact
    # (the actual TLM stub forward does not need a GRPOConfig, but a real A2
    # runner would pass this through to ComposerReplicationTrainer).
    try:
        cfg = make_dr_grpo_config(output_dir="/tmp/b4_ladder_out", report_to=[])
        print(f"[b4] Dr.GRPO config OK: loss_type={cfg.loss_type} "
              f"scale_rewards={cfg.scale_rewards} num_iterations={cfg.num_iterations}")
    except Exception as e:  # noqa: BLE001
        print(f"[b4] (config build skipped: {e!r})")

    # --- REAL collator-built batch with a genuine error turn ---
    tok = _StubTok()
    collator = ComposerDataCollator(
        tokenizer=tok,
        config=CollatorConfig(hint_generator=_hint_gen, enable_replay_dpo=False),
    )
    batch = collator([_error_trace()])

    if batch.get("ctx_teacher_input_ids") is None or batch["ctx_teacher_input_ids"].numel() == 0:
        print("FAIL: collator emitted no error-site teacher context.")
        return 1
    s_idx = batch["student_response_idx"]
    s_valid = batch["student_response_valid"]
    if int(s_valid.sum()) == 0:
        print("FAIL: no valid aligned response positions.")
        return 1
    print(f"[b4] collator emitted real alignment indices: "
          f"student_response_idx shape={tuple(s_idx.shape)}, "
          f"valid positions={int(s_valid.sum())}")

    # --- Make the student tokens differ from teacher at aligned positions ---
    student_ids = batch["input_ids"].clone()
    vocab_ceiling = int(
        max(batch["input_ids"].max(), batch["ctx_teacher_input_ids"].max())
    ) + 8
    for b in range(s_idx.shape[0]):
        for k in range(s_idx.shape[1]):
            if bool(s_valid[b, k]):
                pos = int(s_idx[b, k])
                student_ids[b, pos] = (int(student_ids[b, pos]) + 3) % vocab_ceiling
    batch["input_ids"] = student_ids

    real_model = os.environ.get("ALTERED_MINDS_REAL_MODEL") == "1"
    if real_model:
        try:
            from transformers import AutoModelForCausalLM
            model_id = os.environ.get("SMOKE_MODEL", "Qwen/Qwen2.5-0.5B-Instruct")
            print(f"[b4] loading real model {model_id} (CPU, slow) ...")
            model = AutoModelForCausalLM.from_pretrained(model_id)
            print("[b4] real model loaded; proof path = REAL-MODEL")
        except Exception as e:  # noqa: BLE001
            print(f"[b4] real model unavailable ({e!r}); falling back to TinyLM stub")
            model = _build_tiny_lm(vocab_ceiling)
            real_model = False
    else:
        model = _build_tiny_lm(vocab_ceiling)

    # --- A2 config: SDPO-only small (alpha_sdpo=0.02), strict alignment ---
    obj = ComposerReplicationTrainer.__new__(ComposerReplicationTrainer)
    obj.alpha_sdpo = float(a2["alpha_sdpo"])
    obj.sdpo_jsd_beta = 0.5
    obj.sdpo_temperature = 1.0
    obj.sdpo_token_clip = None
    obj.strict_sdpo_alignment = True

    loss = obj._compute_sdpo_loss(model, batch)
    val = float(loss.detach())
    print("=" * 64)
    print(f"  proof path:            {'REAL-MODEL' if real_model else 'TinyLM-stub-with-differing-tokens'}")
    print(f"  SDPO JSD (sdpo_kl):    {val:.6f}")
    print(f"  requires_grad:         {loss.requires_grad}")

    if not (val == val) or val in (float("inf"), float("-inf")):
        print("  RESULT: FAIL ❌ (loss not finite)")
        return 1
    if val <= 1e-6:
        print("  RESULT: FAIL ❌ (SDPO channel did not fire — JSD ~0)")
        return 1

    (obj.alpha_sdpo * loss).backward()
    grad_norm = sum(
        float(p.grad.norm()) for p in model.parameters() if p.grad is not None
    )
    print(f"  grad norm into model:  {grad_norm:.6f}")
    if grad_norm <= 0.0:
        print("  RESULT: FAIL ❌ (no gradient flowed from SDPO loss)")
        return 1

    print("  RESULT: PASS ✅ (SDPO channel FIRED nonzero via real collator indices)")
    return 0


if __name__ == "__main__":
    sys.exit(main())