"""Cross-architecture parity check: does the lens recipe hold for a given model?
Exports the (backbone, lens) pair to a temp dir and asserts:
1. torch: lens(last block output) == model logits (the identity the whole demo rests on)
2. ONNX fp32 matches torch
3. f16-stored weights keep top-1 at every position of the probe
Usage:
python check_arch_parity.py --model gpt2
python check_arch_parity.py --model HuggingFaceTB/SmolLM2-135M
python check_arch_parity.py --model EleutherAI/pythia-70m --revision step1000
"""
import logging
import tempfile
from pathlib import Path
import fire
import numpy as np
import onnxruntime as ort
import torch
from pydantic import BaseModel
from transformers import AutoModelForCausalLM, AutoTokenizer
from lenslapse.arch import resolve
from lenslapse.export_checkpoints import PROBE, Backbone, build_lens_onnx, export_backbone_onnx
from lenslapse.logging_utils import configure_cli_logging
from lenslapse.onnx_f16 import save_f16
logger = logging.getLogger(__name__)
[docs]
class CheckArchParityConfig(BaseModel):
"""Validated arguments for `check_arch_parity`; see `main`'s docstring for what each means."""
model: str
revision: str = "main"
[docs]
def main(model: str, revision: str = "main") -> None:
"""Check that the lens recipe (lens(last block output) == model logits) holds for a model.
Args:
model: HF id or local directory.
revision: HF revision to check, e.g. a training-checkpoint tag.
"""
cfg = CheckArchParityConfig(model=model, revision=revision)
tok = AutoTokenizer.from_pretrained(cfg.model, revision=cfg.revision)
model_obj = AutoModelForCausalLM.from_pretrained(cfg.model, revision=cfg.revision, dtype=torch.float32)
model_obj.eval()
handles = resolve(model_obj)
logger.info("%s@%s: norm=%s layers=%d", cfg.model, cfg.revision, handles.norm_type, len(handles.layers))
enc = tok(PROBE, return_tensors="pt")
input_ids, attn = enc["input_ids"], enc["attention_mask"]
with torch.no_grad():
ref = model_obj(input_ids=input_ids, attention_mask=attn).logits[0]
hs = Backbone(model_obj)(input_ids, attn)
lens_final = handles.lm_head(handles.final_norm(hs[-1, 0]))
torch_diff = float((ref - lens_final).abs().max())
logger.info("torch lens(last)==logits max diff: %.2e", torch_diff)
assert torch_diff < 1e-3, "lens identity violated — unsupported (e.g. post-norm) architecture"
with tempfile.TemporaryDirectory() as td:
bb, lens = Path(td) / "bb.onnx", Path(td) / "lens.onnx"
export_backbone_onnx(Backbone(model_obj), input_ids, attn, bb)
build_lens_onnx(model_obj, lens)
sb, sl = ort.InferenceSession(str(bb)), ort.InferenceSession(str(lens))
h = sb.run(None, {"input_ids": input_ids.numpy(), "attention_mask": attn.numpy()})[0]
lo = sl.run(None, {"hidden": h[-1, 0]})[0]
onnx_diff = float(np.abs(ref.numpy() - lo).max())
onnx_top1 = bool((lo.argmax(-1) == ref.numpy().argmax(-1)).all())
logger.info("onnx fp32 max diff: %.2e top-1 match: %s", onnx_diff, onnx_top1)
# late checkpoints have large logit magnitudes; match export_checkpoints.py's loose bound
assert onnx_diff < 0.05 and onnx_top1
bb16, lens16 = Path(td) / "bb16.onnx", Path(td) / "lens16.onnx"
save_f16(str(bb), str(bb16))
save_f16(str(lens), str(lens16))
sb16, sl16 = ort.InferenceSession(str(bb16)), ort.InferenceSession(str(lens16))
h16 = sb16.run(None, {"input_ids": input_ids.numpy(), "attention_mask": attn.numpy()})[0]
lo16 = sl16.run(None, {"hidden": h16[-1, 0]})[0]
top1 = bool((lo16.argmax(-1) == ref.numpy().argmax(-1)).all())
logger.info("f16 top-1 match at all positions: %s", top1)
assert top1
# the completion signal, not a diagnostic message: scripts pipe stdout and grep for this,
# so it must stay on print() rather than move to the logger (which defaults to stderr).
print("PARITY_OK")
if __name__ == "__main__":
configure_cli_logging()
fire.Fire(main)