File size: 3,114 Bytes
ae34de5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 | #!/usr/bin/env python3
"""Held-out perplexity eval for Bits vs Standard checkpoints (supports encoder variants).
Same held-out set, same metric exp(mean per-token CE). For Bits variants, pass the SAME
--enc_type/--ffn_depth/--W/--nbits used at training so the module structure matches for loading."""
import argparse, glob, os, math, sys, torch
os.environ['HF_HUB_OFFLINE'] = '1'
os.environ['TRANSFORMERS_OFFLINE'] = '1'
import torch.nn.functional as F
sys.path.insert(0, '/root')
from transformers import AutoConfig, AutoModelForCausalLM
from safetensors.torch import load_file
def build(a):
cfg = AutoConfig.from_pretrained(a.ckpt)
if a.arm == 'standard':
return AutoModelForCausalLM.from_pretrained(
a.ckpt, torch_dtype=torch.bfloat16, attn_implementation='sdpa').cuda().eval()
from bits_llm import swap_in_bits
model = AutoModelForCausalLM.from_config(cfg, torch_dtype=torch.bfloat16, attn_implementation='sdpa')
swap_in_bits(model, a.nbits, a.W, a.seed, enc_type=a.enc_type, ffn_depth=a.ffn_depth, out_scale=0.0)
sd = {}
for f in glob.glob(os.path.join(a.ckpt, '*.safetensors')):
sd.update(load_file(f))
missing, unexpected = model.load_state_dict(sd, strict=False)
enc_keys = [k for k in sd if 'embed_tokens' in k or '.ffn.' in k or 'in_proj' in k or 'out_proj' in k or 'tfenc' in k]
print(f"[load] tensors={len(sd)} enc_keys={len(enc_keys)} missing={len(missing)} unexpected={len(unexpected)}", flush=True)
if len(enc_keys) == 0:
print("[load][FATAL] no encoder weights in checkpoint", flush=True); sys.exit(3)
return model.cuda().eval()
@torch.no_grad()
def evaluate(model, seqs, bs):
tot_loss, tot_tok = 0.0, 0
for i in range(0, len(seqs), bs):
batch = torch.tensor(seqs[i:i+bs]).cuda()
logits = model(input_ids=batch).logits[:, :-1, :].float()
labels = batch[:, 1:]
loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), labels.reshape(-1), reduction='sum')
tot_loss += loss.item(); tot_tok += labels.numel()
ml = tot_loss / tot_tok
return ml, math.exp(ml)
def main():
ap = argparse.ArgumentParser()
ap.add_argument('--arm', required=True, choices=['bits', 'standard'])
ap.add_argument('--ckpt', required=True)
ap.add_argument('--data', default='/root/heldout_c4val_4000x512.pt')
ap.add_argument('--nbits', type=int, default=256)
ap.add_argument('--W', type=int, default=512)
ap.add_argument('--enc_type', default='mlp', choices=['mlp', 'tfenc'])
ap.add_argument('--ffn_depth', type=int, default=2)
ap.add_argument('--seed', type=int, default=42)
ap.add_argument('--n_seq', type=int, default=4000)
ap.add_argument('--bs', type=int, default=8)
a = ap.parse_args()
model = build(a)
seqs = torch.load(a.data)[:a.n_seq]
ml, ppl = evaluate(model, seqs, a.bs)
print(f"RESULT arm={a.arm} ckpt={a.ckpt} enc={a.enc_type} depth={a.ffn_depth} W={a.W} nbits={a.nbits} "
f"tokens={len(seqs)*len(seqs[0])} eval_loss={ml:.4f} ppl={ppl:.2f}", flush=True)
if __name__ == '__main__':
main()
|