File size: 19,129 Bytes
09c13a5 | 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 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 | """Letter readout of google/gemma-4-31B-it served by vLLM in-process, with the same prompts and options as jevlab.reader.
Same token ids as Reader: the rows are built by jevlab.reader.Layout (decider or header layout, labels A..J, single-token
wide labels above 10 options), and the read policy (read_twice, twice_max_tokens, reread_below, reverse_below) follows
Reader.__call__ step for step. tools/vllm_check.py checks both on CPU against Reader with a mocked vLLM.
What differs is the readout. Each row is one vLLM request with max_tokens=1 and `logprob_token_ids` set to the labels of
its options (stock vLLM SamplingParams, at most 128 ids per request, so wider questions are read in chunks of 128 with
the same prompt). vLLM returns log-probabilities of the logits after Gemma's final soft cap (30 * tanh(x / 30)), not the
raw logits Reader reads, so Reader's T(n) = 10.124 - 1.633 ln(n) does not carry over. The probabilities are a softmax of
the label log-probabilities at a temperature per question type:
choice: T = temperature_choice + temperature_choice_log_slope * ln(n options)
noul: T = temperature_noul, and noul_bias is added to the "yes" score after dividing by T
The defaults (choice 4.263, noul 7.609, no slope, no bias) are wcole3/jev-brain's published fit for this model in the
header layout on vLLM (MIT, data/gemma-4-31b/header/calibration.json, fitted outside the suite). They are a starting
point; refit them on our own off-suite data. temperature_scheme="options" applies Reader's T(n) instead (for checks).
Why in-process (vllm.LLM with prompt token ids) and not an OpenAI-compatible server: the kit runs one request at a time
in one process, as the board measures it; the engine loads like Reader (`--engine jevlab.vllm_reader:VllmReader`), with
no second process, port or HTTP client to manage on a shared machine, and token ids go to vLLM unchanged (checked on
every output). vLLM's own engine core still runs in a child process; latency is the wall time of the generate calls.
All rows of one read pass of a request go to vLLM in one generate call. With prefix caching, the rows share the KV blocks
of their common prefix (chat head, question list, state); the re-read and reversed passes hit the same cache. Rows above
the token limit (max_model_len - 1) are refused as Unsupported, never truncated. max_model_len defaults to 131,072 (the
board's maintainers lift context caps to that or more for their private runs; the model's window is 262,144): with FP8
weights a request of that length needs about 23 GiB of KV cache, which the default gpu_memory_utilization=0.70 holds.
python -m decision_index run --engine jevlab.vllm_reader:VllmReader --option layout=header \
--rows sample-2000.jsonl.gz --out runs/vllm-header-sample2000
"""
import math
import os
from decision_index.engines.base import Engine, Unsupported
from jevlab.reader import LAYOUTS, Layout, _flag, render_question
REPO, REVISION = "google/gemma-4-31B-it", "842da3794eaa0b77d5f08bae87a17459d91ff475"
MAX_LABEL_IDS = 128 # vLLM's MAX_LOGPROB_TOKEN_IDS (vllm/sampling_params.py)
TEMPERATURE_SCHEMES = ("type", "options")
# Reader options that only mean something for the transformers engine; accepted so the same command lines work.
HF_ONLY_OPTIONS = ("device", "chunk_tokens", "cache_prefix", "cudnn_sdp", "pack_questions", "pack_tokens", "attn_block")
LENGTH_MARKERS = ("maximum model length", "maximum context length", "longer than the maximum")
def softmax(scores):
top = max(scores)
weights = [math.exp(x - top) for x in scores]
total = sum(weights)
return [w / total for w in weights]
class VllmBackend:
"""Log-probabilities of given token ids at the position after each prompt, from one vLLM decoding step.
`llm` and `sampling_params` can be stand-ins with the same interface (tools/vllm_check.py)."""
def __init__(self, llm=None, sampling_params=None, **llm_kwargs):
if llm is None:
from vllm import LLM
llm = LLM(**llm_kwargs)
if sampling_params is None:
from vllm import SamplingParams as sampling_params
self.llm, self.SamplingParams = llm, sampling_params
def read(self, prompts, label_ids):
"""-> (one list of log-probabilities per prompt, in the order of its label_ids; prompt tokens served from the
prefix cache per prompt). Label ids go in chunks of 128; each chunk is sent with the prompts that read one."""
found = [{} for _ in prompts]
cached = [None] * len(prompts)
width = max(map(len, label_ids), default=0)
for start in range(0, width, MAX_LABEL_IDS):
rows = [k for k, ids in enumerate(label_ids) if len(ids) > start]
params = [self.SamplingParams(max_tokens=1, temperature=0.0, logprob_token_ids=label_ids[k][start:start + MAX_LABEL_IDS], detokenize=False)
for k in rows]
outs = self.llm.generate([{"prompt_token_ids": prompts[k]} for k in rows], params, use_tqdm=False)
if len(outs) != len(rows):
raise RuntimeError(f"vLLM returned {len(outs)} outputs for {len(rows)} prompts")
for k, o in zip(rows, outs):
got = o.prompt_token_ids if isinstance(o.prompt_token_ids, list) else list(o.prompt_token_ids)
if got != prompts[k]:
raise RuntimeError("vLLM read other token ids than the ones sent (order or added tokens)")
step = o.outputs[0].logprobs[0]
for t in label_ids[k][start:start + MAX_LABEL_IDS]:
found[k][t] = float(step[t].logprob)
if cached[k] is None:
cached[k] = int(o.num_cached_tokens or 0)
return [[found[k][t] for t in ids] for k, ids in enumerate(label_ids)], cached
def close(self):
engine = getattr(self.llm, "llm_engine", None)
core = getattr(engine, "engine_core", None)
if core is not None and hasattr(core, "shutdown"):
core.shutdown()
class VllmReader(Engine):
name = "jevlab-vllm-reader"
latency = ("In-process request wall time including prompt construction: one vLLM generate call per read pass of a request "
"(all its rows together, shared prefix from vLLM's prefix cache); excludes model loading.")
def __init__(self, model=REPO, revision=REVISION, quantization="fp8_per_block", dtype="bfloat16", max_model_len=131072, max_tokens=None,
gpu_memory_utilization=0.70, kv_cache_dtype="auto", attention_backend="TRITON_ATTN", prefix_caching=True,
prefix_cache_retention_interval=32, max_num_batched_tokens=8192, max_num_seqs=256, enforce_eager=False, language_model_only=True,
deep_gemm=False, seed=0, llm_options=None, layout="decider", read_twice=False, twice_max_tokens=0, reread_below=0.0, reverse_below=0.0,
temperature_scheme="type", temperature_choice=4.263, temperature_choice_log_slope=0.0, temperature_noul=7.609, noul_bias=0.0,
temperature=10.124, temperature_log_slope=-1.633, backend=None, tokenizer=None, system=None, **options):
ignored = {k: options.pop(k) for k in HF_ONLY_OPTIONS if k in options}
if options:
raise ValueError(f"unknown VllmReader options {sorted(options)}")
super().__init__(**ignored)
if layout not in LAYOUTS:
raise ValueError(f"layout must be one of {LAYOUTS}, not {layout!r}")
if temperature_scheme not in TEMPERATURE_SCHEMES:
raise ValueError(f"temperature_scheme must be one of {TEMPERATURE_SCHEMES}, not {temperature_scheme!r}")
if isinstance(quantization, str) and quantization.strip().lower() in ("", "none", "null", "bf16", "bfloat16"):
quantization = None
if isinstance(attention_backend, str) and attention_backend.strip().lower() in ("", "auto", "none", "null"):
attention_backend = None
self.max_model_len = int(max_model_len)
self.limit = int(max_tokens) if max_tokens else self.max_model_len - 1
if not 0 < self.limit < self.max_model_len:
raise ValueError(f"max_tokens must be below max_model_len ({self.max_model_len}): one position is the answer slot")
self.header = layout == "header"
self.read_twice = _flag(read_twice)
self.twice_max_tokens, self.reread_below = int(twice_max_tokens), float(reread_below)
self.reverse_below = float(reverse_below)
self.scheme = temperature_scheme
self.t_choice, self.t_choice_slope = float(temperature_choice), float(temperature_choice_log_slope)
self.t_noul, self.noul_bias = float(temperature_noul), float(noul_bias)
self.temperature, self.temperature_log_slope = float(temperature), float(temperature_log_slope)
self.model_id = model
if tokenizer is None:
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained(model, revision=revision)
self.tok = tokenizer
self.system = system or None
self.layout = Layout(tokenizer, system=self.system)
self.llm_kwargs = dict(model=model, revision=revision, tokenizer_revision=revision, dtype=dtype, quantization=quantization,
max_model_len=self.max_model_len, gpu_memory_utilization=float(gpu_memory_utilization), kv_cache_dtype=kv_cache_dtype,
enable_prefix_caching=_flag(prefix_caching), prefix_cache_retention_interval=prefix_cache_retention_interval,
max_num_batched_tokens=int(max_num_batched_tokens), max_num_seqs=int(max_num_seqs), enforce_eager=_flag(enforce_eager),
language_model_only=_flag(language_model_only), attention_backend=attention_backend, seed=int(seed),
logprobs_mode="raw_logprobs", generation_config="vllm", **dict(llm_options or {}))
if backend is None and not _flag(deep_gemm):
# block-FP8 on sm_120 picks DeepGEMM, which fails there; off, vLLM uses its CUTLASS block-FP8 kernel (vLLM #60261, #51884)
os.environ.setdefault("VLLM_USE_DEEP_GEMM", "0")
self.backend = backend if backend is not None else VllmBackend(**self.llm_kwargs)
temps = ({"choice": f"{self.t_choice} + {self.t_choice_slope} * ln(n options)", "noul": self.t_noul, "noul_bias": self.noul_bias}
if self.scheme == "type" else {"all": f"max({self.temperature} + {self.temperature_log_slope} * ln(n options), 0.05)"})
self.provenance = {
"kind": "inference technique",
"repo": model,
"revision": revision,
"serving": "vLLM offline LLM, in-process, prompt token ids",
"vllm_options": {k: v for k, v in self.llm_kwargs.items() if k not in ("model", "revision", "tokenizer_revision")},
"VLLM_USE_DEEP_GEMM": os.environ.get("VLLM_USE_DEEP_GEMM"),
"max_row_tokens": self.limit,
"read_twice": self.read_twice,
"twice_max_tokens": self.twice_max_tokens,
"reread_below": self.reread_below,
"reverse_below": self.reverse_below,
"layout": layout,
**({"system": self.system} if self.system else {}),
"temperature_scheme": self.scheme,
"temperatures": temps,
"ignored_options": sorted(ignored),
"policy": ("The state is preceded by the list of every question of the request ('Questions that will be asked about the context:'), "
"the header layout of wcole3/jev-brain. " if self.header else "")
+ (f"A fixed system turn comes first: {self.system!r}. " if self.system else "")
+ "One row per question in the model's chat template with thinking off, ending at 'Answer: ('; vLLM reads one step and returns "
"the log-probabilities (after the final soft cap) of the option labels (logprob_token_ids), which go through a softmax at a "
+ ("temperature per question type. " if self.scheme == "type" else "temperature set by the number of options. ")
+ "Options keep the request's order. Rows above the token limit are refused, never cut."
+ (" The user turn (state, question, options) is written twice, joined by 'Read the same input again before answering:'"
+ (f", in requests of at most {self.twice_max_tokens} tokens read once" if self.twice_max_tokens > 0 else "") + "." if self.read_twice else "")
+ (f" Questions of the other requests whose most likely option has probability below {self.reread_below} are read again, written twice "
"the same way, after the first read." if self.reread_below > 0 else "")
+ (f" A question whose most likely option then has probability below {self.reverse_below} is read once more with its options in "
"reverse order, and the two distributions are averaged." if self.reverse_below > 0 else ""),
}
def runtime(self):
info = {}
try:
import torch
import vllm
info.update(vllm=vllm.__version__, torch=torch.__version__, cuda=torch.version.cuda)
from vllm.platforms import current_platform
info["gpu"] = current_platform.get_device_name(0) # NVML, so the parent process opens no CUDA context
except Exception as exc: # the mocked backend of tools/vllm_check.py has no vLLM
info.setdefault("note", f"runtime details unavailable: {type(exc).__name__}")
try:
import transformers
info["transformers"] = transformers.__version__
except ImportError:
pass
return info
def close(self):
if hasattr(self.backend, "close"):
self.backend.close()
def _read(self, ctx, pieces, ns):
"""-> [(label log-probabilities, prompt tokens served from the prefix cache)] for the rows ctx + piece."""
try:
scores, cached = self.backend.read([ctx + piece for piece in pieces], [self.layout.label_ids[:n] for n in ns])
except ValueError as exc:
if any(m in str(exc) for m in LENGTH_MARKERS):
raise Unsupported(f"maximum context length: {exc}") from exc
raise
return list(zip(scores, cached))
def _probs(self, scores, kind, reverse=False):
"""Softmax over the shown options; with reverse the options are in reverse order (yes first for noul)."""
n = len(scores)
if self.scheme == "options":
t, bias = max(self.temperature + self.temperature_log_slope * math.log(n), 0.05), 0.0
elif kind == "noul":
t, bias = self.t_noul, self.noul_bias
else:
t, bias = max(self.t_choice + self.t_choice_slope * math.log(n), 0.05), 0.0
x = [s / t for s in scores]
if bias:
x[0 if reverse else 1] += bias
return softmax(x)
def _check(self, ctx, piece):
if len(ctx) + len(piece) > self.limit:
raise Unsupported(f"maximum context length: a row of {len(ctx) + len(piece)} tokens exceeds the declared limit of {self.limit}")
def __call__(self, state, questions):
# the same steps as jevlab.reader.Reader.__call__, with log-probabilities from vLLM in place of label logits
rendered = [(key, q["type"], *render_question(q)) for key, q in questions.items()]
texts = [text for _, _, text, _, _ in rendered] if self.header else None
ctx = self.layout.context(state, texts)
again = self.layout.again(state, texts) if self.read_twice or self.reread_below > 0 else None
rows = [(key, kind, self.layout.question(text, options), self.layout.question(text, options, again) if again else None, keys)
for key, kind, text, options, keys in rendered]
# a request whose rows written twice would pass the token limit is read once
fits = [again is not None and len(ctx) + len(twice) <= self.limit for _, _, _, twice, _ in rows]
once_tokens = len(ctx) + sum(len(once) for _, _, once, _, _ in rows)
whole = self.read_twice and all(fits) and (self.twice_max_tokens <= 0 or once_tokens <= self.twice_max_tokens)
first = [twice if whole else once for _, _, once, twice, _ in rows]
for piece in first:
self._check(ctx, piece)
ns = [len(keys) for *_, keys in rows]
kinds = [kind for _, kind, *_ in rows]
read = self._read(ctx, first, ns)
pieces = list(first)
reread = self.reread_below > 0 and not whole
unsure = [i for i, (scores, _) in enumerate(read) if reread and fits[i] and max(self._probs(scores, kinds[i])) < self.reread_below]
firsts = {}
if unsure:
for i, again_read in zip(unsure, self._read(ctx, [rows[i][3] for i in unsure], [ns[i] for i in unsure])):
firsts[i] = read[i][0]
read[i] = again_read
pieces[i] = rows[i][3]
probs = [self._probs(scores, kind) for (scores, _), kind in zip(read, kinds)]
flipped = {}
if self.reverse_below > 0:
low = [i for i, pr in enumerate(probs) if max(pr) < self.reverse_below]
if low:
twice = [whole or i in firsts for i in low]
rev = [self.layout.question(rendered[i][2], rendered[i][3][::-1], again if t else None) for i, t in zip(low, twice)]
for piece in rev:
self._check(ctx, piece)
for i, (scores, _) in zip(low, self._read(ctx, rev, [ns[i] for i in low])):
back = self._probs(scores, kinds[i], reverse=True)[::-1] # reversed positions back to the caller's order
flipped[i] = scores
probs[i] = [(a + b) / 2 for a, b in zip(probs[i], back)]
answers, raw = {}, {}
for i, ((key, kind, _, _, keys), (scores, cached)) in enumerate(zip(rows, read)):
raw[key] = {"label_logprobs": [round(x, 6) for x in scores], "prompt_tokens": len(ctx) + len(pieces[i]), "cached_prefix_tokens": cached,
"read_twice": whole or i in firsts}
if i in firsts:
raw[key]["first_label_logprobs"] = [round(x, 6) for x in firsts[i]]
if i in flipped:
raw[key]["reversed_label_logprobs"] = [round(x, 6) for x in flipped[i]]
probs_i = probs[i]
if kind == "noul":
answers[key] = {"type": "noul", "noul": probs_i[1]}
else:
answers[key] = {"type": "choice", "choice": keys[max(range(len(keys)), key=probs_i.__getitem__)], "probabilities": dict(zip(keys, probs_i))}
return {"model": self.model_id, "answers": answers, "usage": {"input_tokens": sum(v["prompt_tokens"] for v in raw.values())}}, raw
|