"""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 = {"": 0, "": 1, "": 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())