347 lines
15 KiB
Python
347 lines
15 KiB
Python
"""LLM clients: HFTransformersClient (transformers/AWQ, T4, greedy) plus
|
|
NullClient/CallableClient for dev and an optional VLLMClient."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import re
|
|
from pathlib import Path
|
|
from typing import Callable, List, Optional, Sequence, Tuple
|
|
|
|
MAX_NEW_TOKENS = 1024
|
|
DEFAULT_MODEL_DIR = "." # submission: weights ship at the repo root
|
|
|
|
|
|
class LLMClient:
|
|
"""Interface: generate(prompts) -> one completion per prompt."""
|
|
|
|
available: bool = False
|
|
# can this backend score candidate next-tokens (needed for the match_letters
|
|
# assignment solver)? Only the real transformers backend can.
|
|
can_score: bool = False
|
|
deadline: Optional[float] = None # monotonic wall-clock abort (armed by caller)
|
|
|
|
def generate(self, prompts: Sequence[str], max_new_tokens: int = MAX_NEW_TOKENS,
|
|
system: Optional[str] = None, sample: bool = False,
|
|
temperature: float = 0.5) -> List[str]:
|
|
raise NotImplementedError
|
|
|
|
|
|
class NullClient(LLMClient):
|
|
available = False
|
|
|
|
def generate(self, prompts, max_new_tokens=MAX_NEW_TOKENS, system=None,
|
|
sample=False, temperature=0.5):
|
|
return ["" for _ in prompts]
|
|
|
|
|
|
class CallableClient(LLMClient):
|
|
available = True
|
|
|
|
def __init__(self, fn: Callable[[str], str]):
|
|
self.fn = fn
|
|
|
|
def generate(self, prompts, max_new_tokens=MAX_NEW_TOKENS, system=None,
|
|
sample=False, temperature=0.5):
|
|
return [self.fn(p) for p in prompts]
|
|
|
|
|
|
def pack_by_tokens(lengths: Sequence[int], token_budget: int, max_batch: int
|
|
) -> List[List[int]]:
|
|
"""Group prompt indices into batches whose TOTAL token count stays under
|
|
`token_budget`. On the sandbox's transformers version, prefill computes
|
|
float32 logits over every prompt position (batch x seq x 152k vocab), so
|
|
total-tokens-per-batch — not batch count — is what bounds T4 memory.
|
|
An oversized single prompt still gets its own batch (handled by the OOM
|
|
retry path)."""
|
|
batches: List[List[int]] = []
|
|
cur: List[int] = []
|
|
cur_tokens = 0
|
|
for i, n in enumerate(lengths):
|
|
if cur and (cur_tokens + n > token_budget or len(cur) >= max_batch):
|
|
batches.append(cur)
|
|
cur, cur_tokens = [], 0
|
|
cur.append(i)
|
|
cur_tokens += n
|
|
if cur:
|
|
batches.append(cur)
|
|
return batches
|
|
|
|
|
|
class HFTransformersClient(LLMClient):
|
|
"""transformers backend, mirroring the official notebook's loading code
|
|
(fp16, device_map=auto, greedy) plus token-budget batching and OOM
|
|
recovery. Import is lazy."""
|
|
|
|
available = True
|
|
can_score = True # supports score_next_logprobs (match_letters assignment)
|
|
|
|
# total prompt tokens per generation batch. On the sandbox's transformers,
|
|
# prefill computes fp32 logits over every prompt position (batch x seq x
|
|
# 152k vocab), ~0.9 MB/token; total-tokens-per-batch — not batch count — is
|
|
# the T4 memory bound. 3500 is the safe fallback used when we cannot
|
|
# measure free VRAM; __init__ raises it to fit whatever headroom the loaded
|
|
# model actually leaves (≈6500 for a 7B-AWQ, ≈3500 for a 14B-AWQ).
|
|
TOKEN_BUDGET = 3500
|
|
MB_PER_TOKEN = 0.9 # fp32 prefill logits at Qwen's 152k vocab
|
|
VRAM_RESERVE_GB = 1.8 # KV cache + activations + fragmentation slack
|
|
|
|
def __init__(self, model_dir: str = DEFAULT_MODEL_DIR, batch_size: int = 8):
|
|
import inspect
|
|
import time as _time
|
|
|
|
import torch
|
|
from transformers import (AutoModelForCausalLM, AutoTokenizer,
|
|
StoppingCriteria, StoppingCriteriaList)
|
|
|
|
self.torch = torch
|
|
self.tok = AutoTokenizer.from_pretrained(model_dir)
|
|
# Pin to GPU 0: device_map="auto" can silently offload layers to CPU on a
|
|
# tight T4 (~100x slower). Fall back to "auto" if the pinned load fails.
|
|
try:
|
|
self.model = AutoModelForCausalLM.from_pretrained(
|
|
model_dir, torch_dtype=torch.float16,
|
|
device_map={"": 0} if torch.cuda.is_available() else "auto",
|
|
).eval()
|
|
except Exception:
|
|
self.model = AutoModelForCausalLM.from_pretrained(
|
|
model_dir, torch_dtype=torch.float16, device_map="auto",
|
|
).eval()
|
|
if self.tok.pad_token_id is None:
|
|
self.tok.pad_token = self.tok.eos_token
|
|
self.batch_size = batch_size
|
|
# newer transformers can skip full-sequence prefill logits entirely
|
|
fwd_params = inspect.signature(self.model.forward).parameters
|
|
self._logits_kwarg = next(
|
|
(k for k in ("logits_to_keep", "num_logits_to_keep") if k in fwd_params),
|
|
None)
|
|
self._tune_token_budget()
|
|
self.last_truncated: List[bool] = []
|
|
|
|
# per-token wall-clock abort: a batch started near the deadline can't
|
|
# overrun and get the process killed (the budget is otherwise only
|
|
# checked between batches). set self.deadline (monotonic ts) to arm it.
|
|
self.deadline: Optional[float] = None
|
|
|
|
class _Deadline(StoppingCriteria):
|
|
def __call__(self, input_ids, scores, **kw):
|
|
return _time.monotonic() > deadline_holder[0]
|
|
|
|
deadline_holder = [float("inf")]
|
|
self._deadline_holder = deadline_holder
|
|
self._deadline_crit = StoppingCriteriaList([_Deadline()])
|
|
|
|
def _tune_token_budget(self) -> None:
|
|
"""Size the per-batch token budget to the VRAM the loaded weights
|
|
actually leave free. Falls back to the safe class default if the GPU
|
|
can't be queried (CPU dev, older CUDA)."""
|
|
torch = self.torch
|
|
try:
|
|
if not torch.cuda.is_available():
|
|
return
|
|
free_bytes, _ = torch.cuda.mem_get_info()
|
|
free_gb = free_bytes / 1e9
|
|
budget = int((free_gb - self.VRAM_RESERVE_GB) * 1000 / self.MB_PER_TOKEN)
|
|
# the logits_to_keep fast path removes the big allocation entirely,
|
|
# but stay conservative regardless; clamp to a sane window
|
|
self.TOKEN_BUDGET = max(2500, min(7000, budget))
|
|
except Exception:
|
|
pass
|
|
|
|
def _chat_texts(self, prompts: Sequence[str], system: Optional[str]) -> List[str]:
|
|
texts = []
|
|
for p in prompts:
|
|
messages = ([{"role": "system", "content": system}] if system else []) \
|
|
+ [{"role": "user", "content": p}]
|
|
texts.append(self.tok.apply_chat_template(
|
|
messages, add_generation_prompt=True, tokenize=False))
|
|
return texts
|
|
|
|
def _generate_chunk(self, chunk: List[str], max_new_tokens: int,
|
|
sample: bool = False, temperature: float = 0.5
|
|
) -> Tuple[List[str], List[bool]]:
|
|
torch = self.torch
|
|
# pad only when batching more than one prompt (padding can perturb a
|
|
# quantized model's greedy outputs)
|
|
enc = self.tok(chunk, return_tensors="pt", padding=len(chunk) > 1,
|
|
add_special_tokens=False).to(self.model.device)
|
|
kwargs = {self._logits_kwarg: 1} if self._logits_kwarg else {}
|
|
# repetition_penalty=1.0 explicitly: the shipped generation_config sets
|
|
# 1.05, which is applied even under greedy and biases against the
|
|
# repeated characters common in these answers.
|
|
kwargs["repetition_penalty"] = 1.0
|
|
if sample:
|
|
kwargs.update(do_sample=True, temperature=temperature, top_p=0.95)
|
|
else:
|
|
kwargs["do_sample"] = False
|
|
if self.deadline is not None:
|
|
self._deadline_holder[0] = self.deadline
|
|
kwargs["stopping_criteria"] = self._deadline_crit
|
|
with torch.no_grad():
|
|
gen = self.model.generate(
|
|
**enc, max_new_tokens=max_new_tokens,
|
|
pad_token_id=self.tok.pad_token_id, **kwargs,
|
|
)
|
|
new_tokens = gen[:, enc["input_ids"].shape[1]:]
|
|
eos = self.tok.eos_token_id
|
|
texts, truncated = [], []
|
|
for row in new_tokens:
|
|
texts.append(_guard(self.tok.decode(row, skip_special_tokens=True)).strip())
|
|
# no EOS in the generated span => generation was cut at the cap
|
|
truncated.append(eos is None or int((row == eos).sum()) == 0)
|
|
return texts, truncated
|
|
|
|
def generate(self, prompts, max_new_tokens=MAX_NEW_TOKENS, system=None,
|
|
sample=False, temperature=0.5):
|
|
if not prompts:
|
|
self.last_truncated = []
|
|
return []
|
|
torch = self.torch
|
|
texts = self._chat_texts(prompts, system)
|
|
lengths = [len(self.tok(t, add_special_tokens=False)["input_ids"])
|
|
for t in texts]
|
|
out: List[str] = [""] * len(texts)
|
|
trunc: List[bool] = [False] * len(texts)
|
|
prev_side = self.tok.padding_side
|
|
self.tok.padding_side = "left" # sequences must end at the gen position
|
|
try:
|
|
for batch in pack_by_tokens(lengths, self.TOKEN_BUDGET, self.batch_size):
|
|
chunk = [texts[i] for i in batch]
|
|
try:
|
|
results, tflags = self._generate_chunk(
|
|
chunk, max_new_tokens, sample, temperature)
|
|
except torch.cuda.OutOfMemoryError:
|
|
# halve pressure: clear cache, retry one prompt at a time;
|
|
# a prompt that OOMs alone yields "" (symbolic answer stands)
|
|
torch.cuda.empty_cache()
|
|
results, tflags = [], []
|
|
for t in chunk:
|
|
try:
|
|
r1, f1 = self._generate_chunk(
|
|
[t], max_new_tokens, sample, temperature)
|
|
results.append(r1[0])
|
|
tflags.append(f1[0])
|
|
except torch.cuda.OutOfMemoryError:
|
|
torch.cuda.empty_cache()
|
|
results.append("")
|
|
tflags.append(False)
|
|
for i, r, f in zip(batch, results, tflags):
|
|
out[i] = r
|
|
trunc[i] = f
|
|
finally:
|
|
self.tok.padding_side = prev_side
|
|
self.last_truncated = trunc
|
|
return out
|
|
|
|
def score_next_logprobs(self, prompts, cand_token_ids, system=None,
|
|
batch_size=4):
|
|
"""For each prompt, return the max next-token log-prob over each
|
|
candidate group. `cand_token_ids` is a list of token-id groups (one per
|
|
option), shared across prompts. Returns List[List[float]]
|
|
(prompt x option). One forward pass per batch; used by the match_letters
|
|
assignment solver. Missing/OOM/timed-out prompts get all-zero rows so
|
|
the caller can fall back."""
|
|
import time
|
|
if not prompts:
|
|
return []
|
|
torch = self.torch
|
|
n_opt = len(cand_token_ids)
|
|
texts = self._chat_texts(prompts, system)
|
|
out: List[List[float]] = []
|
|
prev_side = self.tok.padding_side
|
|
self.tok.padding_side = "left"
|
|
bs = batch_size
|
|
i = 0
|
|
try:
|
|
while i < len(texts):
|
|
if self.deadline is not None and time.monotonic() > self.deadline:
|
|
out.extend([[0.0] * n_opt for _ in range(len(texts) - i)])
|
|
break
|
|
chunk = texts[i:i + bs]
|
|
try:
|
|
enc = self.tok(chunk, return_tensors="pt", padding=True,
|
|
add_special_tokens=False, truncation=True,
|
|
max_length=6144).to(self.model.device)
|
|
fwd = {self._logits_kwarg: 1} if self._logits_kwarg else {}
|
|
with torch.no_grad():
|
|
logits = self.model(**enc, **fwd).logits[:, -1, :].float()
|
|
logprobs = torch.log_softmax(logits, dim=-1)
|
|
for b in range(len(chunk)):
|
|
row = [max((logprobs[b, t].item() for t in group),
|
|
default=-1e9) if group else -1e9
|
|
for group in cand_token_ids]
|
|
out.append(row)
|
|
i += bs
|
|
except torch.cuda.OutOfMemoryError:
|
|
torch.cuda.empty_cache()
|
|
if bs == 1:
|
|
out.append([0.0] * n_opt)
|
|
i += 1
|
|
else:
|
|
bs = max(1, bs // 2)
|
|
finally:
|
|
self.tok.padding_side = prev_side
|
|
return out
|
|
|
|
|
|
class VLLMClient(LLMClient):
|
|
"""Optional vLLM backend for throughput experiments. Never required."""
|
|
|
|
available = True
|
|
|
|
def __init__(self, model_dir: str = DEFAULT_MODEL_DIR, max_model_len: int = 4096):
|
|
from vllm import LLM
|
|
|
|
self.llm = LLM(model=model_dir, dtype="half", max_model_len=max_model_len,
|
|
gpu_memory_utilization=0.90)
|
|
|
|
def generate(self, prompts, max_new_tokens=MAX_NEW_TOKENS, system=None):
|
|
from vllm import SamplingParams
|
|
|
|
tok = self.llm.get_tokenizer()
|
|
texts = []
|
|
for p in prompts:
|
|
messages = ([{"role": "system", "content": system}] if system else []) \
|
|
+ [{"role": "user", "content": p}]
|
|
texts.append(tok.apply_chat_template(messages, add_generation_prompt=True,
|
|
tokenize=False))
|
|
params = SamplingParams(temperature=0.0, max_tokens=max_new_tokens)
|
|
outs = self.llm.generate(texts, params)
|
|
return [_guard(o.outputs[0].text if o.outputs else "").strip() for o in outs]
|
|
|
|
|
|
def _guard(text: str, max_repeat: int = 4) -> str:
|
|
"""Loop-collapse guard: truncate at the point where a line repeats more
|
|
than `max_repeat` times consecutively."""
|
|
lines = text.splitlines()
|
|
out, streak = [], 0
|
|
for i, l in enumerate(lines):
|
|
if i > 0 and l.strip() and l == lines[i - 1]:
|
|
streak += 1
|
|
if streak >= max_repeat:
|
|
break
|
|
else:
|
|
streak = 0
|
|
out.append(l)
|
|
return "\n".join(out)
|
|
|
|
|
|
def load_client(model_id: Optional[str] = None) -> LLMClient:
|
|
"""Best available client for `model_id` (script.py's MODEL_ID):
|
|
- a local path ("." in the submission, weights/base in dev) is loaded
|
|
when its config.json exists;
|
|
- a Hub name (contains "/" and is not a local dir) is passed straight to
|
|
transformers — the Colab-testing path, mirroring the workshop notebook;
|
|
- anything unloadable degrades to NullClient (symbolic-only pipeline)."""
|
|
if model_id and "/" in model_id and not Path(model_id).exists():
|
|
try:
|
|
return HFTransformersClient(model_dir=model_id)
|
|
except Exception:
|
|
return NullClient()
|
|
for d in ([model_id] if model_id else []) + [DEFAULT_MODEL_DIR, "weights/base"]:
|
|
if d and Path(d, "config.json").exists():
|
|
try:
|
|
return HFTransformersClient(model_dir=d)
|
|
except Exception:
|
|
continue
|
|
return NullClient()
|