"""DELIVERABLE 1 -- inference speed benchmark: ORIGINAL teacher vs BANK student. Teacher = bigcode/starcoder2-3b, bf16, GPU. Student = same model, every layer's MLP replaced by trained E=2048 Bank (loaded from /tmp/banks_cmp_feat.pt; uses the Bank class from /tmp/train_compress2.py). If the checkpoint won't load, build fresh untrained E=2048 banks and SAY SO (speed only). Measures, in bf16, warmed up, median over several runs: (a) generation: batch=1, 256-token prompt -> generate 256 tokens -> tokens/sec and median per-token latency (b) prefill : batch=8, ctx=512 forward pass -> tokens/sec Reports each model's total param count and peak VRAM, prints a table. """ import importlib.util, json, statistics, time import torch from transformers import AutoModelForCausalLM, AutoTokenizer spec = importlib.util.spec_from_file_location("tc2", "/tmp/train_compress2.py") tc2 = importlib.util.module_from_spec(spec); spec.loader.exec_module(tc2) Bank = tc2.Bank import torch.nn.functional as F class Bank16(Bank): """Same trained weights, but compute in native bf16 (no fp32 upcast). Shows the decode speed achievable once the training-code fp32 path is dropped for deployment. Numerically near-identical for inference.""" def forward(self, x): act = F.gelu(x @ self.down.t() + self.b, approximate="tanh") return act @ self.up + self.obias DEV = 0 MODEL = "bigcode/starcoder2-3b" CKPT = "/tmp/banks_cmp_feat.pt" def load_model(): m = AutoModelForCausalLM.from_pretrained(MODEL, dtype=torch.bfloat16, device_map={"": DEV}) m.config.use_cache = True m.eval() for p in m.parameters(): p.requires_grad_(False) return m def n_params(m): return sum(p.numel() for p in m.parameters()) def install_banks(m, cls=Bank, dtype=None): """Replace each layer.mlp with a trained Bank (cls). Returns (used_trained, E).""" layers = m.model.layers used_trained = True E = 2048 try: ck = torch.load(CKPT, map_location="cpu") states = ck["banks"]; E = ck["E"] assert len(states) == len(layers) for l, sd in zip(layers, states): bk = cls(l.mlp, E, "random") bk.load_state_dict(sd) bk = bk.to(DEV) if dtype is not None: bk = bk.to(dtype) l.mlp = bk except Exception as e: print(f"!! checkpoint load failed ({str(e)[:80]}); building FRESH UNTRAINED E={E} banks (speed only)", flush=True) used_trained = False for l in layers: bk = cls(l.mlp, E, "random").to(DEV) if dtype is not None: bk = bk.to(dtype) l.mlp = bk return used_trained, E @torch.no_grad() def bench_generation(m, tok, n_runs=5, prompt_len=256, gen_len=256): torch.cuda.reset_peak_memory_stats(DEV) ids = torch.randint(0, tok.vocab_size, (1, prompt_len), device=DEV) # warmup for _ in range(2): m.generate(ids, max_new_tokens=16, do_sample=False, use_cache=True, pad_token_id=tok.eos_token_id or 0) torch.cuda.synchronize(DEV) tps, lat = [], [] for _ in range(n_runs): torch.cuda.synchronize(DEV); t0 = time.time() out = m.generate(ids, max_new_tokens=gen_len, min_new_tokens=gen_len, do_sample=False, use_cache=True, pad_token_id=tok.eos_token_id or 0) torch.cuda.synchronize(DEV); dt = time.time() - t0 new = out.shape[1] - prompt_len tps.append(new / dt); lat.append(dt / new * 1000.0) peak = torch.cuda.max_memory_allocated(DEV) / 1e9 return statistics.median(tps), statistics.median(lat), peak @torch.no_grad() def bench_prefill(m, n_runs=5, batch=8, ctx=512): torch.cuda.reset_peak_memory_stats(DEV) ids = torch.randint(0, 49000, (batch, ctx), device=DEV) for _ in range(2): m(ids) torch.cuda.synchronize(DEV) tps = [] for _ in range(n_runs): torch.cuda.synchronize(DEV); t0 = time.time() m(ids) torch.cuda.synchronize(DEV); dt = time.time() - t0 tps.append(batch * ctx / dt) peak = torch.cuda.max_memory_allocated(DEV) / 1e9 return statistics.median(tps), peak def run_one(label, build_student, cls=Bank, dtype=None): torch.cuda.empty_cache(); torch.cuda.reset_peak_memory_stats(DEV) tok = AutoTokenizer.from_pretrained(MODEL) m = load_model() note = "" if build_student: used_trained, E = install_banks(m, cls=cls, dtype=dtype) note = f"E={E} {'trained' if used_trained else 'UNTRAINED'}" params = n_params(m) gen_tps, gen_lat, gen_peak = bench_generation(m, tok) pf_tps, pf_peak = bench_prefill(m) peak = max(gen_peak, pf_peak) del m; torch.cuda.empty_cache() return dict(label=label, note=note, params=params, gen_tps=gen_tps, gen_lat=gen_lat, prefill_tps=pf_tps, peak_vram=peak) def main(): torch.cuda.set_device(DEV); torch.cuda.init() res = {} res["teacher"] = run_one("teacher (original)", build_student=False) res["student"] = run_one("student (bank, fp32 mm)", build_student=True) res["student_bf16"] = run_one("student (bank, bf16 mm)", build_student=True, cls=Bank16, dtype=torch.bfloat16) json.dump(res, open("speed_results.json", "w"), indent=2) print("\n" + "=" * 78) print("SPEED BENCHMARK -- StarCoder2-3b : teacher vs E=2048 bank student (bf16, GPU)") print("=" * 78) hdr = f"{'model':<22}{'params':>12}{'gen tok/s':>11}{'lat ms/tok':>12}{'prefill tok/s':>15}{'peak VRAM':>11}" print(hdr); print("-" * len(hdr)) for k in ("teacher", "student", "student_bf16"): r = res[k] print(f"{r['label']:<22}{r['params']/1e9:>11.3f}B{r['gen_tps']:>11.1f}" f"{r['gen_lat']:>12.2f}{r['prefill_tps']:>15.0f}{r['peak_vram']:>9.2f}GB") t, s, s16 = res["teacher"], res["student"], res["student_bf16"] print("-" * len(hdr)) print(f"student note: {s['note']}") print(f"params: {t['params']/1e9:.3f}B -> {s['params']/1e9:.3f}B " f"({(1-s['params']/t['params'])*100:.1f}% smaller)") print(f"fp32-mm bank gen {s['gen_tps']/t['gen_tps']:.2f}x prefill {s['prefill_tps']/t['prefill_tps']:.2f}x") print(f"bf16-mm bank gen {s16['gen_tps']/t['gen_tps']:.2f}x prefill {s16['prefill_tps']/t['prefill_tps']:.2f}x") print("wrote speed_results.json") if __name__ == "__main__": main()