scripts/train_sft.py

"""scripts/train_sft.py - SFT warm-up phase (master spec, sections 1-3).

Loads ``Qwen/Qwen2.5-3B-Instruct`` in 4-bit (NF4) via Unsloth, attaches a
LoRA adapter (rank 16, alpha 32, dropout 0.05, on q/k/v/o projections),
and runs a single epoch of supervised fine-tuning on
``data/sft_dataset.jsonl`` (3,000 examples).

Goal: take the base model from ~0% format compliance to >=95% so the GRPO
trainer has a non-zero probability of getting parseable rewards.

Locked hyperparameters (master spec, section 1):
    * batch=4, grad_accum=4 -> effective batch 16
    * lr=2e-4 with 20-step linear warmup -> constant
    * weight_decay=0.01, optimizer=adamw_8bit, mixed precision=bf16
    * max_seq_len=1024, epochs=1, max_steps=200
    * checkpoint every 50, eval every 50, log every 10
    * seed=42

Designed to run on a Colab T4 in <=30 minutes.

Usage::

    pip install -r requirements-train.txt
    python -m scripts.train_sft \
        --dataset data/sft_dataset.jsonl \
        --val-dataset data/sft_validation.jsonl \
        --output checkpoints/sft_warmup \
        --report-to wandb

W&B logging (master spec, section 2)
------------------------------------
* Every 10 steps: TRL's built-in train/loss, learning_rate, grad_norm,
  epoch, global_step.
* Every 50 steps (validation pass on 100 held-out syndromes):

      eval/format_compliance
      eval/logical_correction_rate
      eval/exact_match_pymatching
      eval/hamming_overlap_mean
      eval/output_length_mean
      eval/output_diversity         (10 samples of one prompt @ T=0.7)
      eval/syndrome_consistency

* End-of-train: ``run.summary`` dump of final eval scores; LoRA adapter
  uploaded as a W&B artifact.

Early stopping (master spec, section 3)
---------------------------------------
Training halts as soon as ANY of these is true after a validation pass:

    1. format_compliance >= 0.95 AND logical_correction_rate >= 0.80
       AND output_diversity >= 3                                (success)
    2. global_step >= 200                                       (hard cap)
    3. wall-clock >= 30 minutes                                 (hard cap)
    4. train/loss has NaN or inf                                (failure)
"""
from __future__ import annotations

import argparse
import json
import os
import random
import re
import statistics
import sys
import time
from collections import defaultdict
from pathlib import Path
from typing import Iterable, Optional


# --------------------------------------------------------------------------- #
# Pre-flight dataset audit                                                    #
# --------------------------------------------------------------------------- #
# Runs as the FIRST step of main(), before any model/tokenizer/heavy imports.
# Catches dataset regressions (class collapse, format drift, parse breakage,
# size mismatches) in a few seconds, before burning ~30 min of GPU on a run
# that was doomed at row 0.
#
# 9 checks, 3 of them duplicated on the validation split. Any failure raises
# SystemExit(2) so the Colab/Lightning shell pipeline exits with a non-zero
# status and won't proceed to model loading.

_FORMAT_ANCHOR_RE = re.compile(r"X_ERRORS=\[[\d,\s]*\]\s*Z_ERRORS=\[[\d,\s]*\]\s*$")
_FORMAT_ONLY_RE = re.compile(r"^\s*X_ERRORS=\[[\d,\s]*\]\s*Z_ERRORS=\[[\d,\s]*\]\s*$")
_TAIL_RE = re.compile(r"X_ERRORS=\[([^\]]*)\]\s*Z_ERRORS=\[([^\]]*)\]\s*$")
_LEVEL_P_RE = re.compile(r"Physical error rate:\s*([\d.]+)")
_LEVEL_D_RE = re.compile(r"Code distance:\s*(\d+)")


def _detect_level_from_prompt(prompt: str) -> str:
    """Return ``"L1"``/``"L2"``/``"L3"``/``"unknown"`` for an SFT prompt.

    Used as a fallback for legacy datasets that didn't write a ``level``
    field into each record. We read the L1/L2/L3 ``p`` and ``distance``
    values straight from :mod:`qubit_medic.config` rather than hardcoding
    them, so the audit keeps working when the curriculum is tuned (e.g.
    L1's ``p`` was bumped from 0.0001 -> 0.0005, which broke the old
    hardcoded check and made every L1 row read as ``unknown``).
    """
    m_p = _LEVEL_P_RE.search(prompt)
    m_d = _LEVEL_D_RE.search(prompt)
    if not m_p or not m_d:
        return "unknown"
    p = float(m_p.group(1))
    d = int(m_d.group(1))
    try:
        from qubit_medic.config import level_by_name
        l3 = level_by_name("L3_stretch")
        l2 = level_by_name("L2_target")
        l1 = level_by_name("L1_warmup")
        if d == l3.distance and abs(p - l3.p) < 1e-9:
            return "L3"
        if d == l2.distance and abs(p - l2.p) < 1e-9:
            return "L2"
        if d == l1.distance and abs(p - l1.p) < 1e-9:
            return "L1"
    except Exception:
        pass
    return "unknown"


def _level_label_from_record(rec: dict) -> str:
    """Return ``"L1"``/``"L2"``/``"L3"``/``"unknown"`` for an SFT record.

    Prefers the explicit ``level`` field written by
    ``scripts/generate_sft_data.py`` (e.g. ``"L1_warmup"``). Falls back
    to :func:`_detect_level_from_prompt` for legacy records that lack
    that field.
    """
    raw = rec.get("level")
    if isinstance(raw, str):
        if raw.startswith("L1"):
            return "L1"
        if raw.startswith("L2"):
            return "L2"
        if raw.startswith("L3"):
            return "L3"
    prompt = rec.get("prompt")
    if isinstance(prompt, str):
        return _detect_level_from_prompt(prompt)
    return "unknown"


def _has_nonempty_correction(completion: str) -> bool:
    """True iff the completion's trailing format line predicts at least one
    error (X or Z). Robust to a leading reasoning prefix.
    """
    m = _TAIL_RE.search(completion.rstrip())
    if m is None:
        return False
    return bool(m.group(1).strip()) or bool(m.group(2).strip())


def _audit_file(path: Path) -> dict:
    """Compute raw audit metrics for one JSONL file."""
    if not path.exists():
        return {"error": f"missing file: {path}"}
    rows: list[dict] = []
    parse_failures = 0
    with path.open() as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            try:
                rec = json.loads(line)
            except json.JSONDecodeError:
                parse_failures += 1
                continue
            if "prompt" not in rec or "completion" not in rec:
                parse_failures += 1
                continue
            rows.append(rec)
    n = len(rows)
    total_lines = n + parse_failures
    parse_rate = (n / total_lines) if total_lines else 0.0
    nonempty = sum(_has_nonempty_correction(r["completion"]) for r in rows)
    anchor = sum(1 for r in rows if _FORMAT_ANCHOR_RE.search(r["completion"].rstrip()))
    levels = {"L1": 0, "L2": 0, "L3": 0, "unknown": 0}
    for r in rows:
        levels[_level_label_from_record(r)] += 1
    plens = [len(r["prompt"]) for r in rows]
    clens = [len(r["completion"]) for r in rows]
    format_only = sum(1 for r in rows if _FORMAT_ONLY_RE.fullmatch(r["completion"].strip()))
    return {
        "n": n,
        "parse_failures": parse_failures,
        "parse_rate": parse_rate,
        "nonempty_frac": (nonempty / n) if n else 0.0,
        "anchor_frac": (anchor / n) if n else 0.0,
        "level_pct": {k: ((v / n) if n else 0.0) for k, v in levels.items()},
        "plens": plens,
        "clens": clens,
        "format_only_frac": (format_only / n) if n else 0.0,
    }


def audit_sft_dataset(
    train_path: str = "data/sft_dataset.jsonl",
    val_path: str = "data/sft_validation.jsonl",
) -> None:
    """Pre-flight audit of the SFT dataset. Halts (SystemExit) on violation.

    Runs 9 checks against ``train_path`` plus 4 parallel checks against
    ``val_path``. Designed to run in seconds on the CPU before any heavy
    ML deps are imported, so a broken dataset never reaches the GPU.

    Locked thresholds:
        Total rows:           train=3000, val=100
        JSON parse rate:      100%
        Non-empty correction: 65-75%
        Format anchor:        100%
        Curriculum L1/L2/L3:  35-45% / 45-55% / 7-15%
        Prompt length:        min>=800, median in [1100,1600], max<=2200
        Completion length:    min>=22, median in [22,80], max<=120
        Format-only target:   100%
        Validation parallel:  same thresholds applied to val split
    """
    EXPECTED_TRAIN = 3000
    EXPECTED_VAL = 100
    NONEMPTY_LO, NONEMPTY_HI = 0.65, 0.75
    # Tightened to match quota-based per-level generation in
    # scripts/generate_sft_data.py, which produces the 40/50/10 split
    # exactly (no rejection-sampling drift).
    L1_LO, L1_HI = 0.38, 0.42
    L2_LO, L2_HI = 0.48, 0.52
    L3_LO, L3_HI = 0.08, 0.12
    PLEN_MIN, PLEN_MED_LO, PLEN_MED_HI, PLEN_MAX = 800, 1100, 1600, 2200
    # Targets are deliberately one-line format strings. The earlier
    # reasoning-prefix targets made the base model burn the full eval token
    # budget on analysis and never reach the required parseable answer line.
    CLEN_MIN, CLEN_MED_LO, CLEN_MED_HI, CLEN_MAX = 22, 22, 80, 120
    FORMAT_ONLY_MIN = 1.0

    train = _audit_file(Path(train_path))
    if "error" in train:
        print(f"[audit] FATAL: {train['error']}")
        raise SystemExit(2)

    # ------------------------------- train checks ------------------------- #
    checks: list[tuple[str, str, bool]] = []

    checks.append((
        "Total rows",
        f"{train['n']} (expected {EXPECTED_TRAIN})",
        train["n"] == EXPECTED_TRAIN,
    ))
    checks.append((
        "JSON parse rate",
        f"{train['parse_rate'] * 100:.1f}% ({train['parse_failures']} failures)",
        abs(train["parse_rate"] - 1.0) < 1e-9,
    ))
    checks.append((
        "Non-empty correction",
        f"{train['nonempty_frac'] * 100:.1f}% (target 65-75%)",
        NONEMPTY_LO <= train["nonempty_frac"] <= NONEMPTY_HI,
    ))
    checks.append((
        "Format anchor",
        f"{train['anchor_frac'] * 100:.1f}%",
        abs(train["anchor_frac"] - 1.0) < 1e-9,
    ))

    p1 = train["level_pct"]["L1"]
    p2 = train["level_pct"]["L2"]
    p3 = train["level_pct"]["L3"]
    p_unknown = train["level_pct"]["unknown"]
    checks.append((
        "Curriculum L1/L2/L3",
        f"{p1*100:.1f}/{p2*100:.1f}/{p3*100:.1f}% (unknown={p_unknown*100:.1f}%)",
        (L1_LO <= p1 <= L1_HI
         and L2_LO <= p2 <= L2_HI
         and L3_LO <= p3 <= L3_HI),
    ))

    pmin = min(train["plens"]) if train["plens"] else 0
    pmed = int(statistics.median(train["plens"])) if train["plens"] else 0
    pmax = max(train["plens"]) if train["plens"] else 0
    checks.append((
        "Prompt length",
        f"min={pmin} median={pmed} max={pmax}",
        (pmin >= PLEN_MIN
         and PLEN_MED_LO <= pmed <= PLEN_MED_HI
         and pmax <= PLEN_MAX),
    ))

    cmin = min(train["clens"]) if train["clens"] else 0
    cmed = int(statistics.median(train["clens"])) if train["clens"] else 0
    cmax = max(train["clens"]) if train["clens"] else 0
    checks.append((
        "Completion length",
        f"min={cmin} median={cmed} max={cmax}",
        (cmin >= CLEN_MIN
         and CLEN_MED_LO <= cmed <= CLEN_MED_HI
         and cmax <= CLEN_MAX),
    ))

    checks.append((
        "Format-only completions",
        f"{train['format_only_frac'] * 100:.1f}% (target 100%)",
        abs(train["format_only_frac"] - FORMAT_ONLY_MIN) < 1e-9,
    ))

    # ------------------------------- val parallel ------------------------- #
    val = _audit_file(Path(val_path))
    if "error" in val:
        checks.append(("Validation parallel", val["error"], False))
    else:
        v1 = val["level_pct"]["L1"]
        v2 = val["level_pct"]["L2"]
        v3 = val["level_pct"]["L3"]
        val_pass = (
            val["n"] == EXPECTED_VAL
            and abs(val["parse_rate"] - 1.0) < 1e-9
            and NONEMPTY_LO <= val["nonempty_frac"] <= NONEMPTY_HI
            and abs(val["anchor_frac"] - 1.0) < 1e-9
            and abs(val["format_only_frac"] - FORMAT_ONLY_MIN) < 1e-9
            and L1_LO <= v1 <= L1_HI
            and L2_LO <= v2 <= L2_HI
            and L3_LO <= v3 <= L3_HI
        )
        val_summary = (
            f"rows={val['n']} parse={val['parse_rate']*100:.0f}% "
            f"nonempty={val['nonempty_frac']*100:.1f}% "
            f"anchor={val['anchor_frac']*100:.0f}% "
            f"format_only={val['format_only_frac']*100:.0f}% "
            f"L1/L2/L3={v1*100:.1f}/{v2*100:.1f}/{v3*100:.1f}%"
        )
        checks.append(("Validation parallel", val_summary, val_pass))

    # ------------------------------- print banner ------------------------- #
    print()
    print("DATASET AUDIT SUMMARY")
    print("=" * 21)
    label_w = max(len(label) for label, _, _ in checks) + 1
    val_w = max(len(val_str) for _, val_str, _ in checks)
    for label, val_str, passed in checks:
        mark = "✓" if passed else "✗"  # ✓ / ✗
        print(f"{(label + ':').ljust(label_w + 1)} {val_str.ljust(val_w)} [{mark}]")

    all_passed = all(passed for _, _, passed in checks)
    print()
    if all_passed:
        print("ALL CHECKS PASSED — DATASET READY FOR TRAINING")
        print()
        return
    print("AUDIT FAILED — FIX DATASET BEFORE TRAINING")
    print()
    raise SystemExit(2)


# --------------------------------------------------------------------------- #
# Validation-record loading                                                   #
# --------------------------------------------------------------------------- #


def _load_jsonl(path: str) -> list[dict]:
    rows: list[dict] = []
    with open(path) as f:
        for line in f:
            rows.append(json.loads(line))
    return rows


def _load_train_dataset(path: str, tokenizer):
    """Load the SFT JSONL into a HuggingFace Dataset.

    Master spec (section 4): the chat template is applied via the
    tokenizer (``apply_chat_template``), NOT by manually inserting
    ``<|im_start|>`` markers - that way the same template works across
    Qwen2.5 / Qwen3 / etc. without surprises.
    """
    from datasets import Dataset

    rows = _load_jsonl(path)
    out = []
    for rec in rows:
        messages = [
            {"role": "user", "content": rec["prompt"]},
            {"role": "assistant", "content": rec["completion"]},
        ]
        try:
            text = tokenizer.apply_chat_template(messages, tokenize=False)
        except Exception:
            # Defensive fallback if apply_chat_template ever misbehaves.
            text = (
                "<|im_start|>user\n"
                f"{rec['prompt']}\n<|im_end|>\n"
                "<|im_start|>assistant\n"
                f"{rec['completion']}<|im_end|>"
            )
        out.append({
            "prompt": rec["prompt"],
            "completion": rec["completion"],
            "text": text,
        })
    return Dataset.from_list(out)


# --------------------------------------------------------------------------- #
# Per-level physics caches (used by the validation callback)                  #
# --------------------------------------------------------------------------- #


def _build_level_caches(needed_levels: set[str]) -> dict[str, dict]:
    """Pre-build circuit / matching / layout / supports per curriculum level."""
    import pymatching

    from qubit_medic.config import level_by_name
    from qubit_medic.server.physics import (
        build_circuit, build_dem, extract_layout, per_round_x_z_counts,
    )
    from qubit_medic.server.rewards import compute_final_detector_supports

    caches: dict[str, dict] = {}
    for name in needed_levels:
        lvl = level_by_name(name)
        circuit = build_circuit(lvl)
        dem = build_dem(circuit)
        matching = pymatching.Matching.from_detector_error_model(dem)
        layout = extract_layout(circuit)
        n_x, n_z = per_round_x_z_counts(layout)
        supports = compute_final_detector_supports(layout)
        caches[name] = {
            "level": lvl,
            "circuit": circuit,
            "dem": dem,
            "matching": matching,
            "layout": layout,
            "supports": supports,
            "num_x_stab": n_x,
            "num_z_stab": n_z,
        }
    return caches


# --------------------------------------------------------------------------- #
# Validation callback (master spec, section 2 + section 3)                    #
# --------------------------------------------------------------------------- #


def _build_validation_callback(
    *,
    model,
    tokenizer,
    val_records: list[dict],
    eval_every: int,
    eval_schedule: tuple[tuple[int, int, str], ...] | None,
    print_sample_outputs: int,
    output_dir: str,
    max_new_tokens: int,
    diversity_n_samples: int,
    diversity_temperature: float,
    early_stop_format: float,
    early_stop_correction: float,
    early_stop_diversity: int,
    max_wall_seconds: float,
    started_wall: float,
    diversity_floor: int = 2,
    diversity_run_len: int = 2,
):
    """Returns a ``TrainerCallback`` that:
       * fires at every step in ``eval_schedule`` (or every ``eval_every``
         steps if no schedule is given) with a per-step sample size,
       * logs the spec metrics + new diagnostic metrics to W&B,
       * prints the first ``print_sample_outputs`` raw model outputs to
         stdout AND to ``{output_dir}/eval_samples_step{N}.txt`` so a
         broken parser / generation drift can be diagnosed in seconds,
       * stops training when the success criterion or hard caps fire.

    Metric semantics changed in this revision:
        * Parse failures NO LONGER default to "predict no errors". Failed
          rows contribute logical_correction=0, hamming=0,
          syndrome_consistency=0 to the aggregates. This stops trivial
          syndromes (~95% at p=0.001) from inflating logical_correction_rate
          to 0.98 while format_compliance sits at 0.01.
        * New ``eval/parse_failure_rate`` = 1 - format_compliance, so a
          broken parser is impossible to miss.
        * New ``eval/format_compliance_strict`` reports the share of
          outputs that hit the canonical ``X_ERRORS=[...] Z_ERRORS=[...]``
          form (Reward 4 == 1.0). The looser ``eval/format_compliance``
          reports the share where the model's answer was extractable at all.
    """
    from transformers import TrainerCallback

    from qubit_medic import wandb_utils
    from qubit_medic.prompts import parse_action
    from qubit_medic.server.physics import SyndromeSample
    from qubit_medic.server.rewards import compute_all_rewards

    if not val_records:
        return None

    # Pre-build per-level physics for fast scoring.
    needed = {r["level"] for r in val_records}
    level_caches = _build_level_caches(needed)

    # Pick one stable prompt for the diversity probe (always the same record
    # so the diversity number is comparable across checkpoints).
    diversity_record = val_records[0]
    diversity_messages = [{"role": "user", "content": diversity_record["prompt"]}]

    # Index the schedule: step -> (sample_size, mode). Sample sizes are
    # capped at len(val_records) so a small held-out set still works.
    if eval_schedule:
        schedule = {
            step: (min(size, len(val_records)), mode)
            for step, size, mode in eval_schedule
        }
    else:
        schedule = {}

    sample_dir = Path(output_dir)
    sample_dir.mkdir(parents=True, exist_ok=True)

    # 2026-04 (FIX 2) diversity-collapse rolling buffer. We track the
    # last ``diversity_run_len`` full-eval ``output_diversity`` values
    # and stop training when every entry is below ``diversity_floor``.
    from collections import deque as _deque
    recent_diversity = _deque(maxlen=diversity_run_len)

    class _ValidationCallback(TrainerCallback):
        # Stamp the most recent eval here so the on_train_end hook can avoid
        # re-running if the eval step coincided with the final step.
        last_eval_step: int = -1

        def on_step_end(self, args, state, control, **kwargs):  # noqa: D401
            now = time.time() - started_wall
            if now >= max_wall_seconds:
                print(f"[sft] wall-clock cap {max_wall_seconds:.0f}s hit at step "
                      f"{state.global_step}; stopping.")
                control.should_training_stop = True
                return

            step = state.global_step
            if step == 0:
                return
            if schedule:
                if step not in schedule:
                    return
            else:
                if step % eval_every != 0:
                    return
            self._run_eval(state, control)

        def on_train_end(self, args, state, control, **kwargs):  # noqa: D401
            if state.global_step != self.last_eval_step:
                self._run_eval(state, control, final=True)

        # ------------------------------------------------------------------ #
        # Core evaluation                                                    #
        # ------------------------------------------------------------------ #
        def _generate_greedy(self, messages: list[dict]) -> tuple[str, int]:
            text = tokenizer.apply_chat_template(
                messages, tokenize=False, add_generation_prompt=True,
            )
            inputs = tokenizer(text, return_tensors="pt").to(model.device)
            try:
                out = model.generate(
                    **inputs,
                    max_new_tokens=max_new_tokens,
                    do_sample=False,
                    eos_token_id=tokenizer.eos_token_id,
                    pad_token_id=tokenizer.pad_token_id or tokenizer.eos_token_id,
                )
                gen_ids = out[0][inputs["input_ids"].shape[1]:]
                completion = tokenizer.decode(gen_ids, skip_special_tokens=True)
                return completion, int(gen_ids.shape[0])
            except Exception as exc:
                return f"<gen-error: {exc}>", 0

        def _generate_sampled(self, messages: list[dict]) -> str:
            text = tokenizer.apply_chat_template(
                messages, tokenize=False, add_generation_prompt=True,
            )
            inputs = tokenizer(text, return_tensors="pt").to(model.device)
            try:
                out = model.generate(
                    **inputs,
                    max_new_tokens=max_new_tokens,
                    do_sample=True,
                    temperature=diversity_temperature,
                    top_p=0.95,
                    eos_token_id=tokenizer.eos_token_id,
                    pad_token_id=tokenizer.pad_token_id or tokenizer.eos_token_id,
                )
                return tokenizer.decode(
                    out[0][inputs["input_ids"].shape[1]:],
                    skip_special_tokens=True,
                )
            except Exception as exc:
                return f"<gen-error: {exc}>"

        def _run_eval(self, state, control, *, final: bool = False) -> None:
            self.last_eval_step = state.global_step
            try:
                from unsloth import FastLanguageModel
                FastLanguageModel.for_inference(model)
            except Exception:
                model.eval()  # type: ignore[attr-defined]

            step = state.global_step
            # Resolve sample size + mode for this step.
            if final and step in schedule:
                sample_size, mode = schedule[step]
            elif final:
                sample_size, mode = len(val_records), "full"
            elif step in schedule:
                sample_size, mode = schedule[step]
            else:
                sample_size, mode = len(val_records), "full"

            # Deterministic slice so the same prompts are used across checkpoints.
            records = val_records[:sample_size]
            n = len(records)
            full_eval = (mode == "full")

            n_format = 0          # lenient parse_success
            n_format_strict = 0   # canonical "=" + "[]"
            n_logical = n_exact = 0
            sum_hamming = 0.0
            sum_syndrome = 0.0
            sum_length = 0
            rows: list[dict] = []
            sample_dump_lines: list[str] = [
                f"=== eval samples @ step {step} (mode={mode}, n={n}) ===",
            ]

            for idx, rec in enumerate(records):
                num_data = int(rec["num_data_qubits"])
                messages = [{"role": "user", "content": rec["prompt"]}]
                completion, n_tokens = self._generate_greedy(messages)
                sum_length += n_tokens

                parsed = parse_action(completion, num_data_qubits=num_data)
                fmt_ok = parsed.parse_success
                fmt_strict_ok = bool(parsed.strict_format)
                n_format += int(fmt_ok)
                n_format_strict += int(fmt_strict_ok)

                # Physics-heavy metrics only in "full" mode AND only when
                # the parse actually succeeded. A failed parse means the
                # model didn't produce a usable prediction; we score that
                # as a miss (0) for every downstream metric instead of
                # silently substituting an empty Pauli frame, which would
                # accidentally score correct on the ~95% of trivial
                # syndromes at p=0.001.
                logical_ok = False
                exact_ok = False
                hamming = 0.0
                syndrome = 0.0
                if full_eval and fmt_ok:
                    cache = level_caches[rec["level"]]
                    layout = cache["layout"]
                    supports = cache["supports"]
                    sample = SyndromeSample(
                        syndrome_bits=list(map(int, rec["syndrome_bits"])),
                        actual_observable_flip=int(rec["actual_observable_flip"]),
                        pymatching_observable_pred=int(rec["pymatching_observable_pred"]),
                        pymatching_x_errors=list(map(int, rec["true_x_errors"])),
                        pymatching_z_errors=list(map(int, rec["true_z_errors"])),
                    )
                    breakdown = compute_all_rewards(parsed, sample, layout, supports)
                    logical_ok = breakdown.logical_correction >= 0.5
                    hamming = float(breakdown.hamming_overlap)
                    syndrome = float(breakdown.syndrome_consistency)
                    exact_ok = (
                        parsed.x_errors == sorted(set(rec["true_x_errors"]))
                        and parsed.z_errors == sorted(set(rec["true_z_errors"]))
                    )

                n_logical += int(logical_ok)
                n_exact += int(exact_ok)
                sum_hamming += hamming
                sum_syndrome += syndrome

                if idx < print_sample_outputs:
                    sample_dump_lines.append(
                        f"\n--- sample {idx} (level={rec['level']}, "
                        f"true_x={rec['true_x_errors']}, true_z={rec['true_z_errors']}, "
                        f"fmt_ok={fmt_ok}, fmt_strict={fmt_strict_ok}, "
                        f"n_tokens={n_tokens}) ---\n"
                        f">>> RAW MODEL OUTPUT:\n{completion}\n"
                        f">>> PARSED: x={parsed.x_errors} z={parsed.z_errors}"
                    )

                if idx < 4:  # keep W&B table tiny
                    rows.append({
                        "step": step,
                        "prompt": rec["prompt"][:600],
                        "gold": rec["completion"],
                        "model": completion[:300],
                        "x_pred": ",".join(map(str, parsed.x_errors)),
                        "z_pred": ",".join(map(str, parsed.z_errors)),
                        "format_ok": fmt_ok,
                        "format_strict_ok": fmt_strict_ok,
                        "logical_ok": logical_ok,
                        "exact_match": exact_ok,
                        "hamming_overlap": hamming,
                    })

            # ---------- print + persist raw output samples -------------- #
            sample_blob = "\n".join(sample_dump_lines)
            print(sample_blob)
            try:
                (sample_dir / f"eval_samples_step{step}.txt").write_text(sample_blob)
            except OSError as exc:
                print(f"[sft][eval@{step}] could not persist sample outputs: {exc}")

            # ---------- diversity probe (skip in format_only mode) ------ #
            if full_eval:
                diverse_outputs: list[str] = []
                for _ in range(diversity_n_samples):
                    diverse_outputs.append(self._generate_sampled(diversity_messages))
                output_diversity = len(set(diverse_outputs))
            else:
                output_diversity = 0  # not measured this step

            # ---------- aggregate + log to W&B ------------------------- #
            metrics: dict[str, float | int] = {
                "eval/format_compliance": n_format / max(1, n),
                "eval/format_compliance_strict": n_format_strict / max(1, n),
                "eval/parse_failure_rate": 1.0 - (n_format / max(1, n)),
                "eval/output_length_mean": sum_length / max(1, n),
                "eval/episodes": n,
                "eval/mode_full": int(full_eval),
            }
            if full_eval:
                metrics.update({
                    "eval/logical_correction_rate": n_logical / max(1, n),
                    "eval/exact_match_pymatching": n_exact / max(1, n),
                    "eval/hamming_overlap_mean": sum_hamming / max(1, n),
                    "eval/syndrome_consistency": sum_syndrome / max(1, n),
                    "eval/output_diversity": output_diversity,
                })
            print(f"[sft][eval@{step}] " + ", ".join(
                f"{k.split('/')[-1]}={v:.3f}" if isinstance(v, float) else f"{k.split('/')[-1]}={v}"
                for k, v in metrics.items()
            ))
            wandb_utils.log(metrics, step=step)
            wandb_utils.log_generation_table(
                rows, step=step,
                table_name=("sft/final_validation" if final else "sft/validation"),
                columns=["step", "prompt", "gold", "model", "x_pred", "z_pred",
                         "format_ok", "format_strict_ok", "logical_ok",
                         "exact_match", "hamming_overlap"],
            )

            # Fail fast on the known broken-SFT pattern: the model burns the
            # whole generation budget on prose and never emits the format line.
            # These thresholds mirror the runbook table in the issue analysis.
            format_floor_by_step = {5: 0.10, 15: 0.30, 30: 0.60, 50: 0.80}
            floor = format_floor_by_step.get(step)
            if (
                floor is not None
                and not final
                and metrics["eval/format_compliance"] < floor
            ):
                print(
                    f"[sft] format guard tripped at step {step}: "
                    f"format_compliance={metrics['eval/format_compliance']:.3f} "
                    f"< {floor:.2f}. Stop and inspect raw outputs / data."
                )
                control.should_training_stop = True
                wandb_utils.update_summary({
                    "sft/early_stop_reason": "format_guard",
                    "sft/format_guard_step": step,
                    "sft/format_guard_floor": floor,
                })

            # ---------- early stop checks ------------------------------ #
            # Only meaningful on full evals: logical_correction_rate and
            # output_diversity are not measured in format_only mode.
            if full_eval:
                success = (
                    metrics["eval/format_compliance"] >= early_stop_format
                    and metrics["eval/logical_correction_rate"] >= early_stop_correction
                    and metrics["eval/output_diversity"] >= early_stop_diversity
                )
                if success and not final:
                    print(f"[sft] success criterion hit at step {state.global_step}: "
                          f"format={metrics['eval/format_compliance']:.3f} >= {early_stop_format}, "
                          f"correction={metrics['eval/logical_correction_rate']:.3f} >= {early_stop_correction}, "
                          f"diversity={int(metrics['eval/output_diversity'])} >= {early_stop_diversity}; "
                          f"stopping.")
                    control.should_training_stop = True
                    wandb_utils.update_summary({"sft/early_stop_reason": "success_criterion"})

                # 2026-04 (FIX 2) diversity-collapse early stop. Pushed
                # AFTER the success check so a model that satisfies both
                # criteria still wins; only sustained low diversity
                # without convergence triggers the regression stop.
                recent_diversity.append(int(metrics["eval/output_diversity"]))
                if (
                    not final
                    and not control.should_training_stop
                    and len(recent_diversity) >= diversity_run_len
                    and all(d < diversity_floor for d in recent_diversity)
                ):
                    history = list(recent_diversity)
                    print(
                        f"[sft] diversity collapse early stop at step "
                        f"{state.global_step}: eval/output_diversity has "
                        f"been < {diversity_floor} for {diversity_run_len} "
                        f"consecutive full evals (history={history}). "
                        f"Stopping. Bump --lora-dropout (e.g. 0.15) or "
                        f"increase label smoothing and rerun."
                    )
                    control.should_training_stop = True
                    wandb_utils.update_summary({
                        "sft/early_stop_reason": "diversity_collapse",
                        "sft/diversity_collapse_step": state.global_step,
                        "sft/diversity_collapse_history": history,
                    })

            try:
                from unsloth import FastLanguageModel
                FastLanguageModel.for_training(model)
            except Exception:
                model.train()  # type: ignore[attr-defined]

    return _ValidationCallback()


# --------------------------------------------------------------------------- #
# Loss-divergence guard (failure mode early stop)                             #
# --------------------------------------------------------------------------- #


def _build_loss_guard_callback():
    import math

    from transformers import TrainerCallback

    class _LossGuard(TrainerCallback):
        def on_log(self, args, state, control, logs=None, **kwargs):  # noqa: D401
            if not logs:
                return
            loss = logs.get("loss")
            if loss is None:
                return
            try:
                lf = float(loss)
            except (TypeError, ValueError):
                return
            if math.isnan(lf) or math.isinf(lf):
                print(f"[sft] loss={loss} is NaN/inf at step {state.global_step}; "
                      f"stopping training.")
                control.should_training_stop = True

    return _LossGuard()


# --------------------------------------------------------------------------- #
# Main                                                                         #
# --------------------------------------------------------------------------- #


def main(argv: Iterable[str] = ()) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--dataset", type=str, default="data/sft_dataset.jsonl")
    parser.add_argument("--val-dataset", type=str,
                        default="data/sft_validation.jsonl",
                        help="held-out validation JSONL (rich records). "
                             "If missing, validation is skipped.")
    parser.add_argument("--output", type=str, default="checkpoints/sft_warmup")
    parser.add_argument("--model", type=str,
                        default=os.getenv("QUBIT_MEDIC_MODEL",
                                          "Qwen/Qwen2.5-3B-Instruct"))
    parser.add_argument("--epochs", type=int, default=None)
    parser.add_argument("--batch-size", type=int, default=None)
    parser.add_argument("--grad-accum", type=int, default=None)
    parser.add_argument("--lr", type=float, default=None)
    parser.add_argument("--max-seq-len", type=int, default=None)
    parser.add_argument("--max-steps", type=int, default=None,
                        help="hard cap on training steps (default 200)")
    parser.add_argument("--seed", type=int, default=None)
    parser.add_argument("--lora-r", type=int, default=None)
    parser.add_argument("--lora-alpha", type=int, default=None)
    parser.add_argument("--lora-dropout", type=float, default=None)
    parser.add_argument("--report-to", type=str, default="wandb")
    parser.add_argument("--wandb-run-name", type=str, default=None)
    parser.add_argument("--wandb-group", type=str, default=None)
    parser.add_argument("--wandb-tags", type=str, nargs="*", default=("sft",))
    parser.add_argument("--wandb-notes", type=str, default=None)
    parser.add_argument("--eval-every", type=int, default=None,
                        help="run validation pass every N steps (legacy "
                             "fallback when --no-eval-schedule is set)")
    parser.add_argument("--no-eval-schedule", action="store_true",
                        help="disable the variable-cadence schedule "
                             "(SFT_EVAL_SCHEDULE) and fall back to "
                             "uniform --eval-every spacing")
    parser.add_argument("--print-sample-outputs", type=int,
                        default=None,
                        help="N raw model outputs to print + persist per eval "
                             "(defaults to SFT_PRINT_SAMPLE_OUTPUTS from config)")
    parser.add_argument("--diversity-samples", type=int, default=10,
                        help="N samples for the output_diversity probe")
    parser.add_argument("--diversity-temperature", type=float, default=0.7)
    parser.add_argument("--no-artifact", action="store_true")
    args = parser.parse_args(list(argv))

    # Pre-flight dataset audit. Runs in seconds on the CPU before any heavy
    # ML deps are imported, so a broken dataset never reaches the GPU. Halts
    # via SystemExit(2) on any threshold violation.
    audit_sft_dataset(args.dataset, args.val_dataset)

    # Heavy imports are lazy so this module is importable without GPU deps.
    try:
        from unsloth import FastLanguageModel
    except ImportError:
        print("ERROR: unsloth not installed. Run `pip install -r requirements-train.txt`",
              file=sys.stderr)
        return 1
    import torch
    from transformers import TrainingArguments
    from trl import SFTTrainer

    from qubit_medic import wandb_utils
    from qubit_medic.config import (
        LORA_ALPHA, LORA_DROPOUT, LORA_R, LORA_TARGET_MODULES, MODEL_ID,
        PRIMARY_SEED, SFT_BATCH_SIZE, SFT_DIVERSITY_COLLAPSE_RUN_LEN,
        SFT_EARLY_STOP_CORRECTION, SFT_EARLY_STOP_DIVERSITY,
        SFT_EARLY_STOP_FORMAT, SFT_EPOCHS, SFT_EVAL_EVERY, SFT_EVAL_SCHEDULE,
        SFT_GRAD_ACCUM, SFT_LABEL_SMOOTHING, SFT_LOG_EVERY, SFT_LR,
        SFT_LR_SCHEDULER, SFT_MAX_NEW_TOKENS, SFT_MAX_SEQ_LEN, SFT_MAX_STEPS,
        SFT_MAX_WALL_SECONDS, SFT_OPTIMIZER, SFT_PREFLIGHT_DIVERSITY_FLOOR,
        SFT_PRINT_SAMPLE_OUTPUTS, SFT_SAVE_EVERY, SFT_WARMUP_STEPS,
        SFT_WEIGHT_DECAY,
    )

    epochs = args.epochs if args.epochs is not None else SFT_EPOCHS
    batch_size = args.batch_size if args.batch_size is not None else SFT_BATCH_SIZE
    grad_accum = args.grad_accum if args.grad_accum is not None else SFT_GRAD_ACCUM
    lr = args.lr if args.lr is not None else SFT_LR
    max_seq_len = args.max_seq_len if args.max_seq_len is not None else SFT_MAX_SEQ_LEN
    max_steps = args.max_steps if args.max_steps is not None else SFT_MAX_STEPS
    seed = args.seed if args.seed is not None else PRIMARY_SEED
    lora_r = args.lora_r if args.lora_r is not None else LORA_R
    lora_alpha = args.lora_alpha if args.lora_alpha is not None else LORA_ALPHA
    lora_dropout = args.lora_dropout if args.lora_dropout is not None else LORA_DROPOUT
    eval_every = args.eval_every if args.eval_every is not None else SFT_EVAL_EVERY
    print_sample_outputs = (
        args.print_sample_outputs
        if args.print_sample_outputs is not None
        else SFT_PRINT_SAMPLE_OUTPUTS
    )
    model_id = args.model if args.model else MODEL_ID

    random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)

    # ---- W&B init (no-op if unavailable / disabled) -------------------- #
    report_to = wandb_utils.derive_report_to(args.report_to)
    run_name = args.wandb_run_name or wandb_utils.make_run_name("sft")
    wandb_utils.init_run(
        run_name=run_name,
        job_type="sft",
        tags=args.wandb_tags,
        notes=args.wandb_notes,
        group=args.wandb_group,
        extra_config={
            "cli": {
                "epochs": epochs,
                "batch_size": batch_size,
                "grad_accum": grad_accum,
                "effective_batch": batch_size * grad_accum,
                "lr": lr,
                "lr_scheduler": SFT_LR_SCHEDULER,
                "warmup_steps": SFT_WARMUP_STEPS,
                "weight_decay": SFT_WEIGHT_DECAY,
                "optimizer": SFT_OPTIMIZER,
                "max_seq_len": max_seq_len,
                "max_steps": max_steps,
                "lora_r": lora_r,
                "lora_alpha": lora_alpha,
                "lora_dropout": lora_dropout,
                "lora_target_modules": list(LORA_TARGET_MODULES),
                "dataset_path": args.dataset,
                "val_dataset_path": args.val_dataset,
                "model": model_id,
                "seed": seed,
                "report_to": report_to,
                "eval_every": eval_every,
                "save_every": SFT_SAVE_EVERY,
                "log_every": SFT_LOG_EVERY,
                "early_stop_format": SFT_EARLY_STOP_FORMAT,
                "early_stop_correction": SFT_EARLY_STOP_CORRECTION,
                "early_stop_diversity": SFT_EARLY_STOP_DIVERSITY,
                "max_wall_seconds": SFT_MAX_WALL_SECONDS,
            },
        },
    )

    # ---- Preflight: refuse to run with the known-bad Unsloth+TF combo  #
    # (unsloth >= 2026.4.0) + (transformers < 4.55.0) silently misparses
    # the Qwen2.5-3B config: it instantiates a 7B-shaped model
    # (hidden=4096) and crashes when the 3B checkpoint (hidden=2048)
    # starts loading, with:
    #   RuntimeError: size mismatch for weight: copying a param with
    #   shape torch.Size([151936, 2048]) from checkpoint, the shape in
    #   current model is torch.Size([151936, 4096]).
    # We catch this BEFORE downloading >5GB of weights so the user does
    # not burn GPU minutes on a deterministic failure.
    import unsloth as _unsloth
    import transformers as _transformers

    def _parse_ver(v: str) -> tuple[int, ...]:
        out: list[int] = []
        for part in v.split("+", 1)[0].split("."):
            digits = "".join(ch for ch in part if ch.isdigit())
            out.append(int(digits) if digits else 0)
        return tuple(out)

    _u = _parse_ver(_unsloth.__version__)
    _t = _parse_ver(_transformers.__version__)
    _is_qwen25_3b = "qwen2.5-3b" in model_id.lower()
    _bad_combo = _u >= (2026, 4, 0) and _t < (4, 55, 0)
    if _is_qwen25_3b and _bad_combo:
        print(
            "[train_sft] FATAL: detected the unsloth/transformers combo that\n"
            f"           silently misparses {model_id} into a 7B-shaped model.\n"
            f"           Installed: unsloth=={_unsloth.__version__}  "
            f"transformers=={_transformers.__version__}\n"
            "           This exact pair produces the\n"
            "             'size mismatch ... [151936, 2048] vs [151936, 4096]'\n"
            "           error during model load on Lightning AI / Colab.\n"
            "           Fix: pin to a known-good combination, e.g.\n"
            "             pip install --no-deps --force-reinstall \\\n"
            "                 unsloth==2025.11.1 unsloth_zoo==2026.4.9\n"
            "             pip install --force-reinstall \\\n"
            "                 transformers==4.57.2 trl==0.20.0\n"
            "           Or re-run scripts/run_lightning_pipeline.sh which\n"
            "           pins these correctly and now hard-fails if the pins\n"
            "           do not stick.",
            file=sys.stderr,
        )
        return 1

    # ---- Load model + datasets --------------------------------------- #
    print(f"loading {model_id} via Unsloth (4-bit NF4)")
    print(f"  unsloth={_unsloth.__version__}  "
          f"transformers={_transformers.__version__}")
    model, tokenizer = FastLanguageModel.from_pretrained(
        model_name=model_id,
        max_seq_length=max_seq_len,
        load_in_4bit=True,
        dtype=None,  # Unsloth auto-selects bf16/fp16
    )
    model = FastLanguageModel.get_peft_model(
        model,
        r=lora_r,
        lora_alpha=lora_alpha,
        target_modules=list(LORA_TARGET_MODULES),
        lora_dropout=lora_dropout,
        bias="none",
        use_gradient_checkpointing="unsloth",
        random_state=seed,
    )

    print(f"loading train dataset from {args.dataset}")
    train_dataset = _load_train_dataset(args.dataset, tokenizer)
    print(f"  {len(train_dataset)} samples; first text len = "
          f"{len(train_dataset[0]['text'])}")

    val_records: list[dict] = []
    val_path = Path(args.val_dataset)
    if val_path.exists():
        val_records = _load_jsonl(args.val_dataset)
        print(f"loaded {len(val_records)} held-out validation records "
              f"from {args.val_dataset}")
    else:
        print(f"WARNING: no validation file at {args.val_dataset}; "
              f"running without eval / early-stop.")

    wandb_utils.log({
        "sft/train_dataset_size": len(train_dataset),
        "sft/val_dataset_size": len(val_records),
        "sft/first_text_len": len(train_dataset[0]["text"]),
    })

    # Dataset preview to W&B (sanity check the chat-template wrapping).
    wandb_utils.log_generation_table(
        [
            {"split": "train", "prompt": train_dataset[i]["prompt"][:600],
             "completion": train_dataset[i]["completion"]}
            for i in range(min(8, len(train_dataset)))
        ],
        step=0,
        table_name="sft/train_preview",
        columns=["split", "prompt", "completion"],
    )

    # ---- TrainingArguments (locked spec) ----------------------------- #
    Path(args.output).mkdir(parents=True, exist_ok=True)
    bf16_supported = (
        torch.cuda.is_available() and torch.cuda.is_bf16_supported()
    )
    training_args = TrainingArguments(
        output_dir=args.output,
        num_train_epochs=epochs,
        max_steps=max_steps,             # hard cap; wins over epochs
        per_device_train_batch_size=batch_size,
        gradient_accumulation_steps=grad_accum,
        learning_rate=lr,
        weight_decay=SFT_WEIGHT_DECAY,
        # Label smoothing was added in the 2026-04 SFT regularisation
        # rewrite (FIX 2) to combat mode collapse: spreading the loss
        # across non-target tokens makes the model less sharply rewarded
        # for memorising one canonical completion, which is what kept
        # output_diversity at 1 across every prior checkpoint.
        label_smoothing_factor=SFT_LABEL_SMOOTHING,
        warmup_steps=SFT_WARMUP_STEPS,
        lr_scheduler_type=SFT_LR_SCHEDULER,
        optim=SFT_OPTIMIZER,
        bf16=bf16_supported,
        fp16=torch.cuda.is_available() and not bf16_supported,
        logging_steps=SFT_LOG_EVERY,
        save_steps=SFT_SAVE_EVERY,
        save_total_limit=4,
        seed=seed,
        report_to=report_to,
        run_name=run_name,
    )

    # ---- Callbacks --------------------------------------------------- #
    started_wall = time.time()
    callbacks = [_build_loss_guard_callback()]
    eval_schedule = None if args.no_eval_schedule else SFT_EVAL_SCHEDULE
    val_cb = _build_validation_callback(
        model=model,
        tokenizer=tokenizer,
        val_records=val_records,
        eval_every=eval_every,
        eval_schedule=eval_schedule,
        print_sample_outputs=print_sample_outputs,
        output_dir=args.output,
        max_new_tokens=SFT_MAX_NEW_TOKENS,
        diversity_n_samples=args.diversity_samples,
        diversity_temperature=args.diversity_temperature,
        early_stop_format=SFT_EARLY_STOP_FORMAT,
        early_stop_correction=SFT_EARLY_STOP_CORRECTION,
        early_stop_diversity=SFT_EARLY_STOP_DIVERSITY,
        max_wall_seconds=SFT_MAX_WALL_SECONDS,
        started_wall=started_wall,
        # 2026-04 (FIX 2) diversity-collapse regression early stop.
        diversity_floor=SFT_PREFLIGHT_DIVERSITY_FLOOR,
        diversity_run_len=SFT_DIVERSITY_COLLAPSE_RUN_LEN,
    )
    if val_cb is not None:
        callbacks.append(val_cb)

    trainer = SFTTrainer(
        model=model,
        tokenizer=tokenizer,
        train_dataset=train_dataset,
        dataset_text_field="text",
        max_seq_length=max_seq_len,
        args=training_args,
        packing=False,
        callbacks=callbacks,
    )

    print(f"training (max_steps={max_steps}, eval_every={eval_every}) ...")
    train_result = trainer.train()
    elapsed = time.time() - started_wall
    metrics = getattr(train_result, "metrics", {}) or {}
    wandb_utils.update_summary({
        "sft/wall_seconds": elapsed,
        **{f"sft/final/{k}": v for k, v in metrics.items()
           if isinstance(v, (int, float))},
    })
    print(f"training finished in {elapsed:.1f}s "
          f"(max_wall_seconds={SFT_MAX_WALL_SECONDS:.0f})")

    print(f"saving adapters to {args.output}")
    model.save_pretrained(args.output)
    tokenizer.save_pretrained(args.output)

    # ---- Upload adapter as W&B artifact ------------------------------ #
    if not args.no_artifact:
        wandb_utils.log_artifact(
            args.output,
            name=f"sft-adapter-{run_name}",
            artifact_type="model",
            description="SFT-warmed Qwen2.5-3B + LoRA adapter (Qubit-Medic).",
        )

    wandb_utils.finish_run()
    print("done")
    return 0


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))