esh-pfi-infrastructure/services/meromero-quant/perlayer_test.py

"""Attempt-4 blocker: is dropping `per_layer_config` the correct fix, or must we
force `allow_global_per_layer_attribute_access`?

The two candidates are NOT equivalent:
  * DROP  -> config becomes homogeneous in transformers' eyes and the global
             `global_head_dim` / `num_global_key_value_heads` fields describe the
             full-attention layers, exactly as zerofata's canonical config does.
  * FORCE -> config stays heterogeneous; `config.head_dim` starts answering 256 to
             every caller, including the ones building the 512-wide full-attention
             layers. That is the hazard transformers' own warning names.

So this is not a "did the traceback go away" test. It builds the model on the meta
device from each candidate and reads the ACTUAL k_proj widths back, against the
checkpoint's measured 2048 (full) / 4096 (sliding). A candidate that constructs but
mis-shapes a layer is a worse outcome than the crash, because it would ship.

Positive control: zerofata's canonical config, which we already quantized
successfully on 2026-08-21, MUST pass every check here. If it doesn't, the
instrument is broken and none of the negatives mean anything.
"""
import json, shutil, tempfile, traceback
from pathlib import Path

import torch
import transformers
from transformers import AutoConfig

print(f"transformers {transformers.__version__}  torch {torch.__version__}", flush=True)

HERETIC = Path("/tank/aimodels/G4-MeroMero-v2-31B-heretic-bf16")
CANON = Path("/tank/aimodels/meromero-v2-nvfp4-work/src")

# Measured off the checkpoints by shape_verify.py; both trees agree.
EXPECT = {"full_attention": 2048, "sliding_attention": 4096}


def variant(name, cfg_dict):
    d = Path(tempfile.mkdtemp(prefix=f"cfg-{name}-"))
    (d / "config.json").write_text(json.dumps(cfg_dict))
    return name, d


heretic = json.loads((HERETIC / "config.json").read_text())
canon = json.loads((CANON / "config.json").read_text())

dropped = json.loads(json.dumps(heretic))
dropped["text_config"].pop("per_layer_config")

forced = json.loads(json.dumps(heretic))
forced["text_config"]["allow_global_per_layer_attribute_access"] = True

variants = [
    variant("A-canonical-POSITIVE-CONTROL", canon),
    variant("B-heretic-asis", heretic),
    variant("C-heretic-drop-per_layer_config", dropped),
    variant("D-heretic-force-global-access", forced),
]

for name, d in variants:
    print(f"\n=== {name} ===", flush=True)
    try:
        cfg = AutoConfig.from_pretrained(d)
    except Exception as e:
        print(f"  CONFIG FAILED: {type(e).__name__}: {str(e)[:160]}")
        continue
    t = cfg.text_config
    fields = {}
    for k in ("head_dim", "num_key_value_heads", "global_head_dim",
              "num_global_key_value_heads"):
        try:
            fields[k] = getattr(t, k, "<absent>")
        except Exception as e:
            fields[k] = f"<{type(e).__name__}>"
    print(f"  config OK: {fields}")

    try:
        from transformers import Gemma4ForConditionalGeneration as M
        with torch.device("meta"):
            model = M(cfg)
    except Exception:
        print("  MODEL BUILD FAILED:")
        print("    " + traceback.format_exc().strip().replace("\n", "\n    ")[-1200:])
        continue

    layer_types = t.layer_types
    probes = [next(i for i, x in enumerate(layer_types) if x == "full_attention"),
              next(i for i, x in enumerate(layer_types) if x == "sliding_attention")]
    layers = model.model.language_model.layers
    verdict = []
    for li in probes:
        got = tuple(layers[li].self_attn.k_proj.weight.shape)
        want = EXPECT[layer_types[li]]
        ok = got[0] == want
        verdict.append(ok)
        print(f"  L{li:>2} {layer_types[li]:<18} k_proj {got}  "
              f"want out={want}  {'OK' if ok else '*** MISMATCH ***'}")
    print(f"  => {'GEOMETRY MATCHES CHECKPOINT' if all(verdict) else 'GEOMETRY WRONG'}")

for _, d in variants:
    shutil.rmtree(d, ignore_errors=True)