File size: 10,388 Bytes
5407bdc | 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 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 | """Decoding strategies implemented on top of ``model.forward()``.
``model.generate()`` is never called. Each strategy drives the forward pass
itself and maintains its own KV cache, which is what keeps the cost linear in
the number of generated tokens instead of quadratic: without a cache every new
token would re-encode the entire prefix.
Prompts in a batch have different lengths, so they are **left**-padded. That
puts every real last token at index -1, which is what the decoder needs, but
it also means the absolute positions are wrong unless ``position_ids`` is
derived from the attention mask rather than from ``arange`` -- which is what
``_initial_forward`` does below.
"""
from __future__ import annotations
import torch
import torch.nn.functional as F
# ---------------------------------------------------------------------------
# Shared plumbing
# ---------------------------------------------------------------------------
def _initial_forward(model, input_ids, attention_mask):
"""Encode the prompt once and return (last-step logits, cache, position)."""
pos = attention_mask.long().cumsum(-1) - 1
pos.clamp_(min=0)
out = model(input_ids=input_ids, attention_mask=attention_mask,
position_ids=pos, use_cache=True)
return out.logits[:, -1, :].float(), out.past_key_values, pos[:, -1]
def _step(model, token, attention_mask, past, pos):
"""One incremental decoding step through the cache."""
out = model(input_ids=token.unsqueeze(-1), attention_mask=attention_mask,
position_ids=pos.unsqueeze(-1), past_key_values=past,
use_cache=True)
return out.logits[:, -1, :].float(), out.past_key_values
def _grow_mask(attention_mask):
ones = attention_mask.new_ones((attention_mask.size(0), 1))
return torch.cat([attention_mask, ones], dim=1)
def _finish(seqs, eos_id, prompt_len):
"""Trim each generated sequence at its first EOS."""
out = []
for s in seqs:
if eos_id is not None and eos_id in s:
s = s[:s.index(eos_id)]
out.append(s)
return out
# ---------------------------------------------------------------------------
# 1. Greedy
# ---------------------------------------------------------------------------
@torch.no_grad()
def greedy_decode(model, input_ids, attention_mask, max_new=64, eos_id=None):
"""Always take the arg-max. Deterministic, and the width-1 limit of beam
search *only* when no length normalisation is applied."""
B = input_ids.size(0)
logits, past, pos = _initial_forward(model, input_ids, attention_mask)
gen = [[] for _ in range(B)]
done = torch.zeros(B, dtype=torch.bool, device=input_ids.device)
scores = torch.zeros(B, device=input_ids.device)
for _ in range(max_new):
lp = F.log_softmax(logits, dim=-1)
nxt = lp.argmax(-1)
scores += lp.gather(1, nxt.unsqueeze(1)).squeeze(1) * (~done)
for i, t in enumerate(nxt.tolist()):
if not done[i]:
gen[i].append(t)
if eos_id is not None:
done |= nxt == eos_id
if bool(done.all()):
break
attention_mask = _grow_mask(attention_mask)
pos = pos + 1
logits, past = _step(model, nxt, attention_mask, past, pos)
return _finish(gen, eos_id, input_ids.size(1)), scores.tolist()
# ---------------------------------------------------------------------------
# 2 & 3. Truncated sampling: top-k and top-p
# ---------------------------------------------------------------------------
def _top_k_filter(logits, k):
"""Keep the k highest logits, mask the rest to -inf."""
if k is None or k <= 0 or k >= logits.size(-1):
return logits
kth = logits.topk(k, dim=-1).values[..., -1, None]
return logits.masked_fill(logits < kth, float("-inf"))
def _top_p_filter(logits, p):
"""Nucleus filtering: keep the smallest prefix of the sorted distribution
whose cumulative probability first reaches p.
The shift below is what guarantees the top-1 token always survives: a
token is removed only if the mass *strictly before* it already reached p,
so even a single token with probability > p is kept.
"""
if p is None or p >= 1.0:
return logits
srt, idx = logits.sort(descending=True, dim=-1)
cum = srt.softmax(-1).cumsum(-1)
remove = cum - srt.softmax(-1) >= p # mass strictly before token
srt = srt.masked_fill(remove, float("-inf"))
return srt.scatter(-1, idx, srt)
@torch.no_grad()
def sample_decode(model, input_ids, attention_mask, max_new=64, eos_id=None,
top_k=None, top_p=None, temperature=1.0, generator=None):
"""Ancestral sampling from a truncated, temperature-scaled distribution."""
B = input_ids.size(0)
logits, past, pos = _initial_forward(model, input_ids, attention_mask)
gen = [[] for _ in range(B)]
done = torch.zeros(B, dtype=torch.bool, device=input_ids.device)
scores = torch.zeros(B, device=input_ids.device)
for _ in range(max_new):
full_lp = F.log_softmax(logits, dim=-1) # score under the *model*
z = logits / max(temperature, 1e-6)
z = _top_k_filter(z, top_k)
z = _top_p_filter(z, top_p)
probs = z.softmax(-1)
nxt = torch.multinomial(probs, 1, generator=generator).squeeze(1)
scores += full_lp.gather(1, nxt.unsqueeze(1)).squeeze(1) * (~done)
for i, t in enumerate(nxt.tolist()):
if not done[i]:
gen[i].append(t)
if eos_id is not None:
done |= nxt == eos_id
if bool(done.all()):
break
attention_mask = _grow_mask(attention_mask)
pos = pos + 1
logits, past = _step(model, nxt, attention_mask, past, pos)
return _finish(gen, eos_id, input_ids.size(1)), scores.tolist()
# ---------------------------------------------------------------------------
# 4. Beam search
# ---------------------------------------------------------------------------
@torch.no_grad()
def beam_search(model, input_ids, attention_mask, beam_width=4, max_new=64,
eos_id=None, length_penalty=1.0):
"""Keep the ``beam_width`` highest-scoring prefixes at every step.
Cost per step is ``O(W * |V|)`` for the scoring plus a ``top-W`` over the
``W * |V|`` candidates, and the whole batch is run as one ``B*W`` forward,
so wall time grows roughly linearly in W (memory does too -- the KV cache
is W copies of the prompt). Finished hypotheses are moved aside and their
beam slot is neutralised with a -inf mask so it cannot be extended again.
Hypotheses are *selected* by summed log-probability divided by
``len ** length_penalty``; without that normalisation beam search
systematically prefers short outputs, since every extra token adds a
negative number. The score *returned* is the raw summed log-probability,
so that it means the same thing as the score the other strategies return
and the caller can derive a per-token figure the same way for all of them.
"""
B, L = input_ids.shape
dev = input_ids.device
W, V = beam_width, model.config.vocab_size
logits, past, pos = _initial_forward(model, input_ids, attention_mask)
lp = F.log_softmax(logits, dim=-1) # (B, V)
# First expansion: from one prompt to W beams.
top_lp, top_tok = lp.topk(W, dim=-1) # (B, W)
beam_scores = top_lp.reshape(B * W) # (B*W,)
beams = [[int(t)] for t in top_tok.reshape(-1).tolist()]
past.batch_repeat_interleave(W)
attention_mask = attention_mask.repeat_interleave(W, dim=0)
pos = pos.repeat_interleave(W, dim=0)
cur = top_tok.reshape(-1)
# Each entry is (selection score, raw summed logprob, tokens).
finished = [[] for _ in range(B)]
alive = torch.ones(B * W, dtype=torch.bool, device=dev)
for i in range(B * W):
if eos_id is not None and beams[i][0] == eos_id:
finished[i // W].append((beam_scores[i].item(),
beam_scores[i].item(), []))
alive[i] = False
for step in range(1, max_new):
if not bool(alive.any()):
break
attention_mask = _grow_mask(attention_mask)
pos = pos + 1
logits, past = _step(model, cur, attention_mask, past, pos)
lp = F.log_softmax(logits, dim=-1) # (B*W, V)
# A dead beam must not spawn children; -inf everywhere keeps it from
# ever winning a top-W slot.
lp = lp.masked_fill(~alive.unsqueeze(1), float("-inf"))
cand = (beam_scores.unsqueeze(1) + lp).view(B, W * V)
top_sc, top_ix = cand.topk(W, dim=-1) # (B, W)
src_beam = top_ix // V # which beam
tok = top_ix % V # which token
flat_src = (torch.arange(B, device=dev).unsqueeze(1) * W
+ src_beam).reshape(-1)
past.reorder_cache(flat_src)
new_beams = [beams[i] + [int(t)]
for i, t in zip(flat_src.tolist(), tok.reshape(-1).tolist())]
beams = new_beams
beam_scores = top_sc.reshape(-1)
alive = alive[flat_src].clone()
cur = tok.reshape(-1)
if eos_id is not None:
hit = (cur == eos_id) & alive
for i in hit.nonzero(as_tuple=True)[0].tolist():
seq = beams[i][:-1]
norm = len(seq) ** length_penalty if seq else 1.0
raw = beam_scores[i].item()
finished[i // W].append((raw / norm, raw, seq))
alive[i] = False
beam_scores[i] = float("-inf")
# Anything still alive at the cap counts as a (length-normalised) candidate.
for i in range(B * W):
if alive[i]:
seq = beams[i]
norm = len(seq) ** length_penalty if seq else 1.0
raw = beam_scores[i].item()
finished[i // W].append((raw / norm, raw, seq))
out, out_scores = [], []
for b in range(B):
if finished[b]:
_, raw, seq = max(finished[b], key=lambda t: t[0])
else:
raw, seq = 0.0, []
out.append(seq)
out_scores.append(raw) # raw sum, not the selection score
return out, out_scores
|