load: one backbone loader, so a vision-language model stops being unr…

pyproject.toml

@@ -93,7 +93,7 @@ pollard-envmatch = "pollard_envmatch:main"
 [tool.setuptools]
 package-dir = { "" = "tools" }
 py-modules = ["pollard_auto", "pollard_calc", "pollard_fit", "pollard_run", "pollard_fit_dit",
-    "pollard_flybrain", "pollard_brainverify", "pollard_brainlanes", "pollard_brain_backends", "pollard_connectome", "pollard_experts", "pollard_route", "pollard_sensitivity", "pollard_smooth", "pollard_rotate", "pollard_precondition", "pollard_eval", "pollard_health", "pollard_pack", "pollard_prune", "pollard_bench", "pollard_export", "pollard_gptq", "pollard_automap", "pollard_abliterate", "pollard_probe", "pollard_probes", "pollard_scorecard", "pollard_kl", "pollard_lowbit", "pollard_verify", "pollard_vllm", "pollard_doctor", "pollard_hf_smooth", "pollard_mx", "pollard_mlx", "pollard_exl3", "pollard_calib", "pollard_palette", "pollard_ls", "pollard_workspace",
+    "pollard_flybrain", "pollard_load", "pollard_brainverify", "pollard_brainlanes", "pollard_brain_backends", "pollard_connectome", "pollard_experts", "pollard_route", "pollard_sensitivity", "pollard_smooth", "pollard_rotate", "pollard_precondition", "pollard_eval", "pollard_health", "pollard_pack", "pollard_prune", "pollard_bench", "pollard_export", "pollard_gptq", "pollard_automap", "pollard_abliterate", "pollard_probe", "pollard_probes", "pollard_scorecard", "pollard_kl", "pollard_lowbit", "pollard_verify", "pollard_vllm", "pollard_doctor", "pollard_hf_smooth", "pollard_mx", "pollard_mlx", "pollard_exl3", "pollard_calib", "pollard_palette", "pollard_ls", "pollard_workspace",
     "pollard_serve_eval", "pollard_card", "pollard_onboard", "pollard_envmatch", "imatrix_fix_gate", "exl3_fix_mtp_ehproj",
     "pollard_archfp", "pollard_errtype", "pollard_errsrc", "pollard_recard",
     "pollard_ggufcompat", "pollard_reclaim", "pollard_exl3_band", "pollard_refcheck",

tests/test_recipes.py

@@ -1415,16 +1415,27 @@ class Qwen2VLConfig for this kind of AutoModel" -- so a VL model could not be lo
         print("    (skipped: torch not installed -- `pip install pollard-weights[flybrain]`)")
         return
     import pollard_flybrain as F
+    import pollard_load as L
 
-    assert hasattr(F, "load_backbone"), "the VL-capable loader is missing"
-    src = pathlib.Path(F.__file__).read_text(encoding="utf-8")
+    assert hasattr(F, "load_backbone") and hasattr(L, "load_backbone")
+    src = pathlib.Path(L.__file__).read_text(encoding="utf-8")
     for cls in ("AutoModelForImageTextToText", "AutoModelForVision2Seq"):
         assert cls in src, f"no fallback to {cls}: a VL model would be unreachable"
-    assert "AutoModelForCausalLM.from_pretrained(a.model" not in src, \
-        "main() still loads the backbone directly, bypassing the VL fallback"
-    # the text stack of a VL model lives under language_model -- the finder must still look there
     assert "model.language_model" in src, "the VL text-stack path was dropped"
 
+    # ONE loader. Eleven tools each called AutoModelForCausalLM directly and each died on a VL
+    # config; a second copy is how they drift back apart.
+    tools = pathlib.Path(__file__).resolve().parent.parent / "tools"
+    offenders = []
+    for f in sorted(tools.glob("pollard_*.py")):
+        if f.name in ("pollard_load.py", "pollard_route.py"):
+            continue
+        body = f.read_text(encoding="utf-8")
+        if "AutoModelForCausalLM.from_pretrained" in body:
+            offenders.append(f.name)
+    assert not offenders, ("these load a backbone directly and will refuse a VL model: "
+                           + ", ".join(offenders))
+
 
 
 def test_brain_query_default_matches_the_verified_prompt():

tools/pollard_abliterate.py

@@ -156,7 +156,8 @@ def main():
         except Exception:
             sys.exit("ERROR: pass --out for the abliterated model.")
 
-    from transformers import AutoModelForCausalLM, AutoTokenizer
+    from transformers import AutoTokenizer
+    from pollard_load import load_backbone
     import pollard_workspace as ws
     trc = ws.resolve_trust_remote_code(a.model, a.trust_remote_code)
     dev = a.device if (a.device != "mps" or torch.backends.mps.is_available()) else "cpu"
@@ -165,8 +166,7 @@ def main():
     if tok.pad_token is None:
         tok.pad_token = tok.eos_token
     tok.padding_side = "left"
-    model = AutoModelForCausalLM.from_pretrained(a.model, dtype=torch.float16,
-                                                trust_remote_code=trc).to(dev).eval()
+    model = load_backbone(a.model, torch.float16, dev, trust_remote_code=trc)
 
     A = _load_lines(a.harmful) if a.harmful else _SMOKE_A
     B = _load_lines(a.harmless) if a.harmless else _SMOKE_B

tools/pollard_doctor.py

@@ -33,9 +33,10 @@
 def scan_outliers(source_dir, device, calib, rows, cols, thresh):
     """Find massive-activation input channels per layer on the fp16 source (the low-bit break predictor)."""
     import torch
-    from transformers import AutoModelForCausalLM, AutoTokenizer
+    from transformers import AutoTokenizer
+    from pollard_load import load_backbone
     tok = AutoTokenizer.from_pretrained(source_dir)
-    model = AutoModelForCausalLM.from_pretrained(source_dir, dtype=torch.float16, device_map=device).eval()
+    model = load_backbone(source_dir, torch.float16, device_map=device)
     try:
         layers = model.model.layers
     except AttributeError:

tools/pollard_flybrain.py

@@ -114,38 +114,10 @@ def _progress(*args, **kw):
     print(*args, **kw)
 
 
-def load_backbone(model_id: str, dtype=None, device: str = "cpu"):
-    """Load any backbone a brain can attach to -- text-only or vision-language.
-
-    AutoModelForCausalLM refuses a vision-language config outright ("Unrecognized configuration
-    class ... for this kind of AutoModel"), so a VL model could not be reached at all even though
-    everything downstream already handles one: _find_stack looks for `model.language_model`, which
-    is exactly where a VL model keeps its text stack, and the token codes come from the same output
-    embedding either way. The brain attaches to the LANGUAGE side of a VL model; the vision tower is
-    untouched, like the rest of the frozen backbone.
-    """
-    import transformers
-    from transformers import AutoModelForCausalLM
-
-    dtype = torch.float32 if dtype is None else dtype
-    try:
-        model = AutoModelForCausalLM.from_pretrained(model_id, dtype=dtype)
-    except ValueError as text_only_err:
-        model = None
-        for name in ("AutoModelForImageTextToText", "AutoModelForVision2Seq"):
-            cls = getattr(transformers, name, None)
-            if cls is None:
-                continue
-            try:
-                model = cls.from_pretrained(model_id, dtype=dtype)
-                break
-            except Exception:
-                continue
-        if model is None:
-            raise SystemExit(
-                f"could not load {model_id!r} as a causal LM or as a vision-language model.\n"
-                f"  {text_only_err}") from None
-    return model.to(device).eval()
+def load_backbone(model_id: str, dtype=None, device: str = "cpu", **kw):
+    """Re-exported from pollard_load so there is ONE loader, not two that drift apart."""
+    from pollard_load import load_backbone as _lb
+    return _lb(model_id, dtype, device, **kw)
 
 
 def probe_basis(model, tok, corpus: str, n_probe: int, device, plen: int = 160):

tools/pollard_gptq.py

@@ -366,13 +366,14 @@ def main():
                     "required to quantize a model bigger than VRAM (e.g. a 7B on 16 GB)")
     a = ap.parse_args()
 
-    from transformers import AutoModelForCausalLM, AutoTokenizer
+    from transformers import AutoTokenizer
+    from pollard_load import load_backbone
     dev = a.device if (a.device != "mps" or torch.backends.mps.is_available()) else "cpu"
     mdev = "cpu" if a.offload else dev                    # where the model itself lives
     print(f"== pollard-gptq :: {a.model}  W{a.bits}g{a.groupsize}  dev={dev}"
           f"{'  (block-offload)' if a.offload else ''}")
     tok = AutoTokenizer.from_pretrained(a.model)
-    model = AutoModelForCausalLM.from_pretrained(a.model, dtype=torch.float16).to(mdev)
+    model = load_backbone(a.model, torch.float16, mdev, eval_mode=False)
 
     print("loading wikitext-2 ...")
     calib = get_wikitext(tok, "train", a.seqlen, a.nsamples, path=a.calib_file)

tools/pollard_hf_smooth.py

@@ -56,11 +56,10 @@ def main():
     import torch
     import pollard_workspace as ws
     trc = ws.resolve_trust_remote_code(a.model, a.trust_remote_code)
-    from transformers import AutoModelForCausalLM, AutoTokenizer
+    from transformers import AutoTokenizer
+    from pollard_load import load_backbone
     tok = AutoTokenizer.from_pretrained(a.model, trust_remote_code=trc)
-    model = AutoModelForCausalLM.from_pretrained(a.model, dtype=torch.float16, device_map=a.device,
-                                                trust_remote_code=trc)
-    model.eval()
+    model = load_backbone(a.model, torch.float16, device_map=a.device, trust_remote_code=trc)
     try:
         layers = model.model.layers
     except AttributeError:

tools/pollard_kl.py

@@ -54,11 +54,12 @@ def main():
     ap.add_argument("--device", default="mps")
     a = ap.parse_args()
 
-    from transformers import AutoModelForCausalLM, AutoTokenizer
+    from transformers import AutoTokenizer
+    from pollard_load import load_backbone
     dev = a.device if (a.device != "mps" or torch.backends.mps.is_available()) else "cpu"
     print(f"== pollard-kl :: {a.model}  recipe={a.recipe} qmode={a.qmode}  dev={dev}", flush=True)
     tok = AutoTokenizer.from_pretrained(a.model)
-    ref = AutoModelForCausalLM.from_pretrained(a.model, dtype=torch.float16).to(dev).eval()
+    ref = load_backbone(a.model, torch.float16, dev)
     calib = _chunks(tok, open(a.calib_file, encoding="utf-8").read(), a.seqlen, a.nsamples)
     ev = _chunks(tok, open(a.eval_file, encoding="utf-8").read(), a.seqlen, a.kl_chunks)
 

tools/pollard_load.py

@@ -0,0 +1,69 @@
+#!/usr/bin/env python3
+"""Load any backbone a Pollard tool needs -- text-only or vision-language.
+
+AutoModelForCausalLM refuses a vision-language config outright:
+
+    ValueError: Unrecognized configuration class Qwen2VLConfig for this kind of
+    AutoModel: AutoModelForCausalLM
+
+Eleven tools called it directly, so every one of them died on a VL model even though nothing they do
+afterwards cares -- probing, smoothing, abliterating and KL all work on the text stack, which a VL
+model keeps under `model.language_model`. The failure reads as an unsupported model rather than an
+unasked question, and it sent a whole quant ladder into the ditch.
+
+One loader, used everywhere: try the causal LM, fall back to the vision-language auto classes, and
+report BOTH failures if neither works rather than only the last one.
+"""
+from __future__ import annotations
+
+
+def load_backbone(model_id: str, dtype=None, device: str = "cpu", eval_mode: bool = True,
+                  **kw):
+    """Load `model_id` under whichever auto class accepts it.
+
+    dtype defaults to float32; pass torch.float16 etc for the lanes that want it. Extra kwargs
+    (trust_remote_code, attn_implementation, ...) pass straight through.
+    """
+    import torch
+    import transformers
+    from transformers import AutoModelForCausalLM
+
+    if dtype is None:
+        dtype = torch.float32
+    try:
+        model = AutoModelForCausalLM.from_pretrained(model_id, dtype=dtype, **kw)
+    except ValueError as text_only_err:
+        model = None
+        errs = [f"AutoModelForCausalLM: {text_only_err}"]
+        for name in ("AutoModelForImageTextToText", "AutoModelForVision2Seq"):
+            cls = getattr(transformers, name, None)
+            if cls is None:
+                continue
+            try:
+                model = cls.from_pretrained(model_id, dtype=dtype, **kw)
+                break
+            except Exception as e:
+                errs.append(f"{name}: {e}")
+        if model is None:
+            raise SystemExit(f"could not load {model_id!r} as a causal LM or a vision-language "
+                             "model:\n  " + "\n  ".join(str(e)[:160] for e in errs)) from None
+    model = model.to(device)
+    return model.eval() if eval_mode else model
+
+
+def text_layers(model):
+    """The decoder layer list, wherever this family keeps it.
+
+    A VL model's text stack lives under `model.language_model`; a text model's is `model.model`.
+    Tools that count layers or iterate them need the same answer for both.
+    """
+    for path in ("model.language_model", "language_model.model", "model"):
+        o = model
+        try:
+            for part in path.split("."):
+                o = getattr(o, part)
+            if hasattr(o, "layers"):
+                return o.layers
+        except AttributeError:
+            continue
+    raise ValueError("could not locate a decoder stack on this model")

tools/pollard_lowbit.py

@@ -155,10 +155,11 @@ def main():
     ap.add_argument("--device", default="cpu")
     a = ap.parse_args()
 
-    from transformers import AutoModelForCausalLM, AutoTokenizer
+    from transformers import AutoTokenizer
+    from pollard_load import load_backbone
     print(f"== pollard-lowbit :: {a.model}  bits={a.bits} keep={a.keep} levels={a.levels}", flush=True)
     tok = AutoTokenizer.from_pretrained(a.model)
-    model = AutoModelForCausalLM.from_pretrained(a.model, dtype=torch.float16).to(a.device)
+    model = load_backbone(a.model, torch.float16, a.device, eval_mode=False)
     calib = _chunks(tok, open(a.calib_file, encoding="utf-8").read(), a.seqlen, a.nsamples)
     test = _chunks(tok, open(a.eval_file, encoding="utf-8").read(), a.seqlen, a.eval_chunks)
 

tools/pollard_palette.py

@@ -111,11 +111,12 @@ def main():
     ap.add_argument("--device", default="mps")
     a = ap.parse_args()
 
-    from transformers import AutoModelForCausalLM, AutoTokenizer
+    from transformers import AutoTokenizer
+    from pollard_load import load_backbone
     dev = a.device if (a.device != "mps" or torch.backends.mps.is_available()) else "cpu"
     print(f"== pollard-palette :: {a.model}  gs={a.groupsize}  dev={dev}", flush=True)
     tok = AutoTokenizer.from_pretrained(a.model)
-    model = AutoModelForCausalLM.from_pretrained(a.model, dtype=torch.float16).to(dev)
+    model = load_backbone(a.model, torch.float16, dev, eval_mode=False)
     calib = _chunks(tok, open(a.calib_file, encoding="utf-8").read(), a.seqlen, a.nsamples)
     test = _chunks(tok, open(a.eval_file, encoding="utf-8").read(), a.seqlen, a.eval_chunks)
 

tools/pollard_probe.py

@@ -142,12 +142,13 @@ def main():
                          "models too big to run layersxgroups forward passes (744B-scale)")
     a = ap.parse_args()
 
-    from transformers import AutoModelForCausalLM, AutoTokenizer
+    from transformers import AutoTokenizer
+    from pollard_load import load_backbone
     dev = a.device if (a.device != "mps" or torch.backends.mps.is_available()) else "cpu"
     groups = [g.strip() for g in a.groups.split(",") if g.strip()]
     print(f"== pollard-probe :: {a.model}  probe={a.probe_bits}bit  dev={dev}", flush=True)
     tok = AutoTokenizer.from_pretrained(a.model)
-    model = AutoModelForCausalLM.from_pretrained(a.model, dtype=torch.float16).to(dev).eval()
+    model = load_backbone(a.model, torch.float16, dev)
     layers = len(model.model.layers)
     ch = _chunks(tok, open(a.eval, encoding="utf-8").read(), a.seqlen, a.chunks)
 

tools/pollard_refcheck.py

@@ -149,13 +149,13 @@ def reference_nll(model_id, rows, max_rows=16, seq=1024, dtype="bfloat16", devic
     explained away as "hard rows"; a rising profile cannot.
     """
     import torch
-    from transformers import AutoModelForCausalLM, AutoTokenizer
+    from transformers import AutoTokenizer
+    from pollard_load import load_backbone
 
     tok = AutoTokenizer.from_pretrained(model_id)
-    model = AutoModelForCausalLM.from_pretrained(
-        model_id, dtype=getattr(torch, dtype),
+    model = load_backbone(
+        model_id, getattr(torch, dtype),
         device_map=device or ("cuda" if torch.cuda.is_available() else "cpu"))
-    model.eval()
 
     tot, ntok, head, head_n, tail, tail_n = 0.0, 0, 0.0, 0, 0.0, 0
     with torch.no_grad():