Spaces:
Running on Zero
Running on Zero
File size: 5,842 Bytes
371d90c | 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 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 | """Parallel decoding must be lossless (greedy) and distribution-preserving (sampling)."""
from __future__ import annotations
import random
import pytest
import torch
from conftest import ARCHS
from decoding.common import make_generator, sample
from decoding.parallel import (
expected_tokens,
ngram_lookup,
predicted_speedup,
run_jacobi,
run_pld,
run_speculative,
verify,
)
def rand_prompt(n: int, vocab: int, seed: int = 0) -> list[int]:
rng = random.Random(seed)
return [rng.randrange(3, vocab) for _ in range(n)]
@pytest.mark.parametrize("arch", ARCHS)
@pytest.mark.parametrize("gamma", [1, 3, 5])
def test_speculative_greedy_equals_autoregressive(tiny_family, arch, gamma):
fam = tiny_family(arch, scale=0.08)
P = {"max_new_tokens": 40, "gamma": gamma, "mode": "greedy", "temperature": 1.0, "seed": 0}
res = run_speculative(fam, {**P, "prompt_ids": rand_prompt(12, 1000, seed=gamma)})
base, meth = res["runs"]["baseline"], res["runs"]["method"]
assert meth["ids"] == base["ids"]
assert res["summary"]["identical"] is True
s = res["summary"]
assert s["drafted"] > 0 and s["target_forwards"] <= len(meth["ids"])
assert len(meth["kinds"]) == len(meth["ids"]) == len(meth["tok_iter"])
def test_speculative_sampling_runs_and_counts(tiny_family):
fam = tiny_family("qwen3", scale=0.05)
P = {"max_new_tokens": 30, "gamma": 4, "mode": "sampling", "temperature": 0.8, "seed": 7}
res = run_speculative(fam, {**P, "prompt_ids": rand_prompt(10, 1000)})
s = res["summary"]
assert len(res["runs"]["method"]["ids"]) <= 30
assert 0.0 <= s["alpha_hat"] <= 1.0
assert s["accepted"] <= s["drafted"]
assert "identical" not in s # identity only holds for greedy
def test_speculative_sampling_preserves_target_distribution():
"""Leviathan et al.: the first emitted token is distributed exactly as p."""
p = torch.tensor([0.05, 0.40, 0.10, 0.25, 0.15, 0.05])
q = torch.tensor([0.30, 0.10, 0.20, 0.10, 0.10, 0.20])
ps = torch.stack([p, p])
gen = make_generator(0)
counts = torch.zeros(6)
trials = 40_000
for _ in range(trials):
d = sample(q, gen)
n, nxt, _, _ = verify(ps, [q], [d], greedy=False, gen=gen)
counts[d if n >= 1 else nxt] += 1
assert torch.allclose(counts / trials, p, atol=0.01)
def test_identical_draft_is_always_accepted():
p = torch.tensor([0.1, 0.6, 0.3])
gen = make_generator(1)
for _ in range(200):
d = sample(p, gen)
n, _, kind, probs = verify(torch.stack([p, p]), [p], [d], greedy=False, gen=gen)
assert n == 1 and kind == "bonus" and probs == [1.0]
@pytest.mark.parametrize("arch", ARCHS)
def test_prompt_lookup_equals_autoregressive(tiny_family, arch):
fam = tiny_family(arch)
motif = rand_prompt(6, 1000, seed=5)
prompt = motif * 3 + rand_prompt(4, 1000, seed=6) + motif[:3]
P = {"max_new_tokens": 40, "num_pred": 6, "ngram_max": 3, "ngram_min": 1}
res = run_pld(fam, {**P, "prompt_ids": prompt})
assert res["runs"]["method"]["ids"] == res["runs"]["baseline"]["ids"]
assert res["summary"]["draft_hits"] > 0
@pytest.mark.parametrize("arch", ARCHS)
@pytest.mark.parametrize("block", [1, 4, 8])
def test_jacobi_equals_autoregressive(tiny_family, arch, block):
fam = tiny_family(arch)
P = {"max_new_tokens": 30, "block": block, "init": "repeat", "seed": 0}
res = run_jacobi(fam, {**P, "prompt_ids": rand_prompt(10, 1000, seed=block)})
meth = res["runs"]["method"]
assert meth["ids"] == res["runs"]["baseline"]["ids"]
assert all(it["jacobi"]["fixed"] >= 1 for it in meth["iters"])
assert res["summary"]["iterations"] <= len(meth["ids"])
def test_jacobi_random_init_equals_autoregressive(tiny_family):
fam = tiny_family("gemma3")
P = {"max_new_tokens": 25, "block": 6, "init": "random", "seed": 3}
res = run_jacobi(fam, {**P, "prompt_ids": rand_prompt(9, 1000)})
assert res["runs"]["method"]["ids"] == res["runs"]["baseline"]["ids"]
class Biased(torch.nn.Module):
"""Wraps a model so that one token always wins: drafts and Jacobi guesses then converge."""
def __init__(self, model, token: int):
super().__init__()
self.inner, self.token = model, token
self.config, self.generation_config = model.config, model.generation_config
@property
def device(self):
return self.inner.device
def get_output_embeddings(self):
return self.inner.get_output_embeddings()
def forward(self, **kw):
out = self.inner(**kw)
out.logits[..., self.token] += 1e4
return out
def test_multi_token_commits(tiny_family):
fam = tiny_family("llama")
fam.main = Biased(fam.main, 7)
prompt = rand_prompt(8, 1000) + [7] * 8
jac = run_jacobi(fam, {"max_new_tokens": 20, "block": 5, "init": "repeat", "seed": 0, "prompt_ids": prompt})
assert jac["runs"]["method"]["ids"] == [7] * 20 == jac["runs"]["baseline"]["ids"]
assert jac["runs"]["method"]["iters"][0]["jacobi"]["fixed"] == 6 # the whole block plus one bonus token
pld = run_pld(fam, {"max_new_tokens": 20, "num_pred": 5, "ngram_max": 3, "ngram_min": 1, "prompt_ids": prompt})
assert pld["runs"]["method"]["ids"] == [7] * 20
assert pld["summary"]["tau"] > 3
def test_ngram_lookup_prefers_longest_recent_match():
S = [1, 2, 3, 9, 1, 2, 3, 8, 7, 1, 2, 3]
cont, (start, n) = ngram_lookup(S, 3, 1, 2)
assert n == 3 and cont == [8, 7] and start == 7
assert ngram_lookup([1, 2, 3], 3, 1, 2) == ([], None)
def test_predicted_speedup_limits():
assert predicted_speedup(1.0, 4, 0.5) == pytest.approx(5 / 3)
assert predicted_speedup(0.0, 4, 0.5) == pytest.approx(1 / 3)
assert predicted_speedup(0.7, 0, 0.5) == 1.0
assert expected_tokens(0.5, 3) == pytest.approx(1.875)
|