diff --git a/qev/engine.py b/qev/engine.py index 2e4c7b0..cb79e4f 100644 --- a/qev/engine.py +++ b/qev/engine.py @@ -10,6 +10,7 @@ tolerance of 0.05 and only guards against structural failures (an unknown cache """ from __future__ import annotations +import contextlib import copy import hashlib import logging @@ -19,6 +20,7 @@ from pathlib import Path from typing import Any from .answers import make_choice, make_noul, make_score +from . import readout as ro from .calibrate import Calibration, softmax from .prompt import (DEFAULT_STYLE, LETTERS, MAX_SLOTS, PROMPT_VERSION, SLOT_LABELS, PromptStyle, RenderedQuestion, build_messages, render_question) @@ -102,11 +104,20 @@ class DecisionEngine: self.model_type = config.model_type self.template_kwargs = {"enable_thinking": False} if "qwen" in self.model_type else {} self.slot_ids = self._verify_slots() + rcfg = ro.ReadoutConfig.load(model_path) if local else ro.ReadoutConfig() + self.readout = ro.Readout.attach(self.model, rcfg, slot_ids=self.slot_ids, + head_path=Path(model_path) / ro.HEAD_FILE if rcfg.head else None, + device=self.device) + if rcfg.head: + self.max_slots = min(self.max_slots, rcfg.n_slots) + if calibration is None and local and (Path(model_path) / "calibration.json").exists(): + self.calibration = Calibration.load(Path(model_path) / "calibration.json") self.pad_id = self.tokenizer.pad_token_id if self.tokenizer.pad_token_id is not None else self.tokenizer.eos_token_id self.info = EngineInfo(model_path=model_path, model_type=self.model_type, dtype=dtype, device=str(self.device), fork_mode="sequential", fork_max_abs_diff=None, transformers_version=transformers.__version__, torch_version=torch.__version__, - extra={"head_dtype": head_dtype}) + extra={"head_dtype": head_dtype, "attention": rcfg.attention, "readout": rcfg.readout, + "calibration": self.calibration.source}) if fork_mode == "auto": self._select_fork_mode() else: @@ -190,6 +201,44 @@ class DecisionEngine: with torch.no_grad(): torch.cuda.empty_cache() if self.device.type == "cuda" else None + def _split(self, state: Any, rendered: list[RenderedQuestion]): + """Prefix ids, suffix ids, and per branch the option token spans relative to its suffix (None for slot_lm). + The layout is the one train.common.batch_scores computes, so serving matches training: under noncausal + attention the cache boundary must be the prefix block boundary.""" + prefix_ids, suffixes = self._prefix_and_suffix_ids(state, rendered) + cfg = self.readout.config + if not cfg.needs_layout: + return prefix_ids, suffixes, [None] * len(rendered) + P = len(prefix_ids) + spans = [] + for r, suf in zip(rendered, suffixes): + prompt = self._render_prompt(state, r.suffix) + start = prompt.rindex(r.suffix) + ids, prefix, sp = ro.token_layout(self.tokenizer, prompt, start, + [(start + a, start + b) for a, b in r.option_spans or []]) + if ids != prefix_ids + suf: + raise ValueError("layout tokenization differs from the branch tokenization") + if cfg.noncausal and prefix != P: + raise ValueError(f"prefix block ends at token {prefix} but the shared cache at {P}") + if any(a < P for a, _ in sp): + raise ValueError("an option span starts inside the shared prefix") + spans.append([(a - P, b - P) for a, b in sp]) + return prefix_ids, suffixes, spans + + def _blocks(self, prefixes: list[int], lengths: list[int], width: int): + if not self.readout.config.noncausal: + return contextlib.nullcontext() + return ro.blocks(ro.block_ids(prefixes, lengths, width, self.device)) + + def _scores(self, o, hidden, row: int, logit_pos: int, end: int, n_slots: int, spans) -> tuple[list[float], float]: + """Scores of one branch from a forward's output: the slot letters' LM logits, or the readout head over the + captured final hidden states (`end` and `spans` index the suffix tokens of this row).""" + if self.readout.head is None: + return self._readout(o.logits[row, logit_pos], n_slots) + h = hidden[row: row + 1] + scores = self.readout.head(h[:, end], ro.pool_options(h, [spans], n_slots))[0] + return scores.tolist(), 1.0 + def _prefix_and_suffix_ids(self, state: Any, rendered: list[RenderedQuestion]) -> tuple[list[int], list[list[int]]]: """Tokenize every branch fully, then find the longest common token prefix that ends at or before the end of the state block, so the shared prefix is exact by construction.""" @@ -218,21 +267,25 @@ class DecisionEngine: def _run_sequential(self, state: Any, rendered: list[RenderedQuestion]) -> list[BranchResult]: import torch - prefix_ids, suffixes = self._prefix_and_suffix_ids(state, rendered) + prefix_ids, suffixes, spans = self._split(state, rendered) dev = self.device + P = len(prefix_ids) with torch.inference_mode(): - out = self.model(input_ids=torch.tensor([prefix_ids], device=dev), - attention_mask=torch.ones((1, len(prefix_ids)), dtype=torch.long, device=dev), - use_cache=True, return_dict=True, logits_to_keep=1) + with self._blocks([P], [P], P): + out = self.model(input_ids=torch.tensor([prefix_ids], device=dev), + attention_mask=torch.ones((1, P), dtype=torch.long, device=dev), + use_cache=True, return_dict=True, logits_to_keep=1) base = out.past_key_values del out results = [] - for r, suf in zip(rendered, suffixes): + for r, suf, sp in zip(rendered, suffixes, spans): cache = copy.deepcopy(base) - o = self.model(input_ids=torch.tensor([suf], device=dev), - attention_mask=torch.ones((1, len(prefix_ids) + len(suf)), dtype=torch.long, device=dev), - past_key_values=cache, use_cache=True, return_dict=True, logits_to_keep=1) - logits, mass = self._readout(o.logits[0, -1], r.n_slots) + with self._blocks([P], [P + len(suf)], P + len(suf)): + o = self.model(input_ids=torch.tensor([suf], device=dev), + attention_mask=torch.ones((1, P + len(suf)), dtype=torch.long, device=dev), + past_key_values=cache, use_cache=True, return_dict=True, logits_to_keep=1) + hidden = self.readout.capture.pop() if self.readout.capture else None + logits, mass = self._scores(o, hidden, 0, -1, len(suf) - 1, r.n_slots, sp) results.append(BranchResult(r.qid, r.kind, r.labels, r.descriptions, logits, mass, len(suf))) del o, cache del base @@ -266,22 +319,23 @@ class DecisionEngine: caller falls back to sequential branches.""" import torch - prefix_ids, suffixes = self._prefix_and_suffix_ids(state, rendered) + prefix_ids, suffixes, spans = self._split(state, rendered) dev = self.device P = len(prefix_ids) chunk = max(1, min(len(rendered), self.fork_token_budget // max(P, 1))) results: list[BranchResult | None] = [None] * len(rendered) with torch.inference_mode(): - out = self.model(input_ids=torch.tensor([prefix_ids], device=dev), - attention_mask=torch.ones((1, P), dtype=torch.long, device=dev), - use_cache=True, return_dict=True, logits_to_keep=1) + with self._blocks([P], [P], P): + out = self.model(input_ids=torch.tensor([prefix_ids], device=dev), + attention_mask=torch.ones((1, P), dtype=torch.long, device=dev), + use_cache=True, return_dict=True, logits_to_keep=1) base = out.past_key_values del out start = 0 while start < len(rendered): idx = list(range(start, min(start + chunk, len(rendered)))) try: - self._run_chunk(base, P, idx, rendered, suffixes, results) + self._run_chunk(base, P, idx, rendered, suffixes, results, spans) except torch.OutOfMemoryError: torch.cuda.empty_cache() if chunk == 1: @@ -294,7 +348,7 @@ class DecisionEngine: del base return results - def _run_chunk(self, base, P: int, idx: list[int], rendered, suffixes, results) -> None: + def _run_chunk(self, base, P: int, idx: list[int], rendered, suffixes, results, spans=None) -> None: import torch dev = self.device @@ -311,13 +365,16 @@ class DecisionEngine: ends.append(len(s) - 1) keep = sorted(set(ends)) cache = self._fork(base, len(idx)) + spans = spans or [None] * len(rendered) try: - o = self.model(input_ids=ids, attention_mask=mask, position_ids=pos, past_key_values=cache, - use_cache=True, return_dict=True, - logits_to_keep=torch.tensor(keep, dtype=torch.long, device=dev)) + with self._blocks([P] * len(idx), [P + len(s) for s in sufs], P + width): + o = self.model(input_ids=ids, attention_mask=mask, position_ids=pos, past_key_values=cache, + use_cache=True, return_dict=True, + logits_to_keep=torch.tensor(keep, dtype=torch.long, device=dev)) + hidden = self.readout.capture.pop() if self.readout.capture else None for row, i in enumerate(idx): r = rendered[i] - logits, mass = self._readout(o.logits[row, keep.index(ends[row])], r.n_slots) + logits, mass = self._scores(o, hidden, row, keep.index(ends[row]), ends[row], r.n_slots, spans[i]) results[i] = BranchResult(r.qid, r.kind, r.labels, r.descriptions, logits, mass, len(sufs[row])) del o finally: diff --git a/qev/prompt.py b/qev/prompt.py index 8eaadd3..de49cca 100644 --- a/qev/prompt.py +++ b/qev/prompt.py @@ -74,13 +74,14 @@ class RenderedQuestion: labels: list[str] descriptions: list[Any] suffix: str + option_spans: list[tuple[int, int]] | None = None n_slots: int = field(init=False) def __post_init__(self) -> None: self.n_slots = len(self.labels) -def _options_block(names: Sequence[str], descriptions: Sequence[Any], show_names: bool = True) -> str: +def _options_lines(names: Sequence[str], descriptions: Sequence[Any], show_names: bool = True) -> list[str]: lines = [] for letter, name, desc in zip(SLOT_LABELS, names, descriptions): text = render_entry(desc) @@ -90,7 +91,11 @@ def _options_block(names: Sequence[str], descriptions: Sequence[Any], show_names lines.append(f"{letter}. {name}: {text}") else: lines.append(f"{letter}. {name}") - return "\n".join(lines) + return lines + + +def _options_block(names: Sequence[str], descriptions: Sequence[Any], show_names: bool = True) -> str: + return "\n".join(_options_lines(names, descriptions, show_names)) def render_question(qid: str, q: Question, order: Sequence[int] | None = None, style: PromptStyle = DEFAULT_STYLE) -> RenderedQuestion: @@ -127,6 +132,7 @@ def render_question(qid: str, q: Question, order: Sequence[int] | None = None, s descs = [descs[i] for i in order] letters = list(SLOT_LABELS[: len(labels)]) + spans = None if style.layout == "json": payload = {"criterion": instructions, "options": [{"letter": l, "description": (f"{n}: {render_entry(d)}" if (style.show_names and n is not None and render_entry(d)) else (render_entry(d) or n))} @@ -134,12 +140,17 @@ def render_question(qid: str, q: Question, order: Sequence[int] | None = None, s suffix = json.dumps(payload, ensure_ascii=False)[1:] else: header = style.question_header if kind != "score" else style.score_header + head = f"{header} {instructions}\n\nOptions:\n" + lines = _options_lines(names, descs, style.show_names) + spans, pos = [], len(head) + for line in lines: + spans.append((pos, pos + len(line))) + pos += len(line) + 1 suffix = ( - f"{header} {instructions}\n\n" - f"Options:\n{_options_block(names, descs, style.show_names)}\n\n" + head + "\n".join(lines) + "\n\n" + (style.answer_line if len(letters) <= len(LETTERS) else style.answer_line_many).format(letters=", ".join(letters)) ) - return RenderedQuestion(qid=qid, kind=kind, labels=labels, descriptions=descs, suffix=suffix) + return RenderedQuestion(qid=qid, kind=kind, labels=labels, descriptions=descs, suffix=suffix, option_spans=spans) def state_block(state: Any, style: PromptStyle = DEFAULT_STYLE) -> str: diff --git a/qev/readout.py b/qev/readout.py new file mode 100644 index 0000000..02707ca @@ -0,0 +1,269 @@ +"""Readout variants, shared by training (train.common / train.sft) and serving (qev.engine). + +attention: + causal the base model's mask. + noncausal in the full-attention layers every token attends bidirectionally within its block and causally to + earlier blocks. A rendered prompt has two blocks: the shared prefix (system turn and state block) and + the question suffix (question, options, answer line, generation prompt). The prefix never sees a + suffix, so its cache is still shared exactly by all questions of a request, and one forward of a + suffix scores all of its options. The GatedDeltaNet (linear-attention) layers are recurrent and stay + causal. + +readout: + slot_lm softmax over the LM-head logits of the slot letters at the answer position. + head a learned head over the final (normed) hidden states: + score_i = + / sqrt(d) + w has one row per slot and starts as the LM-head rows of the slot letters, W_k starts at zero, so the + untrained head reproduces slot_lm exactly; the pointer term scores each option by its content. + +Both settings are recorded next to the weights, so a checkpoint is always served the way it was trained. The +attention mask is declared per layer type under `modeling_overrides` in `config.json`, where serving stacks +read it; `readout.json` holds the rest. Checkpoints without the overrides block carry `attention` in `readout.json`. +""" +from __future__ import annotations + +import bisect +import contextlib +import contextvars +import importlib +import json +import math +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Sequence + +import torch + +from .prompt import MAX_SLOTS + +CONFIG_FILE = "readout.json" +MODEL_CONFIG_FILE = "config.json" +OVERRIDES_KEY = "modeling_overrides" +FULL_ATTENTION_MASK = {"causal": "causal", "noncausal": "block_bidirectional"} +HEAD_FILE = "readout_head.pt" +ATTENTION = ("causal", "noncausal") +READOUT = ("slot_lm", "head") + + +@dataclass +class ReadoutConfig: + attention: str = "causal" + readout: str = "slot_lm" + head_dim: int = 256 + n_slots: int = MAX_SLOTS + + def __post_init__(self) -> None: + if self.attention not in ATTENTION: + raise ValueError(f"attention must be one of {ATTENTION}, got {self.attention!r}") + if self.readout not in READOUT: + raise ValueError(f"readout must be one of {READOUT}, got {self.readout!r}") + + @property + def noncausal(self) -> bool: + return self.attention == "noncausal" + + @property + def head(self) -> bool: + return self.readout == "head" + + @property + def needs_layout(self) -> bool: + return self.noncausal or self.head + + @classmethod + def load(cls, model_dir: str | Path) -> "ReadoutConfig": + path, model_path = Path(model_dir) / CONFIG_FILE, Path(model_dir) / MODEL_CONFIG_FILE + fields = json.loads(path.read_text()) if path.exists() else {} + overrides = json.loads(model_path.read_text()).get(OVERRIDES_KEY) if model_path.exists() else None + if overrides: + mask = overrides["attention"]["full_attention"] + attention = {v: k for k, v in FULL_ATTENTION_MASK.items()}.get(mask) + if attention is None: + raise ValueError(f"unsupported full_attention mask {mask!r} in {model_path}") + if fields.get("attention", attention) != attention: + raise ValueError(f"{path} says attention {fields['attention']!r} but {model_path} says {mask!r}") + fields["attention"] = attention + return cls(**fields) + + def save(self, model_dir: str | Path) -> None: + fields = asdict(self) + model_path = Path(model_dir) / MODEL_CONFIG_FILE + if model_path.exists(): # LoRA adapters have no config.json; their readout.json keeps the attention + attention = fields.pop("attention") + config = json.loads(model_path.read_text()) + config[OVERRIDES_KEY] = {"attention": {"full_attention": FULL_ATTENTION_MASK[attention], + "linear_attention": "causal"}, **fields} + model_path.unlink() # may be a symlink to the source checkpoint (model_v5/mkvariant.py) + model_path.write_text(json.dumps(config, indent=2) + "\n") + (Path(model_dir) / CONFIG_FILE).write_text(json.dumps(fields, indent=2)) + + +_BLOCKS: contextvars.ContextVar = contextvars.ContextVar("qev_readout_blocks", default=None) + + +def install_noncausal(model) -> None: + """Route block ids (set with `blocks`) into the full-attention mask of the model's modeling module. Outside a + `blocks` context the model is unchanged. Needs a mask-based attention implementation (sdpa, eager, flex).""" + impl = getattr(model.config.get_text_config(), "_attn_implementation", None) or model.config._attn_implementation + if impl not in ("sdpa", "eager", "flex_attention"): + raise ValueError(f"noncausal attention needs sdpa, eager or flex_attention, not {impl}") + mod = importlib.import_module(type(model).__module__) + if getattr(mod, "_qev_noncausal", False): + return + if not hasattr(mod, "create_causal_mask"): + raise ValueError(f"{mod.__name__} builds no causal mask to extend") + original = mod.create_causal_mask + + def create_causal_mask(*args, **kwargs): + ids = _BLOCKS.get() + if ids is not None and kwargs.get("block_sequence_ids") is None: + kwargs["block_sequence_ids"] = ids + return original(*args, **kwargs) + + mod.create_causal_mask = create_causal_mask + mod._qev_noncausal = True + + +@contextlib.contextmanager +def blocks(ids): + """Block ids [B, kv_length] over absolute positions (cached prefix included); None leaves the mask causal.""" + token = _BLOCKS.set(ids) + try: + yield + finally: + _BLOCKS.reset(token) + + +def block_ids(prefix_lens: Sequence[int], lengths: Sequence[int], width: int, device) -> torch.Tensor: + """0 on the prefix, 1 on the suffix, -1 on right padding (causal, and masked as a key anyway).""" + pos = torch.arange(width, device=device)[None] + p = torch.tensor(list(prefix_lens), device=device)[:, None] + n = torch.tensor(list(lengths), device=device)[:, None] + return torch.where(pos < p, 0, torch.where(pos < n, 1, -1)) + + +def token_layout(tokenizer, prompt: str, suffix_start: int, option_spans: Sequence[tuple[int, int]] | None): + """Token ids of `prompt`, the prefix length (tokens that end at or before the suffix) and each option line's + token span [start, end), from the tokenizer's character offsets. Option spans are characters of the prompt.""" + enc = tokenizer(prompt, add_special_tokens=False, return_offsets_mapping=True) + ids, offs = enc["input_ids"], enc["offset_mapping"] + ends = [e for _, e in offs] + starts = [s for s, _ in offs] + prefix = bisect.bisect_right(ends, suffix_start) + spans = [] + for cs, ce in option_spans or []: + a = bisect.bisect_right(ends, cs) + b = bisect.bisect_left(starts, ce) + if b <= a: + raise ValueError(f"option span {cs}:{ce} covers no token") + spans.append((a, b)) + return ids, prefix, spans + + +class ReadoutHead(torch.nn.Module): + def __init__(self, hidden: int, n_slots: int = MAX_SLOTS, head_dim: int = 256): + super().__init__() + self.slot = torch.nn.Parameter(torch.zeros(n_slots, hidden)) + self.q = torch.nn.Linear(hidden, head_dim, bias=False) + self.k = torch.nn.Linear(hidden, head_dim, bias=False) + torch.nn.init.normal_(self.q.weight, std=hidden ** -0.5) + torch.nn.init.zeros_(self.k.weight) + + @classmethod + def from_lm_head(cls, lm_head_weight, slot_ids: Sequence[int], head_dim: int) -> "ReadoutHead": + rows = lm_head_weight.detach()[list(slot_ids)].float() + head = cls(rows.shape[1], rows.shape[0], head_dim) + with torch.no_grad(): + head.slot.copy_(rows.to(head.slot.device)) + return head + + def forward(self, h_ans, h_opt): + """h_ans [B, H], h_opt [B, K, H] -> scores [B, K] (fp32).""" + h_ans, h_opt = h_ans.float(), h_opt.float() + k = h_opt.shape[1] + slot = h_ans @ self.slot[:k].T + ptr = torch.einsum("bd,bkd->bk", self.q(h_ans), self.k(h_opt)) / math.sqrt(self.q.out_features) + return slot + ptr + + +def pool_options(hidden, spans: Sequence[Sequence[tuple[int, int]]], kmax: int): + """Mean of hidden [B, T, H] over each option's token span -> [B, kmax, H] fp32 (zeros past a row's options).""" + bsz, _, dim = hidden.shape + dev = hidden.device + rows, toks, segs = [], [], [] + for b, sp in enumerate(spans): + for k, (s, e) in enumerate(sp): + n = e - s + rows.append(torch.full((n,), b, dtype=torch.long)) + toks.append(torch.arange(s, e, dtype=torch.long)) + segs.append(torch.full((n,), b * kmax + k, dtype=torch.long)) + rows, toks, segs = (torch.cat(x).to(dev) for x in (rows, toks, segs)) + flat = hidden[rows, toks].float() + total = torch.zeros(bsz * kmax, dim, dtype=torch.float32, device=dev).index_add(0, segs, flat) + count = torch.zeros(bsz * kmax, dtype=torch.float32, device=dev).index_add( + 0, segs, torch.ones_like(segs, dtype=torch.float32)) + return (total / count.clamp(min=1)[:, None]).view(bsz, kmax, dim) + + +def final_norm(model): + """The final norm of a (possibly multimodal) Qwen-style decoder: its output is the readout hidden state.""" + inner = getattr(model, "model", model) + inner = getattr(inner, "language_model", inner) + norm = getattr(inner, "norm", None) + if norm is None: + raise ValueError(f"cannot find the final norm of {type(model).__name__}") + return norm + + +class HiddenCapture: + """Keeps the final-norm output of the latest forward (with its graph, so the head trains the backbone).""" + + def __init__(self, model): + self.value = None + self.handle = final_norm(model).register_forward_hook(self._hook) + + def _hook(self, module, args, output): + self.value = output + + def pop(self): + out, self.value = self.value, None + if out is None: + raise RuntimeError("no hidden state captured; the forward did not reach the final norm") + return out + + +@dataclass +class Readout: + """A model's readout at runtime: its config, and for `head` the head module and the hidden-state hook.""" + + config: ReadoutConfig + head: ReadoutHead | None = None + capture: HiddenCapture | None = None + + @classmethod + def attach(cls, model, config: ReadoutConfig, slot_ids: Sequence[int] | None = None, + head_path: str | Path | None = None, device=None) -> "Readout": + """Install the variant on `model`. A head is loaded from `head_path` when given, else initialised from the + model's LM-head rows of `slot_ids` (call before the model is sharded).""" + if config.noncausal: + install_noncausal(model) + if not config.head: + return cls(config) + if head_path is not None: + state = torch.load(head_path, map_location="cpu", weights_only=True) + head = ReadoutHead(state["slot"].shape[1], state["slot"].shape[0], state["q.weight"].shape[0]) + head.load_state_dict(state) + else: + head = ReadoutHead.from_lm_head(model.get_output_embeddings().weight, slot_ids[: config.n_slots], + config.head_dim) + head = head.to(device or next(model.parameters()).device) + return cls(config, head, HiddenCapture(model)) + + def save(self, out_dir: str | Path) -> None: + self.config.save(out_dir) + if self.head is not None: + torch.save({k: v.detach().cpu() for k, v in self.head.state_dict().items()}, Path(out_dir) / HEAD_FILE) + + @property + def max_slots(self) -> int: + return self.config.n_slots diff --git a/train/common.py b/train/common.py index 331206f..896e6d3 100644 --- a/train/common.py +++ b/train/common.py @@ -85,6 +85,8 @@ class Example: target: list[float] source: str task: str + suffix_start: int = -1 + option_spans: list[tuple[int, int]] | None = None PAD_FAMILIES = ("next_action",) @@ -185,6 +187,8 @@ def render_example(row: dict, tmpl, idx: int, seed: int | None = None, augment: compact = augment and rng.random() < 0.5 state = serialize_state(row.get("state"), compact) prompt = tmpl(state, rendered.suffix) + suffix_start = prompt.rindex(rendered.suffix) + spans = [(suffix_start + a, suffix_start + b) for a, b in rendered.option_spans or []] return Example( idx=idx, prompt=prompt, @@ -193,6 +197,8 @@ def render_example(row: dict, tmpl, idx: int, seed: int | None = None, augment: target=target_vector(row["target"], rendered.labels), source=str(row.get("source", "unknown")), task=str(row.get("task", "unknown")), + suffix_start=suffix_start, + option_spans=spans or None, ) @@ -244,6 +250,7 @@ class Bundle: template_kwargs: dict processor: Any = None counter: Any = None + readout: Any = None @property def multimodal(self) -> bool: @@ -403,6 +410,65 @@ def readout_logits(model, input_ids, attention_mask, ends): return out.logits[rows, pos] +def batch_scores(bundle: Bundle, model, examples: Sequence[Example], rows: Sequence[dict], device): + """Option scores [B, Kmax] (-inf past each row's options) and their mask for a right-padded batch, under the + bundle's readout (qev.readout): causal or noncausal attention, slot_lm or head scoring. The same function scores + training steps, in-training evals and train.eval_mm, and qev.engine reproduces it for serving.""" + import torch + + from qev import readout as ro + + rd = bundle.readout + cfg = rd.config if rd is not None else ro.ReadoutConfig() + ks = [len(ex.labels) for ex in examples] + if bundle.multimodal: + from train import mm + + inputs, ends = mm.mm_inputs(bundle.processor, [ex.prompt for ex in examples], + [r.get("media") or [] for r in rows], device) + shifts = [bundle.counter.media_extra(r.get("media") or []) for r in rows] + else: + input_ids, attention_mask, ends = pad_batch([bundle.encode(ex.prompt) for ex in examples], bundle.tokenizer.pad_token_id, device) + inputs = {"input_ids": input_ids, "attention_mask": attention_mask} + shifts = [0] * len(examples) + ctx = ro.blocks(None) + spans = None + if cfg.needs_layout: + prefixes, spans = [], [] + for ex, shift, end in zip(examples, shifts, ends): + ids, prefix, sp = ro.token_layout(bundle.tokenizer, ex.prompt, ex.suffix_start, ex.option_spans) + if len(ids) + shift != end + 1: + raise ValueError(f"layout of row {ex.idx} has {len(ids) + shift} tokens, the batch has {end + 1}") + if cfg.head and len(sp) != len(ex.labels): + raise ValueError(f"row {ex.idx}: {len(sp)} option spans for {len(ex.labels)} options") + prefixes.append(prefix + shift) + spans.append([(a + shift, b + shift) for a, b in sp]) + if cfg.noncausal: + width = inputs["input_ids"].shape[1] + ctx = ro.blocks(ro.block_ids(prefixes, [e + 1 for e in ends], width, device)) + with ctx: + if cfg.head: + out = model(**inputs, use_cache=False, return_dict=True, logits_to_keep=1) + hidden = rd.capture.pop() + kmax = max(ks) + rows_ix = torch.arange(len(examples), device=hidden.device) + h_ans = hidden[rows_ix, torch.tensor(ends, device=hidden.device)] + scores = rd.head(h_ans, ro.pool_options(hidden, spans, kmax)) + # every FSDP unit (the LM head too) must take part in the backward + scores = scores + 0.0 * out.logits.float().sum() + mask = torch.zeros((len(ks), kmax), dtype=torch.bool, device=scores.device) + for i, k in enumerate(ks): + mask[i, :k] = True + return scores.masked_fill(~mask, float("-inf")), mask + if bundle.multimodal: + from train import mm + + vocab_logits = mm.mm_readout_logits(model, inputs, ends) + else: + vocab_logits = readout_logits(model, inputs["input_ids"], inputs["attention_mask"], ends) + return slot_view(vocab_logits, bundle.slot_ids, ks) + + def slot_view(vocab_logits, slot_ids, ks): """Gather the first K slot logits per row into a padded [B, Kmax] tensor plus a mask.""" import torch @@ -557,11 +623,7 @@ def collect_slot_logits(bundle: Bundle, rows: Sequence[dict], max_len: int = 409 readout_logits(bundle.model, ids, torch.ones_like(ids), [0]) for bi, batch in enumerate(batches): exs = [examples[i] for i in batch] - id_lists = [bundle.encode(ex.prompt) for ex in exs] - input_ids, attention_mask, ends = pad_batch(id_lists, pad_id, dev) - ks = [len(ex.labels) for ex in exs] - vocab_logits = readout_logits(bundle.model, input_ids, attention_mask, ends) - sel, _ = slot_view(vocab_logits, bundle.slot_ids, ks) + sel, _ = batch_scores(bundle, bundle.model, exs, [rows[ex.idx] for ex in exs], dev) sel = sel.float().cpu().tolist() for r, ex in enumerate(exs): out.append(Scored(kind=ex.kind, k=len(ex.labels), logits=sel[r][: len(ex.labels)], diff --git a/train/eval_mm.py b/train/eval_mm.py index e87e72d..03a6022 100644 --- a/train/eval_mm.py +++ b/train/eval_mm.py @@ -48,6 +48,13 @@ def main() -> int: mm.set_root(args.media_root) bundle = common.load_bundle(args.model, device=args.device, dtype="bfloat16", multimodal=args.multimodal, chat_template=args.chat_template) + from qev.readout import HEAD_FILE, Readout, ReadoutConfig + + rcfg = ReadoutConfig.load(args.model) + head_path = Path(args.model) / HEAD_FILE + bundle.readout = Readout.attach(bundle.model, rcfg, slot_ids=bundle.slot_ids, + head_path=head_path if rcfg.head else None) + print(f"readout: attention {rcfg.attention}, scores {rcfg.readout}", flush=True) results = {} for spec in args.evals: name, _, path = spec.partition("=") diff --git a/train/sft.py b/train/sft.py index 4f6f7bd..ae428d9 100644 --- a/train/sft.py +++ b/train/sft.py @@ -108,6 +108,14 @@ def build_parser() -> argparse.ArgumentParser: "train one step on each listed micro-batch (one GPU), peak memory reset per step, no final eval") p.add_argument("--dump-first-grad", default=None, help="test hook: after the first step's all-reduce, rank 0 saves the adapter gradients here") + p.add_argument("--attention", choices=["causal", "noncausal"], default="causal", + help="noncausal: bidirectional within the prefix and within the question suffix in the full-attention " + "layers (qev.readout)") + p.add_argument("--readout", choices=["slot_lm", "head"], default="slot_lm", + help="head: learned slot + pointer head over the final hidden states instead of LM-head letter logits") + p.add_argument("--head-dim", type=int, default=256, help="pointer projection width of the readout head") + p.add_argument("--head-lr", type=float, default=1e-4, help="peak learning rate of the pointer projections " + "(the head's slot rows use --lr)") p.add_argument("--disable-fla-kernels", action="store_true", help="force the reference PyTorch gated delta rule instead of the fla kernel " "(needed on H100/H200 with Triton in [3.4.0, 3.7.1); correct but much slower)") @@ -326,10 +334,7 @@ def collect_mm(bundle, rows, max_len: int, batch_tokens: int, max_rows: int, dev common.readout_logits(bundle.model, ids, torch.ones_like(ids), [0]) for b in batches: exs = [examples[usable[j]] for j in b] - inputs, ends = mm.mm_inputs(bundle.processor, [ex.prompt for ex in exs], - [rows[ex.idx].get("media") or [] for ex in exs], device) - vocab_logits = mm.mm_readout_logits(bundle.model, inputs, ends) - sel, _ = common.slot_view(vocab_logits, bundle.slot_ids, [len(ex.labels) for ex in exs]) + sel, _ = common.batch_scores(bundle, bundle.model, exs, [rows[ex.idx] for ex in exs], device) sel = sel.float().cpu().tolist() for r, ex in enumerate(exs): out.append(common.Scored(kind=ex.kind, k=len(ex.labels), logits=sel[r][: len(ex.labels)], @@ -339,6 +344,13 @@ def collect_mm(bundle, rows, max_len: int, batch_tokens: int, max_rows: int, dev def save_weights(model, bundle, out_dir: Path, full: bool, sharded: bool = False, is_main: bool = True) -> None: + _save_weights(model, bundle, out_dir, full, sharded, is_main) + if is_main and bundle.readout is not None: + out_dir.mkdir(parents=True, exist_ok=True) + bundle.readout.save(out_dir) + + +def _save_weights(model, bundle, out_dir: Path, full: bool, sharded: bool = False, is_main: bool = True) -> None: """LoRA: the adapter. Full fine-tuning: the whole model in bf16 (from the fp32 master weights, copied to CPU), with tokenizer and processor, loadable like the base checkpoint. Sharded (--fsdp): every rank must call it, the weights are gathered to rank 0's CPU and written there as HF safetensors shards.""" @@ -560,6 +572,8 @@ def main(argv=None) -> int: raise SystemExit("--fsdp is implemented without --share-prefix") if args.disable_fla_kernels: common.block_fla_kernels() + if (args.attention != "causal" or args.readout != "slot_lm") and args.share_prefix: + raise SystemExit("--share-prefix supports only causal slot_lm") import torch from peft import LoraConfig, get_peft_model @@ -600,6 +614,16 @@ def main(argv=None) -> int: if args.plan_only: return plan_only(bundle, rows, args, world) log.info("model_type=%s template_kwargs=%s", bundle.model_type, bundle.template_kwargs) + from qev.readout import Readout, ReadoutConfig + + bundle.readout = Readout.attach(bundle.model, ReadoutConfig(attention=args.attention, readout=args.readout, + head_dim=args.head_dim), + slot_ids=bundle.slot_ids, device=torch.device(args.device)) + head = bundle.readout.head + if head is not None and world > 1: + for prm in head.parameters(): + dist.broadcast(prm.data, 0) + log.info("readout: attention %s, scores %s", args.attention, args.readout) if args.full_finetune: model = bundle.model @@ -696,6 +720,13 @@ def main(argv=None) -> int: opt = torch.optim.AdamW(params, lr=args.lr, weight_decay=args.weight_decay, betas=(0.9, 0.95), eps=1e-8, fused=bool(args.full_finetune)) sched = get_cosine_schedule_with_warmup(opt, num_warmup_steps=warmup, num_training_steps=max(total_steps, 1)) + head_opt = head_sched = None + if head is not None: + head_opt = torch.optim.AdamW([{"params": [head.slot], "lr": args.lr}, + {"params": [head.q.weight, head.k.weight], "lr": args.head_lr}], + weight_decay=0.0, betas=(0.9, 0.95), eps=1e-8) + head_sched = get_cosine_schedule_with_warmup(head_opt, num_warmup_steps=warmup, + num_training_steps=max(total_steps, 1)) logf = open(log_path, "a", encoding="utf-8") if is_main else open(os.devnull, "w") step = 0 @@ -729,25 +760,16 @@ def main(argv=None) -> int: examples.append(common.render_example(rows[i], bundle.template_for(rows[i]), i, seed=seed, augment=args.augment, option_pool=OPTION_POOL, pad_prob=args.pad_options)) - if bundle.multimodal: - from train import mm - - inputs, ends = mm.mm_inputs(bundle.processor, [ex.prompt for ex in examples], - [rows[i].get("media") or [] for i in idx], device) - with torch.autocast("cuda", dtype=torch.bfloat16, enabled=use_autocast(args)): - vocab_logits = mm.mm_readout_logits(model, inputs, ends) - else: - id_lists = [bundle.encode(ex.prompt) for ex in examples] - input_ids, attention_mask, ends = common.pad_batch(id_lists, pad_id, device) - with torch.autocast("cuda", dtype=torch.bfloat16, enabled=use_autocast(args)): - vocab_logits = common.readout_logits(model, input_ids, attention_mask, ends) + with torch.autocast("cuda", dtype=torch.bfloat16, enabled=use_autocast(args)): + sel, mask = common.batch_scores(bundle, model, examples, [rows[i] for i in idx], device) ks = [len(ex.labels) for ex in examples] targets = [ex.target for ex in examples] kmax = max(ks) tgt = torch.zeros((len(ks), kmax), dtype=torch.float32, device=device) for r, t in enumerate(targets): tgt[r, : len(t)] = torch.tensor(t, dtype=torch.float32, device=device) - sel, mask = common.slot_view(vocab_logits, bundle.slot_ids, ks) + if kind == "group": + sel, mask = common.slot_view(vocab_logits, bundle.slot_ids, ks) return sel, mask, tgt if eval_sets and args.eval_at_start: @@ -802,6 +824,17 @@ def main(argv=None) -> int: gnorm = torch.nn.utils.clip_grad_norm_(params, args.max_grad_norm) if hasattr(gnorm, "full_tensor"): gnorm = gnorm.full_tensor() + if head is not None: + for prm in head.parameters(): + if prm.grad is None: + prm.grad = torch.zeros_like(prm) + if world > 1: + dist.all_reduce(prm.grad) + prm.grad /= grad_scale + torch.nn.utils.clip_grad_norm_(list(head.parameters()), args.max_grad_norm) + head_opt.step() + head_sched.step() + head_opt.zero_grad(set_to_none=True) try: opt.step() except torch.OutOfMemoryError: