fix(calib): repair uninitialized rotary inv_freq + skip tied lm_head target

inv_freq (non-persistent, no GGUF counterpart) was left uninitialized by
meta-build + to_empty -> silent wrong-RoPE on full-attn layers (NaN under
use_deterministic_algorithms, which enables fill_uninitialized_memory).
reinit_rotary_buffers() rebuilds it in fp32 every run. Also skip lm_head as
an ml8 target when tie_word_embeddings=True (served by the tied FP8
token_embd; was quantizing an uninitialized buffer = garbage/NaN, 379s wasted).

Verified: deterministic 0.8B calibration now bit-reproducible (187/187 blobs
byte-identical across 3 runs), full 150 ML8 + 37 FP8 coverage, 0 Cholesky
failures (was 21/151). Adds diag_nan_probe.py + diff_calib_blobs.py.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Author: kmbandy

scripts/calibration/calibrate_ml8_paged.py

@@ -53,12 +53,13 @@
 # ── Deterministic calibration — env half (OPT-IN via ML8_DETERMINISTIC=1; MUST precede `import torch`) ───
 # ml8 GPTQ calibration was nondeterministic: the GPU Hessian forward used unforced reduction order
 # and GPTQ's sequential error-feedback amplified the low-bit noise into DIFFERENT weight assignments
-# run-to-run (~0.6 PPL spread — swamps the ±0.05 levers). These flags pin it bit-identically on the
-# linear-attn/MLP path (verified). BUT torch.use_deterministic_algorithms(True) currently breaks
-# full-model calibration: from the FIRST full-self-attention layer onward, the SDPA forward under
-# the hammer yields a degenerate (NaN-suspected) Hessian → GPTQ Cholesky fails → tensors silently
-# skipped → partial-coverage model. Root-cause + fix is deferred (see KG / handoff), so determinism
-# is OPT-IN until then; the default path is the known-good nondeterministic calibrator.
+# run-to-run (~0.6 PPL spread — swamps the ±0.05 levers). These flags pin it bit-identically.
+# VERIFIED 2026-06-02: full-model deterministic calibration is bit-reproducible (187/188 tensors
+# byte-identical across two independent runs; the rest are FP8). The earlier "breaks at the first
+# self-attn layer" failure was NOT a determinism bug — use_deterministic_algorithms(True) enables
+# fill_uninitialized_memory, which surfaced a latent bug: the rotary inv_freq buffer (no GGUF
+# counterpart) was left uninitialized by meta+to_empty → NaN-filled → garbage RoPE. Fixed by
+# reinit_rotary_buffers() (runs every calibration; also repairs the silent non-deterministic case).
 _DETERMINISTIC = os.environ.get("ML8_DETERMINISTIC") == "1"
 if _DETERMINISTIC:
     os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8")  # required for deterministic cuBLAS/hipBLAS GEMM
@@ -86,7 +87,7 @@
         try: setattr(torch.backends.cuda.matmul, _attr, _val)
         except AttributeError: pass
     print("[determinism] OPT-IN ENABLED: use_deterministic_algorithms(True) + seeded + pinned GEMM "
-          "— WARNING: breaks full-model coverage at the first self-attn layer (Cholesky/SDPA bug, WIP)")
+          "— full-model calibration is bit-reproducible (rotary inv_freq reinit applied post-load)")
 
 # Re-use existing calibration helpers (compute_hessian / gptq_quantize_linear / eval_ppl_wikitext / etc.)
 sys.path.insert(0, str(Path(__file__).parent))
@@ -237,7 +238,19 @@ def find_dense_full_targets(model, coverage: str = "ffn"):
                     yield name, mod, Tier.ML8
 
     # Global ML8 roles (lm_head / eh_proj), if present as nn.Linear.
+    # Skip lm_head when the model TIES embeddings: a tied lm_head IS token_embd
+    # (already covered by the FP8 embedding tier), has no independent GGUF tensor
+    # (the converter correctly emits no output.weight), and meta+to_empty leaves
+    # its weight UNINITIALIZED — so quantizing it wastes time on silent garbage
+    # (NaN under determinism) and produces a blob the converter discards anyway.
+    _cfg = getattr(model, "config", None)
+    _tied = bool(getattr(getattr(_cfg, "text_config", _cfg), "tie_word_embeddings", False)
+                 or getattr(_cfg, "tie_word_embeddings", False))
     for role in _FULL_ML8_GLOBAL_ORDER:
+        if role == "lm_head" and _tied:
+            print("[targets] skipping lm_head ML8 target: tie_word_embeddings=True "
+                  "(served by the FP8-tier tied token_embd at inference)")
+            continue
         for name, mod in by_name.items():
             if name.rsplit(".", 1)[-1] != role:
                 continue
@@ -649,6 +662,46 @@ def load_resident_to_model(model: nn.Module, gguf_path: str,
     return n_loaded
 
 
+def reinit_rotary_buffers(model: nn.Module, device: str) -> int:
+    """Recompute every text rotary-embedding ``inv_freq`` buffer from config.
+
+    inv_freq is a NON-PERSISTENT buffer derived from rope_theta — it has no GGUF
+    counterpart, so the resident loader (which walks ``named_parameters``) never
+    fills it, and ``model.to_empty()`` leaves it UNINITIALIZED. The damage is
+    silent without determinism: uninitialized memory is finite garbage, and since
+    cos()/sin() clamp to [-1, 1] the forward never NaNs — it just applies WRONG
+    positional encoding to every full-attention layer (linear-attn layers don't
+    use RoPE). With ``use_deterministic_algorithms(True)`` that same uninitialized
+    memory is NaN-filled, so the first RoPE turns q/k into NaN and the whole
+    forward (and every downstream Hessian) blows up. Both are the SAME bug; the
+    fix is to rebuild inv_freq the way the module's __init__ does and keep it in
+    fp32 (HF never downcasts it; our blanket .to(dtype) did). Idempotent — safe to
+    run on every calibration. Returns the number of rotary modules repaired.
+    """
+    n = 0
+    for mod in model.modules():
+        cls = mod.__class__.__name__
+        if "RotaryEmbedding" not in cls or "Vision" in cls:
+            continue   # text rotary only; vision rotary has a different __init__
+        if not hasattr(mod, "inv_freq") or getattr(mod, "config", None) is None:
+            continue
+        try:
+            fresh = type(mod)(mod.config, device=torch.device(device))
+        except Exception as e:   # noqa: BLE001 — never let a probe-fix abort calibration
+            print(f"[rotary-fix] WARN: could not rebuild {cls} inv_freq: {e}")
+            continue
+        with torch.no_grad():
+            # assign (not copy_) so the buffer keeps fresh's fp32 dtype, not the
+            # bf16 the model-wide .to(dtype) left it as.
+            mod.inv_freq = fresh.inv_freq.to(device=torch.device(device))
+            if hasattr(mod, "original_inv_freq") and hasattr(fresh, "original_inv_freq"):
+                mod.original_inv_freq = fresh.original_inv_freq.to(device=torch.device(device))
+        if hasattr(fresh, "attention_scaling"):
+            mod.attention_scaling = fresh.attention_scaling
+        n += 1
+    return n
+
+
 def _detect_moe_n_experts(config) -> int | None:
     """Return n_experts if `config` describes an MoE model, else None.
 
@@ -1163,6 +1216,73 @@ def main():
     # installed. MUST run before the first forward. ───
     apply_fla_arch_shim(model, args.device)
 
+    # ─── Repair rotary inv_freq: meta-build + to_empty leaves this non-persistent,
+    # non-GGUF buffer uninitialized → silently-wrong RoPE (and a hard NaN under
+    # determinism). MUST run before the first forward, on every calibration. ───
+    _n_rotary = reinit_rotary_buffers(model, args.device)
+    print(f"[rotary-fix] reinitialized inv_freq on {_n_rotary} text rotary module(s)")
+
+    # ─── NaN probe (opt-in ML8_NAN_PROBE=1): hook every decoder layer + the submodules
+    # of the first full_attention layer; on the FIRST forward, print per-module
+    # in-nan/out-nan/out-inf/finite-absmax in completion order, then self-remove. Used to
+    # pinpoint the determinism-induced NaN in the first full-attn layer (deferred bug). ───
+    if os.environ.get("ML8_NAN_PROBE") == "1":
+        _dec = [m for _n, m in model.named_modules()
+                if m.__class__.__name__.endswith("DecoderLayer")]
+        _lt = list(getattr(getattr(model.config, "text_config", model.config),
+                           "layer_types", ["?"] * len(_dec)))
+        _first_full = next((i for i, t in enumerate(_lt) if "full" in str(t)), 3)
+        _pb = {"ev": [], "done": False, "h": []}
+
+        def _pstat(t):
+            if isinstance(t, (tuple, list)) and t and torch.is_tensor(t[0]):
+                t = t[0]
+            if not torch.is_tensor(t):
+                return (None, None, float("nan"))
+            tf = t.float()
+            fin = tf[torch.isfinite(tf)]
+            return (int(torch.isnan(tf).sum()), int(torch.isinf(tf).sum()),
+                    float(fin.abs().max()) if fin.numel() else float("nan"))
+
+        def _mk(label):
+            def _h(mod, inp, out):
+                if _pb["done"]:
+                    return
+                it = inp[0] if isinstance(inp, (tuple, list)) and inp else inp
+                inn = int(torch.isnan(it.float()).sum()) if torch.is_tensor(it) else -1
+                nan, inf, amax = _pstat(out)
+                _pb["ev"].append((label, inn, nan, inf, amax))
+            return _h
+
+        for i, m in enumerate(_dec):
+            _pb["h"].append(m.register_forward_hook(
+                _mk(f"LAYER[{i:02d}] ({_lt[i] if i < len(_lt) else '?'})")))
+            if i == _first_full:
+                for sn, sm in m.named_modules():
+                    if sn:
+                        _pb["h"].append(sm.register_forward_hook(
+                            _mk(f"  L{i}.{sn} <{sm.__class__.__name__}>")))
+
+        def _dump(mod, inp, out):
+            if _pb["done"]:
+                return
+            _pb["done"] = True
+            print("\n=== [NAN-PROBE] first forward, completion order ===")
+            fb = None
+            for label, inn, nan, inf, amax in _pb["ev"]:
+                bad = (nan and nan > 0) or (inf and inf > 0)
+                mark = "  <<< FIRST NaN/Inf" if bad and fb is None else ""
+                if bad and fb is None:
+                    fb = label
+                print(f"{label:<50} in_nan={inn} out_nan={nan} out_inf={inf} "
+                      f"absmax={amax:.4g}{mark}")
+            print(f"[NAN-PROBE] first NaN/Inf at: {fb}\n", flush=True)
+            for h in _pb["h"]:
+                h.remove()
+        _dec[-1].register_forward_hook(_dump)
+        print(f"[NAN-PROBE] installed: {len(_dec)} decoder layers + submodules of "
+              f"layer {_first_full} (first full_attention)")
+
     # ─── Calibration corpus + baseline PPL (paged forward proves the swap works) ───
     print(f"[calib] loading {'budget ' + str(args.token_budget) + ' tok' if args.token_budget else str(args.n_samples) + ' samples'} "
           f"seq_len={args.seq_len} corpus={args.corpus}")

scripts/calibration/diag_nan_probe.py

@@ -0,0 +1,164 @@
+#!/usr/bin/env python3
+"""Pinpoint the determinism-induced NaN in the first full-attention layer of the 0.8B
+dense-hybrid Qwen3.5 bed. Honors ML8_DETERMINISTIC=1 (replicates the driver's determinism
+block EXACTLY, before torch import). Loads the model resident (0.8B fits), registers:
+  (a) a forward hook on EVERY decoder layer  -> first layer whose output goes NaN
+  (b) a forward hook on EVERY submodule of the first full_attention layer -> the exact op
+Runs ONE forward over a small real-token batch and prints, in execution-completion order,
+each module's  input-nan? / output-nan? / output-inf? / finite-absmax.
+
+Usage:  ML8_DETERMINISTIC=1 /usr/bin/python3 diag_nan_probe.py      # hammer ON
+        /usr/bin/python3 diag_nan_probe.py                          # hammer OFF (control)
+"""
+import os
+import sys
+
+# ── replicate the driver determinism block (env half MUST precede torch import) ──
+_DET = os.environ.get("ML8_DETERMINISTIC") == "1"
+if _DET:
+    os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8")
+    os.environ.setdefault("HIPBLASLT_DETERMINISTIC", "1")
+
+import torch  # noqa: E402
+from transformers import AutoModelForCausalLM, AutoTokenizer  # noqa: E402
+
+if _DET:
+    torch.manual_seed(0)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(0)
+    torch.use_deterministic_algorithms(True)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+    for _attr, _val in (("allow_tf32", False),
+                        ("allow_fp16_reduced_precision_reduction", False),
+                        ("allow_bf16_reduced_precision_reduction", False)):
+        try: setattr(torch.backends.cuda.matmul, _attr, _val)
+        except AttributeError: pass
+    try: torch.backends.cudnn.allow_tf32 = False
+    except AttributeError: pass
+
+print(f"[probe] ML8_DETERMINISTIC={'1 (HAMMER ON)' if _DET else '0 (control)'}", flush=True)
+
+MODEL = "/home/kmbandy/models/Qwen3.5-0.8B-hf"
+DEVICE = "cuda:0"
+DTYPE = torch.bfloat16
+
+tok = AutoTokenizer.from_pretrained(MODEL)
+model = AutoModelForCausalLM.from_pretrained(MODEL, dtype=DTYPE).to(DEVICE).eval()
+
+# match the driver: apply the RDNA fla fp32-scan shim (env ML8_FLA_SHIM=1)
+if os.environ.get("ML8_FLA_SHIM") == "1":
+    from fla_compat import apply_fla_arch_shim
+    apply_fla_arch_shim(model, DEVICE)
+
+# ── locate the decoder-layer list (robust to multimodal carrier nesting) ──
+def find_layers(m):
+    for name, mod in m.named_modules():
+        if mod.__class__.__name__.endswith("DecoderLayer"):
+            # parent ModuleList: strip trailing ".<idx>"
+            return name.rsplit(".", 1)[0], mod
+    raise RuntimeError("no *DecoderLayer found")
+
+first_layer_name, _ = find_layers(model)
+layers = dict(model.named_modules())[first_layer_name]
+print(f"[probe] decoder layers at '{first_layer_name}'  n={len(layers)}", flush=True)
+
+# layer_types from config to label which is the first full_attention
+lt = getattr(model.config, "text_config", model.config)
+layer_types = getattr(lt, "layer_types", None)
+if layer_types is None:
+    layer_types = ["?"] * len(layers)
+first_full = next((i for i, t in enumerate(layer_types) if "full" in t), 3)
+print(f"[probe] layer_types[:8]={layer_types[:8]}  first_full_attention=layer {first_full}", flush=True)
+
+events = []  # (label, in_nan, out_nan, out_inf, finite_absmax)
+
+def stat(t):
+    if not torch.is_tensor(t):
+        # tuple/list output: inspect the first tensor element
+        if isinstance(t, (tuple, list)) and len(t) and torch.is_tensor(t[0]):
+            t = t[0]
+        else:
+            return (None, None, None, None)
+    tf = t.float()
+    n = int(torch.isnan(tf).sum())
+    inf = int(torch.isinf(tf).sum())
+    finite = tf[torch.isfinite(tf)]
+    amax = float(finite.abs().max()) if finite.numel() else float("nan")
+    return (n, inf, amax, t.numel())
+
+def mk_hook(label):
+    def h(mod, inp, out):
+        in_t = inp[0] if isinstance(inp, (tuple, list)) and len(inp) else inp
+        in_nan = int(torch.isnan(in_t.float()).sum()) if torch.is_tensor(in_t) else -1
+        n, inf, amax, numel = stat(out)
+        events.append((label, in_nan, n, inf, amax, numel))
+    return h
+
+handles = []
+# (a) every decoder layer
+for i in range(len(layers)):
+    handles.append(layers[i].register_forward_hook(mk_hook(f"LAYER[{i:02d}] ({layer_types[i]})")))
+# (b) every submodule of the first full_attention layer
+for sub_name, sub_mod in layers[first_full].named_modules():
+    if sub_name == "":
+        continue
+    handles.append(sub_mod.register_forward_hook(
+        mk_hook(f"  L{first_full}.{sub_name} <{sub_mod.__class__.__name__}>")))
+
+# ── input: either a benign synthetic sentence, or the EXACT driver calib cache ──
+samples = []
+_calib_pt = os.environ.get("ML8_CALIB_PT")
+if _calib_pt:
+    obj = torch.load(_calib_pt, map_location="cpu", weights_only=False)
+    raw = obj["samples"] if isinstance(obj, dict) and "samples" in obj else obj
+    for s in raw:
+        t = s if torch.is_tensor(s) else torch.as_tensor(s)
+        if t.dim() == 1:
+            t = t.unsqueeze(0)
+        samples.append(t.to(DEVICE))
+    print(f"[probe] loaded {len(samples)} cached calib samples from {_calib_pt}", flush=True)
+else:
+    text = ("The history of scientific discovery is a long and winding road, "
+            "full of unexpected turns and surprising connections between fields. ") * 8
+    samples = [tok(text, return_tensors="pt", truncation=True,
+                   max_length=512).input_ids.to(DEVICE)]
+
+print(f"[probe] forward over {len(samples)} sample(s), shapes={[tuple(s.shape) for s in samples[:3]]}...",
+      flush=True)
+
+def first_bad_in(ev):
+    for rec in ev:
+        label, in_nan, n, inf, amax, numel = rec
+        if (n and n > 0) or (inf and inf > 0):
+            return label
+    return None
+
+bad_sample = None
+bad_events = None
+with torch.no_grad():
+    for si, ids in enumerate(samples):
+        events.clear()
+        model(input_ids=ids)
+        fb = first_bad_in(events)
+        print(f"[probe] sample {si} shape={tuple(ids.shape)}  first_NaN={fb}", flush=True)
+        if fb is not None and bad_sample is None:
+            bad_sample, bad_events = si, list(events)
+            break
+
+for h in handles:
+    h.remove()
+
+# ── full trace of the FIRST sample that went bad (or the last clean one) ──
+trace = bad_events if bad_events is not None else events
+print(f"\n=== forward trace (sample {bad_sample if bad_sample is not None else 'last'}, completion order) ===")
+print(f"{'module':<52}{'in_nan':>8}{'out_nan':>9}{'out_inf':>9}{'absmax':>14}")
+first_bad = None
+for label, in_nan, n, inf, amax, numel in trace:
+    bad = (n and n > 0) or (inf and inf > 0)
+    mark = "  <<< FIRST NaN/Inf" if bad and first_bad is None else ""
+    if bad and first_bad is None:
+        first_bad = label
+    print(f"{label:<52}{str(in_nan):>8}{str(n):>9}{str(inf):>9}{amax:>14.4g}{mark}")
+
+print(f"\n[probe] OVERALL first NaN/Inf: sample={bad_sample} module={first_bad}", flush=True)

scripts/calibration/diff_calib_blobs.py

@@ -0,0 +1,47 @@
+#!/usr/bin/env python3
+"""Bit-reproducibility check for two ml8 calibration runs. For every matching blob
+in dirA and dirB, compare the quantization-defining tensors (indices, centroids,
+scales) for EXACT equality. Prints per-blob match and an overall verdict.
+
+Usage: diff_calib_blobs.py <dirA> <dirB>
+"""
+import sys
+import pathlib
+import torch
+
+dirA, dirB = pathlib.Path(sys.argv[1]), pathlib.Path(sys.argv[2])
+KEYS = ("indices", "centroids_per_group", "scale_per_group")
+
+blobsA = sorted(p.name for p in dirA.glob("*.pt"))
+blobsB = set(p.name for p in dirB.glob("*.pt"))
+
+n_match = n_diff = n_missing = 0
+diffs = []
+for name in blobsA:
+    if name not in blobsB:
+        n_missing += 1
+        diffs.append(f"  MISSING in B: {name}")
+        continue
+    a = torch.load(dirA / name, map_location="cpu", weights_only=False)
+    b = torch.load(dirB / name, map_location="cpu", weights_only=False)
+    blob_ok = True
+    for k in KEYS:
+        if k in a or k in b:
+            ta, tb = a.get(k), b.get(k)
+            if ta is None or tb is None or ta.shape != tb.shape or not torch.equal(ta, tb):
+                blob_ok = False
+                diffs.append(f"  DIFF {name}: tensor '{k}' differs")
+    if blob_ok:
+        n_match += 1
+    else:
+        n_diff += 1
+
+print(f"blobs compared: {len(blobsA)}  identical: {n_match}  differing: {n_diff}  missing: {n_missing}")
+for d in diffs[:20]:
+    print(d)
+if n_diff == 0 and n_missing == 0:
+    print("\n✅ REPRODUCIBLE — every blob's indices/centroids/scales are bit-identical across runs.")
+    sys.exit(0)
+else:
+    print("\n❌ NOT reproducible — see diffs above.")
+    sys.exit(1)