File size: 4,300 Bytes
1e251f3 | 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 | """WIN2 greedy refinement (Fable review): distribution check at the first greedy divergence.
A near-tie only says the two competing tokens were close. This adds the question "did the whole next-token
distribution change?": in a control/candidate/control window the candidate arm's greedy check records the top-20 and
the prefix at each divergence (nbw_cells.greedy(record_distribution=True)); the following control arm re-probes the
same prefixes (fresh cache salt) and top20_kl measures KL(control || candidate) over the shared top-20 tokens plus one
OTHER bucket, the same coarse estimator UP2's screen uses (up2_checks.numerical). The bound is a registration choice
of each runner; there is no default.
"""
from __future__ import annotations
import math
from collections.abc import Mapping, Sequence
from typing import cast
from offload.gb10_budget import window_cells as wc
from offload.launch.windows import nbw_cells as nbw
from offload.launch.windows import up2_checks as ck
from offload.prefill_maxlen.evidence import Json, number, record
MIN_COMMON_MASS = .99
def _probabilities(top: Mapping[str, Json]) -> dict[int, float]:
ck.require(len(top) == 20, 'exactly twenty top-logprob tokens required')
out: dict[int, float] = {}
for key, value in top.items():
ck.require(key.startswith('token_id:') and key[9:].isdigit(), 'malformed top-20 token id')
lp = number(value)
ck.require(-9999 < lp <= 0, 'clamped or invalid logprob')
out[int(key[9:])] = math.exp(lp)
ck.require(sum(out.values()) <= 1 + 1e-6, 'top-20 probabilities exceed one')
return out
def top20_kl(reference: Mapping[str, Json], experiment: Mapping[str, Json]) -> dict[str, float]:
"""KL(reference || experiment) over shared top-20 tokens + OTHER; common_mass = smaller shared probability mass."""
p, q = _probabilities(reference), _probabilities(experiment)
keys = sorted(p.keys() & q.keys())
px, qx = [p[k] for k in keys], [q[k] for k in keys]
common = min(sum(px), sum(qx))
px.append(max(0.0, 1 - sum(px)))
qx.append(max(0.0, 1 - sum(qx)))
kl = sum(v * math.log(v / w) if w > 0 else math.inf for v, w in zip(px, qx) if v > 0)
ck.require(math.isfinite(kl), 'nonfinite coarse KL at the divergence')
return {'kl_nats': kl, 'common_mass': common, 'shared_tokens': float(len(keys))}
def passes(result: Mapping[str, float], bound: float) -> bool:
if not (bound > 0 and math.isfinite(bound)):
raise ValueError('a positive registered KL bound is required')
return result['common_mass'] >= MIN_COMMON_MASS and result['kl_nats'] <= bound
def probe(api: wc.Http, prefix: Sequence[int], salt: str) -> dict[str, Json]:
"""Top-20 at an exact prefix from the server under test, with a fresh cache salt."""
ck.require(bool(salt.strip()), 'probe cache salt required')
r = api.post('/v1/completions', {'model': nbw.MODEL, 'prompt': cast(Json, list(prefix)), 'max_tokens': 1,
'temperature': 0, 'logprobs': 20, 'return_tokens_as_token_ids': True,
'priority': 10, 'cache_salt': salt})
tops = record(record(ck.rows(r['choices'])[0])['logprobs'])['top_logprobs']
ck.require(isinstance(tops, list) and len(tops) == 1, 'one top-logprob position expected')
return record(cast(list[Json], tops)[0])
def compare_recorded(api: wc.Http, candidate_rows: Sequence[Mapping[str, Json]], salt_prefix: str,
bound: float) -> list[dict[str, Json]]:
"""Re-probe every recorded candidate divergence on the CURRENT (control) server and measure the KL."""
out: list[dict[str, Json]] = []
for row in candidate_rows:
if row.get('identical') is True:
continue
ck.require('top20_logprobs' in row and 'probe_prefix_token_ids' in row,
'candidate divergence lacks a recorded distribution (record_distribution=True)')
prefix = nbw.token_ids(row['probe_prefix_token_ids'])
control = probe(api, prefix, f'{salt_prefix}:kl:{row.get("i")}')
result = top20_kl(control, record(row['top20_logprobs']))
out.append({'i': row.get('i'), **{k: v for k, v in result.items()}, 'bound': bound,
'pass': passes(result, bound)})
return out
|