src/exp_claim2b_draft_length.py

#!/usr/bin/env python3
"""
Official Claim 2 (part B): "SPEED-Bench identifies batch-size dependent optimal
draft lengths ... across speculative decoding evaluation."

Toy/calibrated reproduction. We (1) MEASURE the per-token acceptance rate alpha
of the distilgpt2->gpt2 draft/target pair on real SPEED-Bench prompts, and
(2) MEASURE the target model's forward latency L(batch b, seq s) on this
machine. We then combine them in the standard speculative-decoding throughput
model and show the *optimal draft length k\*(b) decreases as batch size b grows*:
long drafts win in the memory-bound (small-batch) regime, short drafts win once
verification becomes compute-bound (large batch).

Throughput per accepted token ~ E[generated_per_round(k,alpha)] / latency(k,b),
with E[gen] = (1 - alpha^{k+1})/(1 - alpha) (Leviathan et al. 2023) and
latency(k,b) = k * L_draft(b) + L_target(b, ctx+k).
Label: TOY (regime/crossover reproduced with measured alpha + measured latency;
exact crossover is model/engine/hardware specific).
"""
import json, os, argparse, time
import numpy as np, torch

def alpha_from_meanAL(mean_al, k):
    """Invert E[AL] = alpha(1-alpha^k)/(1-alpha) for per-token accept prob alpha,
    using the acceptance length measured in the Claim-1 experiment (real prompts)."""
    def E(a):
        return k * a if abs(a - 1) < 1e-9 else a * (1 - a**k) / (1 - a)
    lo, hi = 1e-6, 0.999999
    for _ in range(100):
        mid = 0.5 * (lo + hi)
        if E(mid) < mean_al: lo = mid
        else: hi = mid
    return 0.5 * (lo + hi)

@torch.no_grad()
def measure_latency(model, b, s, device, reps=5):
    x = torch.randint(0, 50257, (b, s), device=device)
    for _ in range(2):  # warmup
        model(x)
    if device == "cuda":
        torch.cuda.synchronize()
    t = time.time()
    for _ in range(reps):
        model(x)
    if device == "cuda":
        torch.cuda.synchronize()
    return (time.time() - t) / reps

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--kmeas", type=int, default=8)
    ap.add_argument("--n_prompts", type=int, default=12)
    ap.add_argument("--max_new", type=int, default=24)
    ap.add_argument("--ctx", type=int, default=64)
    ap.add_argument("--batches", default="1,2,4,8,16,32,64")
    ap.add_argument("--device", default="cpu")
    ap.add_argument("--out", default="outputs/claim2b_draft_length.json")
    a = ap.parse_args()
    dev = a.device
    torch.manual_seed(0)
    t0 = time.time()
    from transformers import AutoModelForCausalLM
    tgt = AutoModelForCausalLM.from_pretrained("gpt2").to(dev).eval()
    drf = AutoModelForCausalLM.from_pretrained("distilgpt2").to(dev).eval()

    # (1) acceptance rate alpha inverted from Claim-1 measured real-prompt mean AL
    c1 = json.load(open("outputs/claim1_throughput_bias.json", encoding="utf-8"))
    mean_al_real = c1["real"]["mean_AL"]; k_c1 = c1["config"]["k"]
    alpha = alpha_from_meanAL(mean_al_real, k_c1)
    acc, cmp = c1["config"]["n_real_rounds"], k_c1

    # (2) measured latencies L(b, s) for target and draft
    batches = [int(x) for x in a.batches.split(",")]
    Lt, Ld = {}, {}
    for b in batches:
        Lt[b] = measure_latency(tgt, b, a.ctx + a.kmeas, dev)
        Ld[b] = measure_latency(drf, b, a.ctx, dev)

    # (3) throughput model: optimal k per batch
    ks = list(range(1, a.kmeas + 1))
    def egen(k): return (1 - alpha**(k + 1)) / (1 - alpha)
    perbatch = {}
    for b in batches:
        tp = {}
        for k in ks:
            lat = k * Ld[b] + Lt[b]           # draft k steps + one verify
            tp[k] = egen(k) / lat
        kstar = max(ks, key=lambda k: tp[k])
        perbatch[b] = {"kstar": kstar,
                       "throughput_by_k": {str(k): tp[k] for k in ks},
                       "L_target_s": Lt[b], "L_draft_s": Ld[b]}
    kstars = {b: perbatch[b]["kstar"] for b in batches}
    decreasing = all(kstars[batches[i]] >= kstars[batches[i + 1]] for i in range(len(batches) - 1))
    strict_drop = kstars[batches[0]] > kstars[batches[-1]]

    res = {"alpha": alpha, "alpha_source": "inverted from Claim-1 real mean AL=%.3f at k=%d" % (mean_al_real, k_c1),
           "kmeas": a.kmeas, "device": dev, "ctx": a.ctx,
           "kstar_by_batch": {str(b): kstars[b] for b in batches},
           "optimal_draft_decreases_with_batch": bool(decreasing and strict_drop),
           "monotone_nonincreasing": bool(decreasing),
           "per_batch": {str(b): perbatch[b] for b in batches},
           "runtime_s": time.time() - t0}
    os.makedirs("outputs", exist_ok=True)
    json.dump(res, open(a.out, "w"), indent=2)
    print("alpha=%.4f (%d/%d)" % (alpha, acc, cmp))
    print("k*(batch):", {b: kstars[b] for b in batches})
    print("decreases with batch:", res["optimal_draft_decreases_with_batch"])
    print("saved", a.out, "in %.1fs" % res["runtime_s"])

if __name__ == "__main__":
    main()