"""dist_align_gate.py — teacher-bank alignment verification + model loading for the distillation bed. The bank stores UNNORMALIZED pooled image features; every consumer L2-normalizes at load. Feature parity requires exact preprocessing: plain-GELU CLIP forward (NOT QuickGELU — the open_clip default for some checkpoints differs at the 0.975-cosine level), torchvision Resize(224, bicubic shortest-edge) + CenterCrop(224) + CLIP mean/std, CLS -> projection readout, fp32 compute. SigLIP uses its own processor (384x384 warp). Two hub checkpoints ship open_clip-format weights only (LAION B/16, DataComp B/32); this module converts them deterministically into transformers CLIPModel format, and reproduction of the stored bank features (cos >= 0.9999 on reference rows) is the proof of conversion. """ import argparse import glob import json import os import sys import zlib import numpy as np import torch import torch.nn.functional as F from PIL import Image torch.backends.cuda.matmul.allow_tf32 = False torch.backends.cudnn.allow_tf32 = False DEV = "cuda" if torch.cuda.is_available() else "cpu" DATA = os.path.join(os.environ.get("DIST_DATA_ROOT", "./data"), "dist") BANK = os.path.join(DATA, "bank") VAL_DIR = os.path.join(DATA, "coco", "val2017") CAPS_JSON = os.path.join(DATA, "coco", "annotations", "captions_val2017.json") GATE_CONFIGS = ("clip_b16_laion2b", "siglip_b16_384", "clip_b32_openai") CFG_MODEL = { "clip_b16_laion2b": "laion/CLIP-ViT-B-16-laion2B-s34B-b88K", "siglip_b16_384": "google/siglip-base-patch16-384", "clip_b32_openai": "openai/clip-vit-base-patch32", "clip_b32_laion2b": "laion/CLIP-ViT-B-32-laion2B-s34B-b79K", "clip_b32_datacomp": "laion/CLIP-ViT-B-32-DataComp.XL-s13B-b90K", "clip_b16_openai": "openai/clip-vit-base-patch16", "dinov2_b14": "facebook/dinov2-base", } PRIMARY = "clip_b16_laion2b" LAION_HF_DIR = os.path.join(DATA, "hf_clip_b16_laion2b") DATACOMP_HF_DIR = os.path.join(DATA, "hf_clip_b32_datacomp") def seed_for(name): return zlib.crc32(name.encode()) & 0x7FFFFFFF def _convert_open_clip(model_id, out_dir, patch): """open_clip-only hub checkpoints (laion b16, datacomp b32 — no transformers config.json, verified) -> deterministic key-map into an HF CLIPModel, saved once to `out_dir`. The bank-parity gate is the end-to-end verifier (a wrong map cannot produce cos ~1 vs stored bank features). ViT-B geometry is shared; only patch_size differs.""" from huggingface_hub import hf_hub_download from safetensors.torch import load_file from transformers import AutoTokenizer, CLIPConfig, CLIPImageProcessor, CLIPModel w = load_file(hf_hub_download(model_id, "open_clip_model.safetensors")) cfg = CLIPConfig( projection_dim=512, # laion-family open_clip ViT-B = PLAIN gelu (quick_gelu is the # OpenAI-checkpoint quirk and HF CLIPConfig's default — must override) vision_config=dict(hidden_size=768, intermediate_size=3072, num_hidden_layers=12, num_attention_heads=12, image_size=224, patch_size=patch, projection_dim=512, hidden_act="gelu"), text_config=dict(hidden_size=512, intermediate_size=2048, num_hidden_layers=12, num_attention_heads=8, max_position_embeddings=77, vocab_size=49408, projection_dim=512, hidden_act="gelu")) model = CLIPModel(cfg) sd = {} def blocks(src, dst, d): for i in range(12): p, q = f"{src}.resblocks.{i}.", f"{dst}.layers.{i}." for s, t in (("ln_1", "layer_norm1"), ("ln_2", "layer_norm2"), ("attn.out_proj", "self_attn.out_proj"), ("mlp.c_fc", "mlp.fc1"), ("mlp.c_proj", "mlp.fc2")): sd[q + t + ".weight"] = w[p + s + ".weight"] sd[q + t + ".bias"] = w[p + s + ".bias"] for j, t in enumerate(("q_proj", "k_proj", "v_proj")): sd[q + f"self_attn.{t}.weight"] = \ w[p + "attn.in_proj_weight"][j * d:(j + 1) * d] sd[q + f"self_attn.{t}.bias"] = \ w[p + "attn.in_proj_bias"][j * d:(j + 1) * d] blocks("visual.transformer", "vision_model.encoder", 768) blocks("transformer", "text_model.encoder", 512) sd["vision_model.embeddings.class_embedding"] = w["visual.class_embedding"] sd["vision_model.embeddings.position_embedding.weight"] = \ w["visual.positional_embedding"] sd["vision_model.embeddings.patch_embedding.weight"] = w["visual.conv1.weight"] pre = [k[:-7] for k in model.state_dict() if "pre_lay" in k and k.endswith(".weight")][0] # pre_layrnorm typo-BC sd[pre + ".weight"] = w["visual.ln_pre.weight"] sd[pre + ".bias"] = w["visual.ln_pre.bias"] sd["vision_model.post_layernorm.weight"] = w["visual.ln_post.weight"] sd["vision_model.post_layernorm.bias"] = w["visual.ln_post.bias"] sd["visual_projection.weight"] = w["visual.proj"].t().contiguous() sd["text_model.embeddings.token_embedding.weight"] = w["token_embedding.weight"] sd["text_model.embeddings.position_embedding.weight"] = w["positional_embedding"] sd["text_model.final_layer_norm.weight"] = w["ln_final.weight"] sd["text_model.final_layer_norm.bias"] = w["ln_final.bias"] sd["text_projection.weight"] = w["text_projection"].t().contiguous() sd["logit_scale"] = w["logit_scale"].reshape(()) missing, unexpected = model.load_state_dict(sd, strict=False) missing = [m for m in missing if "position_ids" not in m] assert not missing and not unexpected, (missing, unexpected) os.makedirs(out_dir, exist_ok=True) model.save_pretrained(out_dir) AutoTokenizer.from_pretrained(CFG_MODEL[PRIMARY]).save_pretrained(out_dir) CLIPImageProcessor().save_pretrained(out_dir) # defaults == the repo's # open_clip preprocess_cfg: 224 bicubic shortest-edge + center-crop, # OpenAI CLIP mean/std (verified against open_clip_config.json) print(f"[GATE] converted open_clip -> HF CLIPModel at {out_dir}") def laion_local_dir(): """Path to the transformers-format LAION B/16 (converted once, cached). dist_bed.py imports this for its text/vision teacher heads.""" if not os.path.isfile(os.path.join(LAION_HF_DIR, "model.safetensors")): _convert_open_clip(CFG_MODEL[PRIMARY], LAION_HF_DIR, 16) return LAION_HF_DIR def datacomp_local_dir(): """Transformers-format DataComp-XL B/32 (open_clip-only upstream).""" if not os.path.isfile(os.path.join(DATACOMP_HF_DIR, "model.safetensors")): _convert_open_clip(CFG_MODEL["clip_b32_datacomp"], DATACOMP_HF_DIR, 32) return DATACOMP_HF_DIR def load_bank(config, split="val"): """-> (ids int64 np, feats float32 torch UNNORMALIZED) in parquet row order.""" import pyarrow as pa import pyarrow.parquet as pq files = sorted(glob.glob(os.path.join(BANK, config, f"{split}-*.parquet"))) if not files: print(f"[GATE] MISSING bank parquet: {config}/{split} under {BANK}") sys.exit(2) t = pa.concat_tables([pq.read_table(f, columns=["image_id", "features"]) for f in files]) ids = t.column("image_id").to_numpy() fl = t.column("features").combine_chunks() D = fl.type.list_size feats = fl.flatten().to_numpy(zero_copy_only=False).reshape(len(ids), D) return ids, torch.from_numpy(np.array(feats, dtype=np.float32, copy=True)) def load_vision(config): """(model, image_processor) for a bank config, fp32 eval on DEV. The LAION primary loads from the local converted dir (open_clip-only upstream). Image processor only — SigLIP's slow tokenizer needs sentencepiece (absent in this venv); the image side never touches a tokenizer. THE BANK ACT FINDING (measured 2026-07-26, this gate's discovery run): the bank's CLIP towers were embedded with PLAIN GELU — for the OpenAI checkpoints that is the open_clip create_model('ViT-B-32', 'openai') silent QuickGELU mismatch (plain-gelu local re-embed: cos mean 1.00000 min 0.99999 vs quick_gelu 0.97515; geometry/precision/pooling/crop all falsified first). Preprocessing = the default 224 shortest-edge bicubic + center-crop + CLIP mean/std; siglip = its default 384 warp; storage is fp16-precision in fp32. The bed MUST match the bank, so every CLIP tower here runs hidden_act='gelu' (correct anyway for laion2b checkpoints).""" from transformers import AutoConfig, AutoImageProcessor, AutoModel if config == PRIMARY: name = laion_local_dir() elif config == "clip_b32_datacomp": name = datacomp_local_dir() else: name = CFG_MODEL[config] kw, proc = {}, None if config.startswith("clip_"): c = AutoConfig.from_pretrained(name) c.vision_config.hidden_act = "gelu" c.text_config.hidden_act = "gelu" kw["config"] = c # bank CLIP geometry is torchvision-exact (measured 1.00000/0.99999 # on b32 vs 0.99919 through the HF processor — crop-rounding differs) proc = clip_torchvision_pipe(224) else: proc = AutoImageProcessor.from_pretrained(name) model = _from_pretrained_fp32(name, **kw).to(DEV).eval() return model, proc def _from_pretrained_fp32(name, **kw): """transformers 5.x takes dtype=, 4.x takes torch_dtype= — accept both (the pod runs 4.49 pinned under torch 2.4.1; local runs 5.x).""" from transformers import AutoModel try: return AutoModel.from_pretrained(name, dtype=torch.float32, **kw) except TypeError: return AutoModel.from_pretrained(name, torch_dtype=torch.float32, **kw) CLIP_MEAN = (0.48145466, 0.4578275, 0.40821073) CLIP_STD = (0.26862954, 0.26130258, 0.27577711) class clip_torchvision_pipe: """Callable matching the AutoImageProcessor interface for the ONE call shape this gate uses: proc(images=[PIL...], return_tensors='pt').""" def __init__(self, size): from torchvision import transforms as T from torchvision.transforms import InterpolationMode self.tf = T.Compose([ T.Resize(size, interpolation=InterpolationMode.BICUBIC), T.CenterCrop(size), T.ToTensor(), T.Normalize(CLIP_MEAN, CLIP_STD)]) def __call__(self, images, return_tensors="pt", **kw): return {"pixel_values": torch.stack([self.tf(im) for im in images])} def load_siglip_text_tokenizer(): """SigLIP fast tokenizer from the repo's own tokenizer.json (no sentencepiece dependency). Exported for dist_bed.py's siglip text head.""" from huggingface_hub import hf_hub_download from transformers import PreTrainedTokenizerFast tj = hf_hub_download(CFG_MODEL["siglip_b16_384"], "tokenizer.json") return PreTrainedTokenizerFast(tokenizer_file=tj, pad_token="", eos_token="", unk_token="", model_max_length=64) def _feat(out): """transformers 5.x returns the full output object from get_*_features (projected features live in .pooler_output); 4.x returned the tensor.""" return out if torch.is_tensor(out) else out.pooler_output @torch.no_grad() def embed_images(model, proc, pils, bs=16): outs = [] for i in range(0, len(pils), bs): px = proc(images=pils[i:i + bs], return_tensors="pt")["pixel_values"] outs.append(_feat(model.get_image_features( pixel_values=px.to(DEV, torch.float32))).float().cpu()) return torch.cat(outs) @torch.no_grad() def embed_texts(model, tok, texts, bs=64, siglip=False): outs = [] for i in range(0, len(texts), bs): kw = (dict(padding="max_length", truncation=True, max_length=64) if siglip else dict(padding=True, truncation=True)) tk = tok(texts[i:i + bs], return_tensors="pt", **kw) tk = {k: v.to(DEV) for k, v in tk.items() if k in ("input_ids", "attention_mask")} outs.append(_feat(model.get_text_features(**tk)).float().cpu()) return torch.cat(outs) def first_captions(): if not os.path.isfile(CAPS_JSON): print(f"[GATE] MISSING captions json: {CAPS_JSON}") sys.exit(2) with open(CAPS_JSON, "r", encoding="utf-8") as f: anns = json.load(f)["annotations"] cap = {} for a in anns: # first caption in annotation order cap.setdefault(int(a["image_id"]), a["caption"].strip()) return cap def main(): ap = argparse.ArgumentParser() ap.add_argument("--n", type=int, default=64) ap.add_argument("--n-text", type=int, default=256) ap.add_argument("--margin", type=float, default=0.05) a, _ = ap.parse_known_args() if DEV == "cuda": torch.cuda.set_per_process_memory_fraction(0.73) if not os.path.isdir(VAL_DIR) or not glob.glob(os.path.join(VAL_DIR, "*.jpg")): print(f"[GATE] val2017 images absent/empty at {VAL_DIR} — still " f"downloading? Re-run once populated.") sys.exit(2) rows, any_fail = [], False for cfg in GATE_CONFIGS: ids, feats = load_bank(cfg) g = torch.Generator().manual_seed(seed_for(f"dist-align:{cfg}")) sel = torch.randperm(len(ids), generator=g)[:a.n].numpy() pils = [] for i in sel: p = os.path.join(VAL_DIR, f"{int(ids[i]):012d}.jpg") if not os.path.isfile(p): print(f"[GATE] missing jpg {p} — val2017 incomplete?") sys.exit(2) pils.append(Image.open(p).convert("RGB")) model, proc = load_vision(cfg) loc = F.normalize(embed_images(model, proc, pils), dim=-1) bank = F.normalize(feats[sel], dim=-1) cos = (loc * bank).sum(-1) mean_c, min_c = float(cos.mean()), float(cos.min()) verdict = ("PASS" if mean_c >= 0.999 else "WARN" if mean_c >= 0.99 else "FAIL") any_fail |= verdict == "FAIL" rows.append((cfg, feats.shape[1], len(sel), mean_c, min_c, verdict)) del model, proc if DEV == "cuda": torch.cuda.empty_cache() # ---- LAION text-head sanity ------------------------------------------ from transformers import AutoModel, AutoTokenizer ldir = laion_local_dir() tok = AutoTokenizer.from_pretrained(ldir) model = _from_pretrained_fp32(ldir).to(DEV).eval() probes = ["a photo of a cat", "a photo of a dog", "a photo of a car"] praw = embed_texts(model, tok, probes) raw_norms = praw.norm(dim=-1) pnorm = F.normalize(praw, dim=-1) norm_dev = float((pnorm.norm(dim=-1) - 1.0).abs().max()) ids, feats = load_bank(PRIMARY) cap = first_captions() g = torch.Generator().manual_seed(seed_for("dist-align-text")) sel = torch.randperm(len(ids), generator=g)[:a.n_text].numpy() texts = [cap[int(ids[i])] for i in sel] temb = F.normalize(embed_texts(model, tok, texts), dim=-1) img = F.normalize(feats[sel], dim=-1) matched = float((temb * img).sum(-1).mean()) shuffled = float((temb.roll(1, 0) * img).sum(-1).mean()) margin = matched - shuffled tverdict = ("PASS" if (margin >= a.margin and norm_dev < 1e-3 and bool(torch.isfinite(praw).all())) else "FAIL") any_fail |= tverdict == "FAIL" print("=== DIST ALIGNMENT GATE (local re-embed vs stored bank, val2017) ===") print(f"{'config':<18} {'dim':>4} {'n':>4} {'mean_cos':>9} {'min_cos':>9} verdict") for cfg, D, n, mc, mn, v in rows: print(f"{cfg:<18} {D:>4} {n:>4} {mc:>9.5f} {mn:>9.5f} {v}") print(f"text-head sanity (laion b16): probe raw norms " f"[{float(raw_norms.min()):.3f}..{float(raw_norms.max()):.3f}] " f"normalized max|1-n| {norm_dev:.1e}; " f"caption cos matched {matched:.4f} vs shuffled {shuffled:.4f} " f"margin {margin:.4f} (bar {a.margin}) {tverdict}") overall = "FAIL" if any_fail else "PASS" print(f"OVERALL: {overall}") sys.exit(1 if any_fail else 0) if __name__ == "__main__": main()