testbed/bench_speed.py

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