| |
| """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() |
|
|
