| """deg model: Qwen3.5-0.8B trunk + decision heads + span pointer + LM head (SPEC.md §1, §12). |
| |
| Layout (token ids; plain text delimiters, no new tokens):: |
| |
| trunk : PREFIX <context tokens> "\\n" -- read once per request |
| query : "FIELD /path\\n" ... "OPTIONS:\\n- a\\n- b\\n" "ANSWER:" -- one per decide/extract |
| field, array count, or jev question; continues a copy of the trunk state |
| gen : "KNOWN VALUES:\\n..." "WRITE /path\\n" ... "VALUE:\\n" <text> <|im_end|> |
| |
| Field independence: the trunk holds the context only (never the schema). Every query is a |
| *fork* that continues its own copy of the trunk state (Gated DeltaNet recurrent state + |
| causal-conv tail, attention K/V); forks never see each other. At inference forks run in |
| fixed :data:`SLOTS`-row micro-batches padded to a :data:`BUCKET` multiple of their own |
| length, so a field's kernels and shapes depend only on that field (bitwise independence; |
| ``tests/test_model.py``). |
| |
| Padding is exact: a padded position has ``beta = 0`` and ``g = 0`` in the delta rule (no |
| write, no decay), its K/V are masked, and the conv tail is gathered at each row's length. |
| |
| Heads (fp32, on the final-norm hidden state; ``a`` = hidden at the query's last token, |
| ``o_i`` = hidden at the last token of option line i, ``c_t`` = trunk hidden at context |
| token t): |
| |
| * choice: ``q(a) . k(o_i) / sqrt(d) + u . o_i`` over the options (jev v1 pointer head); |
| * score: a separate pointer head over the level lines; |
| * noul: ``w . a + b`` -> 2-way logits ``[l, 0]`` (``[p_true, p_false]``); |
| * extract: start/end pointers ``qs(a) . ks(c_t)`` over the context tokens only, plus a null |
| logit ``ns(a)`` (``ne(a)``) as one extra outcome; |
| * count: a linear head over 0..:data:`MAX_COUNT` items, masked to the array bounds (index |
| addressed arrays only: run-1 checkpoints, and arrays whose items have no extract leaf); |
| * anchor: per context token ``t`` an independent "an item starts here" logit |
| ``qa(a) . ka(c_t) / sqrt(d) + ua . c_t + b`` (sigmoid, threshold 0.5), masked to the |
| array's region (the whole context, or the parent item's region for a nested array). The |
| item count is the number of anchors: no cap (SPEC.md §12.6); |
| * anchored extract: the span logits of a field inside an anchored item get a learned bias |
| by bucket of ``t - anchor`` and by whether ``t`` lies inside the item's region; |
| * generate: the base LM head (tied embeddings), only for ``x-generate`` fields. |
| |
| Calibration: one temperature per head kind (``choice, score, noul, extract, count``), fit |
| on validation data after training and stored with the checkpoint. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import bisect |
| import glob |
| import json |
| import math |
| import os |
| import re |
| from dataclasses import dataclass, field |
| from pathlib import Path |
| from typing import Any |
|
|
| import torch |
| import torch.nn.functional as F |
| from torch import Tensor, nn |
| from torch.utils.checkpoint import checkpoint |
|
|
| from deg import profile |
|
|
| |
| HF_HUB = os.environ.get("DEG_HF_HUB", "/data/gwillen-cache/huggingface/hub") |
| BASE_GLOB = f"{HF_HUB}/models--Qwen--Qwen3.5-0.8B/snapshots/*" |
| PREFIX = ("<|im_start|>system\nFill JSON Schema fields from the context: decide, copy spans " |
| "from the context, and write text only for fields marked x-generate.<|im_end|>\n" |
| "<|im_start|>user\nCONTEXT:\n") |
| TRUNK_SUFFIX = "\n" |
| EOS_TEXT = "<|im_end|>" |
| CPU_BUCKET = 64 |
| FORK_TOKENS = 16384 |
| FORK_BYTES_PER_TOKEN = 128 * 1024 |
| FORK_MEM_FRACTION = 0.5 |
| HEAD_LOGIT_BYTES = 64 * 2**20 |
| DENSE_PREFILL_MAX = 8192 |
| PREFILL_CHUNK = 4096 |
| HEAD_ROWS = 8192 |
| MAX_COUNT = 64 |
| ARRAY_MODES = ("index", "anchor") |
| ANCHOR_INIT_BIAS = -6.0 |
| |
| REL_EDGES = (-256, -64, -16, -4, 0, 1, 2, 4, 8, 16, 32, 64, 128, 256) |
| HEAD_DIM = 256 |
| KINDS = ("choice", "score", "noul", "extract", "count") |
|
|
|
|
|
|
|
|
| @dataclass(frozen=True) |
| class BackboneSpec: |
| """One supported base model (``deg.train --backbone NAME``, ``deg.json["backbone"]``). |
| |
| The deg layout (queries, KNOWN VALUES) is the same text for every backbone; a spec |
| says how to load the weights and tokenizer and the chat tokens that differ: ``bos`` |
| (prepended to every trunk), ``pad`` (padding token), ``gen_turn`` (closes the user turn |
| and opens the assistant turn for base-model generation, |
| :func:`deg.infer.base_generate_header`; thinking off), ``prefix`` (system turn + |
| opened user turn before the context; ChatML :data:`PREFIX` by default) and ``eos`` |
| (the end-of-turn token every generated value ends with; ChatML :data:`EOS_TEXT`).""" |
|
|
| name: str |
| snapshot_glob: str |
| lm_class: str |
| bos: str |
| pad: str |
| gen_turn: str |
| prefix: str = PREFIX |
| eos: str = EOS_TEXT |
|
|
|
|
| QWEN_GEN_TURN = "<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n" |
|
|
|
|
| def qwen35_spec(size: str) -> BackboneSpec: |
| """Qwen3.5 of any size (``0.8B``, ``2B``, ``4B``, ...): one family, one tokenizer and |
| chat format; layer counts, hidden/MLP widths and heads come from the snapshot's |
| config (:class:`Backbone` reads them).""" |
| return BackboneSpec( |
| f"qwen3.5-{size.lower()}", f"{HF_HUB}/models--Qwen--Qwen3.5-{size.upper()}/snapshots/*", |
| "Qwen3_5ForCausalLM", "", "<|endoftext|>", QWEN_GEN_TURN) |
|
|
|
|
| |
| |
| |
| |
| GEMMA_PREFIX = ("<|turn>system\nFill JSON Schema fields from the context: decide, copy spans " |
| "from the context, and write text only for fields marked x-generate.<turn|>\n" |
| "<|turn>user\nCONTEXT:\n") |
| GEMMA_GEN_TURN = "<turn|>\n<|turn>model\n" |
| |
| |
| GEMMA4_E2B_IT_REVISION = "3e22461f65e89153144f8adb70e3b8c2cc9845a7" |
|
|
|
|
| def gemma4_spec(repo: str) -> BackboneSpec: |
| """google/gemma-4-E2B(-it): the text model only (Gemma4ForCausalLM; vision/audio |
| towers are not loaded), Gemma's chat format.""" |
| return BackboneSpec(repo.split("/")[-1].lower(), |
| f"{HF_HUB}/models--{repo.replace('/', '--')}/snapshots/*", |
| "Gemma4ForCausalLM", "<bos>", "<pad>", GEMMA_GEN_TURN, GEMMA_PREFIX, |
| "<turn|>") |
|
|
|
|
| BACKBONES = { |
| "qwen3.5-0.8b": qwen35_spec("0.8B"), |
| "qwen3.5-2b": qwen35_spec("2B"), |
| |
| |
| "lfm2.5-1.2b": BackboneSpec( |
| "lfm2.5-1.2b", f"{HF_HUB}/models--LiquidAI--LFM2.5-1.2B-Instruct/snapshots/*", |
| "Lfm2ForCausalLM", "<|startoftext|>", "<|pad|>", |
| "<|im_end|>\n<|im_start|>assistant\n"), |
| |
| |
| "gemma-4-e2b-it": gemma4_spec("google/gemma-4-E2B-it"), |
| "gemma-4-e2b": gemma4_spec("google/gemma-4-E2B"), |
| } |
| DEFAULT_BACKBONE = "qwen3.5-0.8b" |
|
|
|
|
| def backbone_name(name: str) -> str: |
| """Canonical backbone name: ``Qwen/Qwen3.5-2B`` (HF repo id, ``--base-model``) and |
| ``qwen3.5-2b`` are the same backbone; ``LiquidAI/LFM2.5-1.2B-Instruct`` is |
| ``lfm2.5-1.2b``.""" |
| low = name.strip().lower() |
| if low == "liquidai/lfm2.5-1.2b-instruct": |
| return "lfm2.5-1.2b" |
| for org in ("qwen/", "google/"): |
| if low.startswith(org): |
| low = low[len(org):] |
| return low |
|
|
|
|
| def backbone_spec(name: str | None = None) -> BackboneSpec: |
| name = backbone_name(name or DEFAULT_BACKBONE) |
| if name in BACKBONES: |
| return BACKBONES[name] |
| m = re.fullmatch(r"qwen3\.5-(\d+(?:\.\d+)?b)", name) |
| if m: |
| return qwen35_spec(m.group(1)) |
| raise ValueError(f"backbone must be one of {sorted(BACKBONES)} or qwen3.5-<size>: " |
| f"{name!r}") |
|
|
|
|
| def base_snapshot(backbone: str | None = None) -> Path: |
| spec = backbone_spec(backbone) |
| found = sorted(glob.glob(spec.snapshot_glob)) |
| if not found: |
| raise FileNotFoundError(f"{spec.name} snapshot not found under {spec.snapshot_glob}") |
| return Path(found[-1]) |
|
|
|
|
| def load_base_lm(backbone: str | None = None, dtype: torch.dtype = torch.bfloat16 |
| ) -> nn.Module: |
| """The base causal LM of ``backbone`` (text model only) from the local HF cache.""" |
| import transformers |
|
|
| spec = backbone_spec(backbone) |
| snap = str(base_snapshot(spec.name)) |
| if spec.lm_class == "Gemma4ForCausalLM": |
| |
| |
| cfg = transformers.AutoConfig.from_pretrained(snap) |
| return transformers.Gemma4ForCausalLM.from_pretrained( |
| snap, config=getattr(cfg, "text_config", cfg), dtype=dtype, |
| key_mapping={r"^model\.language_model\.": "model."}) |
| return getattr(transformers, spec.lm_class).from_pretrained(snap, dtype=dtype) |
|
|
|
|
| def spec_of_lm(lm: nn.Module, backbone: str | None = None) -> BackboneSpec: |
| """The backbone spec whose architecture ``lm`` is (tests pass tiny random models): |
| ``backbone`` when it names that architecture's family (a tiny 2B-shaped Qwen3.5 may |
| stand in for ``qwen3.5-2b``), else the family's default size.""" |
| kind = lm.config.model_type |
| cls = {"lfm2": "Lfm2ForCausalLM", "gemma4_text": "Gemma4ForCausalLM"}.get( |
| kind, "Qwen3_5ForCausalLM") |
| if backbone is not None: |
| spec = backbone_spec(backbone) |
| if spec.lm_class != cls: |
| raise ValueError(f"lm is a {kind} model, not {spec.name}") |
| return spec |
| return BACKBONES[{"Lfm2ForCausalLM": "lfm2.5-1.2b", |
| "Gemma4ForCausalLM": "gemma-4-e2b-it"}.get(cls, "qwen3.5-0.8b")] |
|
|
|
|
| |
|
|
|
|
| @dataclass(frozen=True) |
| class Query: |
| """One fork's input: what the field is, rendered from the field alone. |
| |
| ``anchor``/``region`` (context token indices; ``region`` = [lo, hi), hi = -1 means the |
| end of the context) locate the item this query belongs to: an anchor query's logits |
| are masked to ``region``; an extract field inside an anchored item gets the relative |
| span bias from ``anchor`` and ``region[1]``. ``anchor = -1``: not inside an anchored |
| item.""" |
|
|
| kind: str |
| header: str |
| options: tuple[str, ...] = () |
| anchor: int = -1 |
| region: tuple[int, int] = (0, -1) |
|
|
|
|
| @dataclass(frozen=True) |
| class EncodedQuery: |
| kind: str |
| ids: tuple[int, ...] |
| option_positions: tuple[int, ...] |
| answer_position: int |
| anchor: int = -1 |
| region: tuple[int, int] = (0, -1) |
|
|
|
|
| @dataclass(frozen=True) |
| class EncodedContext: |
| ids: tuple[int, ...] |
| ctx_start: int |
| offsets: tuple[tuple[int, int], ...] |
| |
| |
| media: tuple[tuple[int, Tensor], ...] = () |
|
|
|
|
| @dataclass(frozen=True) |
| class Media: |
| """One image / audio item of a multimodal context: ``embeds`` [n, hidden] are its soft |
| tokens in the language model's space (:meth:`MediaTowers.image_embeds` / |
| :meth:`MediaTowers.audio_embeds`, or a precomputed cache).""" |
|
|
| kind: str |
| embeds: Tensor |
|
|
|
|
| class Renderer: |
| def __init__(self, tokenizer: Any, spec: BackboneSpec | None = None) -> None: |
| spec = spec or backbone_spec() |
| self.tokenizer = tokenizer |
| self.eos_id = tokenizer.convert_tokens_to_ids(spec.eos) |
| self.pad_id = tokenizer.convert_tokens_to_ids(spec.pad) |
| unk = getattr(tokenizer, "unk_token_id", None) |
| if unk is not None and unk in (self.eos_id, self.pad_id): |
| raise ValueError(f"{spec.name}: eos {spec.eos!r} / pad {spec.pad!r} is not a " |
| "token of this tokenizer") |
| self.gen_turn = spec.gen_turn |
| self._prefix = self.enc(spec.bos + spec.prefix) |
| self._suffix = self.enc(TRUNK_SUFFIX) |
|
|
| def enc(self, text: str) -> list[int]: |
| return list(self.tokenizer(text, add_special_tokens=False)["input_ids"]) if text else [] |
|
|
| def context(self, text: str | list[str | Media]) -> EncodedContext: |
| """The trunk of a context: text, or (multimodal, Gemma 4) a list of text parts and |
| :class:`Media` items. Offsets are code points of the joined text parts; a media |
| item renders as HF Gemma 4's placeholder span (``<boi> <image>*n <eoi>``, audio |
| ``<boa> <audio>*n <eoa>``) whose tokens get zero-width offsets at their place in |
| that text, so the span pointer never addresses them (text-only extraction).""" |
| if isinstance(text, str): |
| enc = self.tokenizer(text, add_special_tokens=False, return_offsets_mapping=True) |
| ids = [*self._prefix, *enc["input_ids"], *self._suffix] |
| return EncodedContext(tuple(ids), len(self._prefix), |
| tuple((int(a), int(b)) for a, b in enc["offset_mapping"])) |
| tok = self.tokenizer |
| ids, offsets, media, pos = list(self._prefix), [], [], 0 |
| for part in text: |
| if isinstance(part, str): |
| enc = tok(part, add_special_tokens=False, return_offsets_mapping=True) |
| ids += enc["input_ids"] |
| offsets += [(pos + int(a), pos + int(b)) for a, b in enc["offset_mapping"]] |
| pos += len(part) |
| continue |
| names = {"image": ("boi_token", "image_token", "eoi_token"), |
| "audio": ("boa_token", "audio_token", "eoa_token")}.get(part.kind) |
| if names is None or not all(getattr(tok, n, None) for n in names): |
| raise ValueError(f"no {part.kind!r} placeholder tokens in this tokenizer") |
| bo, soft, eo = (tok.convert_tokens_to_ids(getattr(tok, n)) for n in names) |
| n = int(part.embeds.shape[0]) |
| media.append((len(ids) + 1, part.embeds)) |
| ids += [bo, *[soft] * n, eo] |
| offsets += [(pos, pos)] * (n + 2) |
| ids += self._suffix |
| return EncodedContext(tuple(ids), len(self._prefix), tuple(offsets), tuple(media)) |
|
|
| def query(self, q: Query) -> EncodedQuery: |
| ids = self.enc(q.header) |
| positions = [] |
| if q.options: |
| ids += self.enc("OPTIONS:\n") |
| for text in q.options: |
| ids += self.enc(f"- {text}") |
| positions.append(len(ids) - 1) |
| ids += self.enc("\n") |
| ids += self.enc("ANSWER:") |
| return EncodedQuery(q.kind, tuple(ids), tuple(positions), len(ids) - 1, q.anchor, |
| q.region) |
|
|
| def known_block(self, known: list[tuple[str, Any]]) -> list[int]: |
| """The shared generate prefix: every decided/extracted value, one line each.""" |
| lines = ["KNOWN VALUES:"] |
| lines += [f"{ptr} = {json.dumps(v, ensure_ascii=False)}" for ptr, v in known] |
| return self.enc("\n".join(lines) + "\n") |
|
|
| def generate_prompt(self, known: list[tuple[str, Any]], header: str) -> list[int]: |
| """Training layout of one generate fork: the shared known block, then the field's |
| own header (inference prefills the block once and forks per field).""" |
| return self.known_block(known) + self.enc(header + "VALUE:\n") |
|
|
|
|
| def load_tokenizer(snapshot: Path | None = None, backbone: str | None = None): |
| from transformers import AutoTokenizer |
|
|
| return AutoTokenizer.from_pretrained(str(snapshot or base_snapshot(backbone))) |
|
|
|
|
| _SPAN_INDEX: list[Any] = [None, None, None] |
|
|
|
|
| def token_span(offsets: tuple[tuple[int, int], ...], start: int, end: int) -> tuple[int, int]: |
| """Char span [start, end) -> inclusive context-token span covering it: the first token |
| ending after ``start`` and the last token starting before ``end``. O(log T) by |
| bisection (tokenizer offsets are non-decreasing); a non-monotone offset list falls |
| back to the linear scan.""" |
| if _SPAN_INDEX[0] is not offsets: |
| starts = [a for a, _ in offsets] |
| ends = [b for _, b in offsets] |
| mono = all(x <= y for x, y in zip(starts, starts[1:])) and all( |
| x <= y for x, y in zip(ends, ends[1:])) |
| _SPAN_INDEX[:] = [offsets, starts if mono else None, ends] |
| starts, ends = _SPAN_INDEX[1], _SPAN_INDEX[2] |
| if starts is None: |
| i = next(k for k, (_, b) in enumerate(offsets) if b > start) |
| j = max(k for k, (a, _) in enumerate(offsets) if a < end) |
| return i, j |
| i = bisect.bisect_right(ends, start) |
| j = bisect.bisect_left(starts, end) - 1 |
| if i >= len(offsets): |
| raise StopIteration |
| if j < 0: |
| raise ValueError("no token starts before end") |
| return i, j |
|
|
|
|
| _OPEN = {"(": ")", "[": "]", "{": "}", "<": ">", '"': '"', "“": "”", "‘": "’", "«": "»"} |
| _CLOSE = {v: k for k, v in _OPEN.items()} |
| _TRAIL = set(",;:'\"”’»") |
| _LEAD = set(">'\"“‘«") |
| _QUOTES = set('"“‘«[<{') |
|
|
|
|
| def _run(text: str, i: int, step: int, lo: int, hi: int) -> int: |
| """Length of the '*' run starting at ``i`` walking by ``step`` within [lo, hi).""" |
| n = 0 |
| while lo <= i < hi and text[i] == "*": |
| n, i = n + 1, i + step |
| return n |
|
|
|
|
| def _emphasis(text: str, start: int, end: int) -> tuple[int, int] | None: |
| """Strip markdown ``*``/``**`` emphasis only as a symmetric wrapper: the same run of |
| ``*`` opens and closes it, either inside the span or (when the token span stopped |
| short) right outside it. A one-sided run (``****4821``, ``AB*``) is value text.""" |
| lead = _run(text, start, 1, start, end) |
| trail = _run(text, end - 1, -1, start, end) |
| if lead and lead == trail and lead < end - start - lead: |
| return start + lead, end - lead |
| if lead and not trail and _run(text, end, 1, end, len(text)) == lead \ |
| and end - start > lead: |
| return start + lead, end |
| if trail and not lead and _run(text, start - 1, -1, 0, start) == trail \ |
| and end - start > trail: |
| return start, end - trail |
| return None |
|
|
|
|
| def trim_span(text: str, start: int, end: int) -> tuple[int, int]: |
| """Token span chars -> value chars: strip whitespace and boundary punctuation that a |
| token carried in (unbalanced brackets/quotes, trailing ``,;:``, symmetric markdown |
| ``*`` emphasis). ``*`` is never stripped one-sided: it can be value (``****4821``).""" |
| while True: |
| while start < end and text[start].isspace(): |
| start += 1 |
| while end > start and text[end - 1].isspace(): |
| end -= 1 |
| if end - start < 2: |
| return start, end |
| a, z = text[start], text[end - 1] |
| inner = text[start:end] |
| emph = _emphasis(text, start, end) if "*" in (a, z) else None |
| if emph is not None: |
| start, end = emph |
| elif a in _QUOTES and z == _OPEN[a]: |
| start, end = start + 1, end - 1 |
| elif a in _OPEN and _OPEN[a] not in inner[1:]: |
| start += 1 |
| elif a in _LEAD and a not in _OPEN: |
| start += 1 |
| elif z in _CLOSE and _CLOSE[z] not in inner[:-1]: |
| end -= 1 |
| elif z in _TRAIL and not (z in _CLOSE and _CLOSE[z] in inner[:-1]): |
| end -= 1 |
| else: |
| return start, end |
|
|
|
|
| |
|
|
|
|
| @dataclass |
| class State: |
| """Prefix state for B rows. ``lin[i]`` = (recurrent state fp32, conv tail [B,K-1,C]) |
| per state layer (Qwen3.5 Gated DeltaNet: recurrent [B,H,K,V]; LFM2 short conv: no |
| recurrent state, an empty [B,0] placeholder); ``kv[i]`` = (K, V) [B,Hkv,P,D] per attention layer; |
| ``kv_mask`` [B,P] marks real prefix positions; ``lengths`` [B] = next position.""" |
|
|
| lin: list[tuple[Tensor, Tensor]] |
| kv: list[tuple[Tensor, Tensor]] |
| kv_mask: Tensor |
| lengths: Tensor |
|
|
| def select(self, index: Tensor) -> State: |
| return State([(r.index_select(0, index), t.index_select(0, index)) for r, t in self.lin], |
| [(k.index_select(0, index), v.index_select(0, index)) for k, v in self.kv], |
| self.kv_mask.index_select(0, index), self.lengths.index_select(0, index)) |
|
|
|
|
| def _l2norm(x: Tensor, eps: float = 1e-6) -> Tensor: |
| return x * torch.rsqrt((x * x).sum(-1, keepdim=True) + eps) |
|
|
|
|
| def delta_rule_reference(q, k, v, g, beta, initial_state): |
| """Exact per-token gated delta rule (CPU path; matches HF's recurrent reference).""" |
| dtype = v.dtype |
| q, k, v = _l2norm(q.float()), _l2norm(k.float()), v.float() |
| q = q * (q.shape[-1] ** -0.5) |
| B, L, H, Dk = k.shape |
| state = (initial_state.float() if initial_state is not None |
| else torch.zeros(B, H, Dk, v.shape[-1], device=v.device)) |
| outs = [] |
| for t in range(L): |
| state = state * g[:, t].float().exp()[..., None, None] |
| mem = (state * k[:, t, :, :, None]).sum(-2) |
| delta = (v[:, t] - mem) * beta[:, t].float()[..., None] |
| state = state + k[:, t, :, :, None] * delta[:, :, None, :] |
| outs.append((state * q[:, t, :, :, None]).sum(-2)) |
| return torch.stack(outs, 1).to(dtype), state |
|
|
|
|
| def _delta_rule(q, k, v, g, beta, initial_state): |
| if q.is_cuda: |
| from fla.ops.gated_delta_rule import chunk_gated_delta_rule |
|
|
| return chunk_gated_delta_rule(q, k, v, g=g, beta=beta, initial_state=initial_state, |
| output_final_state=True, use_qk_l2norm_in_kernel=True) |
| return delta_rule_reference(q, k, v, g, beta, initial_state) |
|
|
|
|
| |
| |
|
|
|
|
| def _rms(x: Tensor, w: Tensor, eps: float) -> Tensor: |
| xf = x.float() |
| return (xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + eps) * (1.0 + w.float()) |
| ).type_as(x) |
|
|
|
|
| def _swiglu(gate: Tensor, up: Tensor) -> Tensor: |
| return F.silu(gate) * up |
|
|
|
|
| def _dn_gates(b: Tensor, a: Tensor, A_log: Tensor, dt_bias: Tensor, valid: Tensor |
| ) -> tuple[Tensor, Tensor]: |
| """beta, g (log decay) with padded positions made exact no-ops (beta = 0, g = 0).""" |
| v = valid[..., None] |
| beta = b.sigmoid() * v |
| g = -A_log.float().exp() * F.softplus(a.float() + dt_bias) * v |
| return beta.to(b.dtype), g |
|
|
|
|
| def _conv_silu(ext: Tensor, w: Tensor) -> Tensor: |
| """Depthwise causal conv over ext [B, K-1+L, C] with w [C, K] -> silu, [B, L, C].""" |
| K = w.shape[1] |
| L = ext.shape[1] - K + 1 |
| w = w.to(ext.dtype) |
| out = ext[:, 0:L] * w[:, 0] |
| for k in range(1, K): |
| out = out + ext[:, k:k + L] * w[:, k] |
| return F.silu(out) |
|
|
|
|
| def _gated_rms(x: Tensor, gate: Tensor, w: Tensor, eps: float) -> Tensor: |
| xf = x.float() |
| xf = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + eps) |
| return ((w * xf.to(x.dtype)) * F.silu(gate.float())).to(x.dtype) |
|
|
|
|
| def _attn_gate(o: Tensor, gate: Tensor) -> Tensor: |
| return o * torch.sigmoid(gate) |
|
|
|
|
| def _rms_w(x: Tensor, w: Tensor, eps: float) -> Tensor: |
| """LFM2 RMSNorm: ``w * norm(x).to(dtype)`` (weight not offset by 1, as in HF).""" |
| xf = x.float() |
| return w * (xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + eps)).to(x.dtype) |
|
|
|
|
| def _short_conv(ext: Tensor, w: Tensor) -> Tensor: |
| """LFM2 depthwise causal conv over ext [B, K-1+L, C] with w [C, K], no activation, |
| fp32 accumulation in a fixed tap order (elementwise: batch-invariant), [B, L, C].""" |
| K = w.shape[1] |
| L = ext.shape[1] - K + 1 |
| e, wf = ext.float(), w.float() |
| out = e[:, 0:L] * wf[:, 0] |
| for k in range(1, K): |
| out = out + e[:, k:k + L] * wf[:, k] |
| return out.to(ext.dtype) |
|
|
|
|
| def _rms_g(x: Tensor, w: Tensor, eps: float) -> Tensor: |
| """Gemma 4 RMSNorm (HF ``Gemma4RMSNorm``): ``norm(x) * w`` in fp32, then cast.""" |
| xf = x.float() |
| return (xf * torch.pow(xf.pow(2).mean(-1, keepdim=True) + eps, -0.5) * w.float() |
| ).type_as(x) |
|
|
|
|
| def _rms_n(x: Tensor, eps: float) -> Tensor: |
| """Gemma 4 RMSNorm without a scale (``with_scale=False``: the value norm).""" |
| xf = x.float() |
| return (xf * torch.pow(xf.pow(2).mean(-1, keepdim=True) + eps, -0.5)).type_as(x) |
|
|
|
|
| def _geglu(gate: Tensor, up: Tensor) -> Tensor: |
| """``gelu_pytorch_tanh(gate) * up`` (Gemma's MLP and its per-layer-input gate).""" |
| return F.gelu(gate, approximate="tanh") * up |
|
|
|
|
| _FUSED: dict[str, Any] = {} |
| _RAW = {"rms": _rms, "swiglu": _swiglu, "dn_gates": _dn_gates, "conv_silu": _conv_silu, |
| "gated_rms": _gated_rms, "attn_gate": _attn_gate, "rms_w": _rms_w, |
| "short_conv": _short_conv, "rms_g": _rms_g, "rms_n": _rms_n, "geglu": _geglu} |
|
|
|
|
| def _op(name: str, x: Tensor): |
| if not x.is_cuda or not COMPILE: |
| return _RAW[name] |
| if name not in _FUSED: |
| import torch._dynamo.config as dc |
|
|
| dc.recompile_limit = max(dc.recompile_limit, 64) |
| dc.accumulated_recompile_limit = max(dc.accumulated_recompile_limit, 512) |
| _FUSED[name] = torch.compile(_RAW[name], dynamic=True) |
| return _FUSED[name] |
|
|
|
|
| COMPILE = True |
|
|
|
|
| def _norm(module: nn.Module, x: Tensor) -> Tensor: |
| return _op("rms", x)(x, module.weight, module.eps) |
|
|
|
|
| |
| |
| |
| |
| |
| |
| |
| |
| LORA_TARGETS = {"linear_attn": ("in_proj_qkv", "in_proj_z", "out_proj"), |
| "self_attn": ("q_proj", "k_proj", "v_proj", "o_proj"), |
| "mlp": ("gate_proj", "up_proj", "down_proj")} |
| |
| ADAPTER_OF_KIND = {"choice": "decide", "score": "decide", "noul": "decide", |
| "extract": "extract", "count": "extract", "anchor": "extract"} |
|
|
|
|
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
|
|
| @dataclass(frozen=True) |
| class WriteRoute: |
| """``sel`` of a writer (W-expert) call: ``lora`` [B, r] ones (the writer LoRA on) or |
| None (an MoE without a writer LoRA). Indexing slices the rows.""" |
|
|
| lora: Tensor | None = None |
|
|
| def __getitem__(self, idx: Any) -> WriteRoute: |
| return WriteRoute(None if self.lora is None else self.lora[idx]) |
|
|
|
|
| |
| |
| |
| |
| |
| |
| |
| TASK_EXPERT_OF_KIND = {"choice": "d", "score": "d", "noul": "d", "extract": "e", |
| "count": "e", "anchor": "e", "generate": "g"} |
|
|
|
|
| @dataclass(frozen=True) |
| class TaskRoute: |
| """``sel`` of a D/E/G MoE call: every row runs FFN expert ``expert`` (``e`` = the |
| layer's ``mlp``, else ``mlp_<expert>``). Indexing keeps the route.""" |
|
|
| expert: str |
|
|
| def __getitem__(self, idx: Any) -> TaskRoute: |
| return self |
|
|
|
|
| def _lsel(sel: Tensor | WriteRoute | TaskRoute | None) -> Tensor | None: |
| """The LoRA selector inside a ``sel`` (a WriteRoute's writer LoRA columns).""" |
| if isinstance(sel, TaskRoute): |
| return None |
| return sel.lora if isinstance(sel, WriteRoute) else sel |
|
|
|
|
| def _expert(layer, sel: Tensor | WriteRoute | TaskRoute | None, name: str = "mlp" |
| ) -> nn.Module: |
| """The FFN expert of this call: W (``<name>_w``) for a WriteRoute, ``<name>_<x>`` for a |
| TaskRoute to expert x (``e``: ``<name>`` itself), else R / the dense FFN.""" |
| if isinstance(sel, WriteRoute): |
| return getattr(layer, name + "_w") |
| if isinstance(sel, TaskRoute) and sel.expert != "e": |
| return getattr(layer, f"{name}_{sel.expert}") |
| return getattr(layer, name) |
|
|
|
|
| def _lin(module: nn.Module, x: Tensor, sel: Tensor | WriteRoute | None) -> Tensor: |
| """``module(x)`` plus the selected LoRA delta (training / trunk path, autograd).""" |
| y = module(x) |
| sel = _lsel(sel) |
| if sel is None: |
| return y |
| A = getattr(module, "lora_A", None) |
| if A is None: |
| return y |
| t = F.linear(x, A) * sel[:, None, :].to(x.dtype) |
| return y + F.linear(t, module.lora_B) * module.lora_scale |
|
|
|
|
| def _flora(module: nn.Module, x: Tensor, y: Tensor, sel: Tensor | WriteRoute | None |
| ) -> Tensor: |
| """Fork path: ``y`` (= x @ W^T) plus the selected LoRA delta through the same |
| batch-invariant GEMM config as the base projection.""" |
| sel = _lsel(sel) |
| if sel is None: |
| return y |
| A = getattr(module, "lora_A", None) |
| if A is None: |
| return y |
| from deg.kernels import bi_linear, bi_linear_small |
|
|
| mm = bi_linear_small if _SMALL_GEMM else bi_linear |
| t = mm(x, A) * sel[:, None, :].to(x.dtype) |
| return y + mm(t, module.lora_B) * module.lora_scale |
|
|
|
|
| def _mlp(layer, h: Tensor, sel: Tensor | WriteRoute | None = None) -> Tensor: |
| m = _expert(layer, sel) |
| x = _norm(layer.post_attention_layernorm, h) |
| a = _op("swiglu", x)(_lin(m.gate_proj, x, sel), _lin(m.up_proj, x, sel)) |
| return h + _lin(m.down_proj, a, sel) |
|
|
|
|
| def _linear_layer(layer, h, valid, lengths_new, rec, tail, rows=None, sel=None): |
| """``rows``: ``rec``/``tail`` are the un-selected prefix states and row b continues |
| prefix row ``rows[b]``; selecting here (inside the checkpoint) keeps the per-row copies |
| out of the saved activations: they are rebuilt during the backward recompute.""" |
| attn = layer.linear_attn |
| B, L, _ = h.shape |
| if rows is not None: |
| rec, tail = rec.index_select(0, rows), tail.index_select(0, rows) |
| xn = _norm(layer.input_layernorm, h) |
| mixed = _lin(attn.in_proj_qkv, xn, sel) |
| z = _lin(attn.in_proj_z, xn, sel).reshape(B, L, -1, attn.head_v_dim) |
| beta, g = _op("dn_gates", h)(attn.in_proj_b(xn), attn.in_proj_a(xn), attn.A_log, |
| attn.dt_bias, valid) |
| ext = torch.cat([tail.to(mixed.dtype), mixed], 1) |
| K1 = tail.shape[1] |
| conv = _op("conv_silu", h)(ext, attn.conv1d.weight[:, 0]) |
| gather = lengths_new[:, None] + torch.arange(K1, device=h.device)[None] |
| new_tail = ext.gather(1, gather[..., None].expand(-1, -1, ext.shape[-1])) |
| qq, kk, vv = torch.split(conv, [attn.key_dim, attn.key_dim, attn.value_dim], dim=-1) |
| qq = qq.reshape(B, L, -1, attn.head_k_dim) |
| kk = kk.reshape(B, L, -1, attn.head_k_dim) |
| vv = vv.reshape(B, L, -1, attn.head_v_dim) |
| rep = attn.num_v_heads // attn.num_k_heads |
| if rep > 1: |
| qq, kk = qq.repeat_interleave(rep, dim=2), kk.repeat_interleave(rep, dim=2) |
| core, new_rec = _delta_rule(qq, kk, vv, g, beta.to(qq.dtype), rec) |
| core = _op("gated_rms", h)(core, z, attn.norm.weight, attn.norm.variance_epsilon) |
| h = h + _lin(attn.out_proj, core.reshape(B, L, -1), sel) |
| return _mlp(layer, h, sel), new_rec, new_tail |
|
|
|
|
| def _attn_qkv(layer, h, cos, sin, sel=None): |
| """Attention projections: q, k, v [B, H, L, D] (rotary applied) and the output gate.""" |
| from transformers.models.qwen3_5.modeling_qwen3_5 import apply_rotary_pos_emb |
|
|
| a = layer.self_attn |
| B, L, _ = h.shape |
| xn = _norm(layer.input_layernorm, h) |
| q, gate = torch.chunk(_lin(a.q_proj, xn, sel).view(B, L, -1, a.head_dim * 2), 2, dim=-1) |
| q = _norm(a.q_norm, q).transpose(1, 2) |
| k = _norm(a.k_norm, _lin(a.k_proj, xn, sel).view(B, L, -1, a.head_dim)).transpose(1, 2) |
| v = _lin(a.v_proj, xn, sel).view(B, L, -1, a.head_dim).transpose(1, 2) |
| q, k = apply_rotary_pos_emb(q, k, cos, sin) |
| return q, k, v, gate |
|
|
|
|
| def _attn_out(layer, h, o, gate, sel=None): |
| """o [B, L, Hq*D] -> gated output projection, residual, MLP.""" |
| B, L, _ = h.shape |
| h = h + _lin(layer.self_attn.o_proj, _op("attn_gate", h)(o, gate.reshape(B, L, -1)), sel) |
| return _mlp(layer, h, sel) |
|
|
|
|
| def _attn_layer(fam, layer, h, cos, sin, k_prev, v_prev, mask, rows=None, sel=None): |
| """``rows``: as in :func:`_linear_layer` (prefix K/V selected inside the checkpoint); |
| ``fam``: the :class:`Family` whose projections run.""" |
| a = layer.self_attn |
| B, L, _ = h.shape |
| if rows is not None: |
| k_prev, v_prev = k_prev.index_select(0, rows), v_prev.index_select(0, rows) |
| q, k, v, gate = fam.attn_qkv(layer, h, cos, sin, sel) |
| keys = torch.cat([k_prev.to(k.dtype), k], 2) |
| values = torch.cat([v_prev.to(v.dtype), v], 2) |
| o = F.scaled_dot_product_attention(q, keys, values, attn_mask=mask[:, None], |
| scale=a.scaling, enable_gqa=True) |
| o = o.transpose(1, 2).reshape(B, L, -1) |
| return fam.attn_out(layer, h, o, gate, sel), k, v |
|
|
|
|
| def _chunk_attn_layer(fam, layer, h, cos, sin, kparts, vparts, kv_mask, s, sel=None): |
| """One attention layer over one chunk of a long right-padded input (training, |
| :meth:`Backbone._forward_chunked_ckpt`): keys are the incoming prefix (``kv_mask``), |
| the ``s`` earlier tokens of this call (all real for a real query: right padding) and |
| the chunk's own tokens (causal). The [B, chunk, P+s+chunk] mask lives only inside the |
| checkpoint. Returns (h, k, v) of the chunk.""" |
| a = layer.self_attn |
| B, c, _ = h.shape |
| q, k, v, gate = fam.attn_qkv(layer, h, cos, sin, sel) |
| keys = torch.cat([*(t.to(k.dtype) for t in kparts), k], 2) |
| values = torch.cat([*(t.to(v.dtype) for t in vparts), v], 2) |
| causal = torch.ones(c, c, dtype=torch.bool, device=h.device).tril() |
| mask = torch.cat([kv_mask[:, None, :].expand(-1, c, -1), |
| torch.ones(B, c, s, dtype=torch.bool, device=h.device), |
| causal[None].expand(B, -1, -1)], 2) |
| o = F.scaled_dot_product_attention(q, keys, values, attn_mask=mask[:, None], |
| scale=a.scaling, enable_gqa=True) |
| return fam.attn_out(layer, h, o.transpose(1, 2).reshape(B, c, -1), gate, sel), k, v |
|
|
|
|
| @dataclass |
| class ForkState: |
| """B fork rows, each continuing one of N prefixes (SPEC.md §12.2). |
| |
| Prefixes (batch-1, unpadded trunk states, e.g. the contexts of several requests) are |
| read by reference: attention reads row b's prefix K/V at ``pk``/``pv`` [Hkv, total, D] |
| positions ``poff[b] .. poff[b]+plen[b]-1`` and the recurrent kernel starts row b from |
| ``prefix_rec[prow[b]]``, so no per-row copy of a context exists. Per row: the |
| recurrent state after the row's own tokens (``rec``; read once ``fresh`` is False), |
| the conv tail, own K/V buffers [B, Hkv, cap, D] and ``own_len`` [B]. ``static`` |
| states (CUDA-graph buffers) are updated strictly in place.""" |
|
|
| prefix_rec: list[Tensor] |
| pk: list[Tensor] |
| pv: list[Tensor] |
| poff: Tensor |
| plen: Tensor |
| prow: Tensor |
| rec: list[Tensor | None] |
| tail: list[Tensor] |
| own_k: list[Tensor] |
| own_v: list[Tensor] |
| own_len: Tensor |
| fresh: bool = True |
| static: bool = False |
| |
| |
| wpoff: Tensor | None = None |
| wplen: Tensor | None = None |
| ple_rows: Tensor | None = None |
|
|
| @property |
| def rows(self) -> int: |
| return int(self.own_len.shape[0]) |
|
|
|
|
| _SMALL_GEMM = False |
|
|
|
|
| def _fop(name: str, x: Tensor): |
| """Fork-path ops, batch-invariant: row norms are fixed-config Triton kernels (one |
| program per row) and elementwise ops run eagerly. torch.compile is NOT used here: it |
| specializes size-1 dims, so one row alone and the same row in a batch would run |
| different generated kernels.""" |
| if not x.is_cuda: |
| return _RAW[name] |
| from deg.kernels import gated_rms_rows, rms_rows |
|
|
| if name == "rms": |
| return rms_rows |
| if name == "gated_rms": |
| return gated_rms_rows |
| if name == "rms_w": |
| from deg.kernels import rms_rows_w |
|
|
| return rms_rows_w |
| if name in ("rms_g", "rms_n"): |
| from deg.kernels import rms_rows_g |
|
|
| return rms_rows_g if name == "rms_g" else (lambda x, eps: rms_rows_g(x, None, eps)) |
| return _RAW[name] |
|
|
|
|
| def _fnorm(module: nn.Module, x: Tensor) -> Tensor: |
| return _fop("rms", x)(x, module.weight, module.eps) |
|
|
|
|
| def _bi(module: nn.Module, x: Tensor) -> Tensor: |
| """Batch-invariant projection. Each call site always uses ONE kernel config: decode |
| steps the small-M split kernel, every other fork pass the wide one.""" |
| from deg.kernels import bi_linear, bi_linear_small |
|
|
| if _SMALL_GEMM: |
| return bi_linear_small(x, module.weight, module.bias) |
| return bi_linear(x, module.weight, module.bias) |
|
|
|
|
| def _bi_fused(owner: nn.Module, names: tuple[str, ...], x: Tensor) -> tuple[Tensor, ...]: |
| """Several bias-free projections of the same input as ONE batch-invariant GEMM. The |
| weights are concatenated once and the original parameters become views into the |
| fused matrix (no extra memory; a later load_state_dict writes through). Each output |
| column is still its own K-loop, so fusing does not change any row's independence.""" |
| from deg.kernels import bi_linear, bi_linear_small |
|
|
| key = "_deg_fused_" + "_".join(names) |
| cached = owner.__dict__.get(key) |
| if cached is None: |
| mods = [getattr(owner, n) for n in names] |
| assert all(m.bias is None for m in mods) |
| w = torch.cat([m.weight.data for m in mods]).contiguous() |
| sizes, off = [], 0 |
| for m in mods: |
| n = m.weight.shape[0] |
| m.weight.data = w[off:off + n] |
| sizes.append(n) |
| off += n |
| cached = owner.__dict__[key] = (w, sizes) |
| w, sizes = cached |
| y = (bi_linear_small if _SMALL_GEMM else bi_linear)(x, w) |
| return tuple(torch.split(y, sizes, -1)) |
|
|
|
|
| def _fork_mlp(layer, h: Tensor, sel: Tensor | WriteRoute | None = None) -> Tensor: |
| m = _expert(layer, sel) |
| x = _fnorm(layer.post_attention_layernorm, h) |
| gate, up = _bi_fused(m, ("gate_proj", "up_proj"), x) |
| gate, up = _flora(m.gate_proj, x, gate, sel), _flora(m.up_proj, x, up, sel) |
| a = _fop("swiglu", x)(gate, up) |
| return h + _flora(m.down_proj, a, _bi(m.down_proj, a), sel) |
|
|
|
|
| def _fork_linear_layer(layer, h, valid, lengths, fs: ForkState, li: int, keep: bool, |
| sel: Tensor | None = None): |
| from deg.kernels import fork_delta_rule |
|
|
| attn = layer.linear_attn |
| B, L, _ = h.shape |
| xn = _fnorm(layer.input_layernorm, h) |
| mixed, z, pb, pa = _bi_fused(attn, ("in_proj_qkv", "in_proj_z", "in_proj_b", |
| "in_proj_a"), xn) |
| mixed = _flora(attn.in_proj_qkv, xn, mixed, sel) |
| z = _flora(attn.in_proj_z, xn, z, sel).reshape(B, L, -1, attn.head_v_dim) |
| beta, g = _fop("dn_gates", h)(pb, pa, attn.A_log, attn.dt_bias, valid) |
| tail = fs.tail[li] |
| ext = torch.cat([tail.to(mixed.dtype), mixed], 1) |
| K1 = tail.shape[1] |
| conv = _fop("conv_silu", h)(ext, attn.conv1d.weight[:, 0]) |
| if keep: |
| gather = lengths[:, None] + torch.arange(K1, device=h.device)[None] |
| new_tail = ext.gather(1, gather[..., None].expand(-1, -1, ext.shape[-1])) |
| if fs.static: |
| fs.tail[li].copy_(new_tail) |
| else: |
| fs.tail[li] = new_tail |
| qq, kk, vv = torch.split(conv, [attn.key_dim, attn.key_dim, attn.value_dim], dim=-1) |
| qq = qq.reshape(B, L, -1, attn.head_k_dim) |
| kk = kk.reshape(B, L, -1, attn.head_k_dim) |
| vv = vv.reshape(B, L, -1, attn.head_v_dim) |
| own = fs.rec[li] |
| if fs.fresh: |
| h0, h0_idx = fs.prefix_rec[li], fs.prow |
| else: |
| assert own is not None |
| h0, h0_idx = own, torch.arange(B, device=h.device) |
| store = None |
| if keep: |
| store = own if own is not None else torch.empty( |
| B, *h0.shape[1:], device=h.device, dtype=torch.float32) |
| core = fork_delta_rule(qq, kk, vv, g, beta.to(qq.dtype), h0, h0_idx, store) |
| if keep: |
| fs.rec[li] = store |
| core = _fop("gated_rms", h)(core, z, attn.norm.weight, attn.norm.variance_epsilon) |
| c2 = core.reshape(B, L, -1) |
| h = h + _flora(attn.out_proj, c2, _bi(attn.out_proj, c2), sel) |
| return _fork_mlp(layer, h, sel) |
|
|
|
|
| def _fork_attn_layer(layer, h, cos, sin, fs: ForkState, ai: int, sel: Tensor | None = None): |
| from transformers.models.qwen3_5.modeling_qwen3_5 import apply_rotary_pos_emb |
|
|
| from deg.kernels import fork_attention |
|
|
| a = layer.self_attn |
| B, L, _ = h.shape |
| xn = _fnorm(layer.input_layernorm, h) |
| pq, pk_, pv_ = _bi_fused(a, ("q_proj", "k_proj", "v_proj"), xn) |
| pq, pk_ = _flora(a.q_proj, xn, pq, sel), _flora(a.k_proj, xn, pk_, sel) |
| pv_ = _flora(a.v_proj, xn, pv_, sel) |
| q, gate = torch.chunk(pq.reshape(B, L, -1, a.head_dim * 2), 2, dim=-1) |
| q = _fnorm(a.q_norm, q).transpose(1, 2) |
| k = _fnorm(a.k_norm, pk_.reshape(B, L, -1, a.head_dim)).transpose(1, 2) |
| v = pv_.reshape(B, L, -1, a.head_dim).transpose(1, 2) |
| q, k = apply_rotary_pos_emb(q, k, cos, sin) |
| o = _fork_attend(fs, ai, q, k, v, a.scaling) |
| o = _fop("attn_gate", h)(o.reshape(B, L, -1), gate.reshape(B, L, -1)) |
| h = h + _flora(a.o_proj, o, _bi(a.o_proj, o), sel) |
| return _fork_mlp(layer, h, sel) |
|
|
|
|
| def _fork_attend(fs: ForkState, ai: int, q: Tensor, k: Tensor, v: Tensor, scaling: float |
| ) -> Tensor: |
| """Fork attention of attention layer ``ai``: q/k/v [B, H, L, D] (rotary applied); the |
| rows' new K/V are written into their own buffers at ``own_len``, then every query |
| attends its prefix (by reference) and its own causal keys -> [B, L, Hq, D].""" |
| from deg.kernels import fork_attention |
|
|
| L = q.shape[2] |
| ok, ov = fs.own_k[ai], fs.own_v[ai] |
| idx = (fs.own_len[:, None] + torch.arange(L, device=q.device)[None]) |
| idx = idx[:, None, :, None].expand(-1, k.shape[1], -1, k.shape[-1]) |
| ok.scatter_(2, idx, k.to(ok.dtype)) |
| ov.scatter_(2, idx, v.to(ov.dtype)) |
| return fork_attention(q.transpose(1, 2).contiguous(), fs.pk[ai], fs.pv[ai], ok, ov, |
| fs.own_len, fs.poff, fs.plen, scaling) |
|
|
|
|
| |
| |
| |
| |
| |
| |
| |
|
|
|
|
| def _lnorm(module: nn.Module, x: Tensor) -> Tensor: |
| return _op("rms_w", x)(x, module.weight, module.variance_epsilon) |
|
|
|
|
| def _lfnorm(module: nn.Module, x: Tensor) -> Tensor: |
| return _fop("rms_w", x)(x, module.weight, module.variance_epsilon) |
|
|
|
|
| def _rotate_half(x: Tensor) -> Tensor: |
| x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2:] |
| return torch.cat((-x2, x1), dim=-1) |
|
|
|
|
| def _rope_1d(q: Tensor, k: Tensor, cos: Tensor, sin: Tensor) -> tuple[Tensor, Tensor]: |
| """HF ``apply_rotary_pos_emb`` (full rotary) for q/k [B, H, L, D], cos/sin [B, L, D].""" |
| cos, sin = cos[:, None], sin[:, None] |
| return q * cos + _rotate_half(q) * sin, k * cos + _rotate_half(k) * sin |
|
|
|
|
| def _lfm_mlp(layer, h: Tensor, sel: Tensor | None = None) -> Tensor: |
| m = layer.feed_forward |
| x = _lnorm(layer.ffn_norm, h) |
| a = _op("swiglu", x)(_lin(m.w1, x, sel), _lin(m.w3, x, sel)) |
| return h + _lin(m.w2, a, sel) |
|
|
|
|
| def _lfm_conv_layer(layer, h, valid, lengths_new, rec, tail, rows=None, sel=None): |
| """Trunk/training path of a short-conv layer; same contract as :func:`_linear_layer` |
| (``rec`` is the empty placeholder and passes through).""" |
| cv = layer.conv |
| if rows is not None: |
| rec, tail = rec.index_select(0, rows), tail.index_select(0, rows) |
| xn = _lnorm(layer.operator_norm, h) |
| b, c, x = _lin(cv.in_proj, xn, sel).chunk(3, -1) |
| bx = b * x |
| ext = torch.cat([tail.to(bx.dtype), bx], 1) |
| K1 = tail.shape[1] |
| conv = _op("short_conv", h)(ext, cv.conv.weight[:, 0]) |
| gather = lengths_new[:, None] + torch.arange(K1, device=h.device)[None] |
| new_tail = ext.gather(1, gather[..., None].expand(-1, -1, ext.shape[-1])) |
| h = h + _lin(cv.out_proj, c * conv, sel) |
| return _lfm_mlp(layer, h, sel), rec, new_tail |
|
|
|
|
| def _lfm_attn_qkv(layer, h, cos, sin, sel=None): |
| a = layer.self_attn |
| B, L, _ = h.shape |
| xn = _lnorm(layer.operator_norm, h) |
| q = _lnorm(a.q_layernorm, _lin(a.q_proj, xn, sel).view(B, L, -1, a.head_dim)) |
| k = _lnorm(a.k_layernorm, _lin(a.k_proj, xn, sel).view(B, L, -1, a.head_dim)) |
| v = _lin(a.v_proj, xn, sel).view(B, L, -1, a.head_dim).transpose(1, 2) |
| q, k = _rope_1d(q.transpose(1, 2), k.transpose(1, 2), cos, sin) |
| return q, k, v, None |
|
|
|
|
| def _lfm_attn_out(layer, h, o, gate, sel=None): |
| return _lfm_mlp(layer, h + _lin(layer.self_attn.out_proj, o, sel), sel) |
|
|
|
|
| def _lfm_fork_mlp(layer, h: Tensor, sel: Tensor | None = None) -> Tensor: |
| m = layer.feed_forward |
| x = _lfnorm(layer.ffn_norm, h) |
| g, u = _bi_fused(m, ("w1", "w3"), x) |
| g, u = _flora(m.w1, x, g, sel), _flora(m.w3, x, u, sel) |
| a = _fop("swiglu", x)(g, u) |
| return h + _flora(m.w2, a, _bi(m.w2, a), sel) |
|
|
|
|
| def _lfm_fork_conv_layer(layer, h, valid, lengths, fs: ForkState, li: int, keep: bool, |
| sel: Tensor | None = None): |
| """Fork path of a short-conv layer: row b continues its own conv tail (copied per row |
| by :meth:`Backbone.fork`, 2 tokens); everything is per row (bi_linear GEMMs, |
| elementwise conv), so batch-invariant.""" |
| cv = layer.conv |
| B = h.shape[0] |
| xn = _lfnorm(layer.operator_norm, h) |
| b, c, x = _flora(cv.in_proj, xn, _bi(cv.in_proj, xn), sel).chunk(3, -1) |
| bx = b * x |
| tail = fs.tail[li] |
| ext = torch.cat([tail.to(bx.dtype), bx], 1) |
| K1 = tail.shape[1] |
| conv = _fop("short_conv", h)(ext, cv.conv.weight[:, 0]) |
| if keep: |
| gather = lengths[:, None] + torch.arange(K1, device=h.device)[None] |
| new_tail = ext.gather(1, gather[..., None].expand(-1, -1, ext.shape[-1])) |
| if fs.static: |
| fs.tail[li].copy_(new_tail) |
| else: |
| fs.tail[li] = new_tail |
| fs.rec[li] = h.new_zeros(B, 0, dtype=torch.float32) |
| y = c * conv |
| h = h + _flora(cv.out_proj, y, _bi(cv.out_proj, y), sel) |
| return _lfm_fork_mlp(layer, h, sel) |
|
|
|
|
| def _lfm_fork_attn_layer(layer, h, cos, sin, fs: ForkState, ai: int, |
| sel: Tensor | None = None): |
| a = layer.self_attn |
| B, L, _ = h.shape |
| xn = _lfnorm(layer.operator_norm, h) |
| pq, pk_, pv_ = _bi_fused(a, ("q_proj", "k_proj", "v_proj"), xn) |
| pq, pk_ = _flora(a.q_proj, xn, pq, sel), _flora(a.k_proj, xn, pk_, sel) |
| pv_ = _flora(a.v_proj, xn, pv_, sel) |
| q = _lfnorm(a.q_layernorm, pq.reshape(B, L, -1, a.head_dim)).transpose(1, 2) |
| k = _lfnorm(a.k_layernorm, pk_.reshape(B, L, -1, a.head_dim)).transpose(1, 2) |
| v = pv_.reshape(B, L, -1, a.head_dim).transpose(1, 2) |
| q, k = _rope_1d(q, k, cos, sin) |
| o = _fork_attend(fs, ai, q, k, v, a.scaling).reshape(B, L, -1) |
| h = h + _flora(a.out_proj, o, _bi(a.out_proj, o), sel) |
| return _lfm_fork_mlp(layer, h, sel) |
|
|
|
|
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
|
|
| def _gnorm(module: nn.Module, x: Tensor) -> Tensor: |
| """Trunk/training-path Gemma norm, eager (not torch.compile'd: the compiled backward of |
| this norm returned NaN weight gradients on long checkpointed reads on A100, torch |
| 2.14 + dynamic shapes; eager is exact and the norms are a small share of the step).""" |
| if getattr(module, "with_scale", True): |
| return _rms_g(x, module.weight, module.eps) |
| return _rms_n(x, module.eps) |
|
|
|
|
| def _gfnorm(module: nn.Module, x: Tensor) -> Tensor: |
| if getattr(module, "with_scale", True): |
| return _fop("rms_g", x)(x, module.weight, module.eps) |
| return _fop("rms_n", x)(x, module.eps) |
|
|
|
|
| def _rope_g(x: Tensor, cos: Tensor, sin: Tensor) -> Tensor: |
| """HF ``apply_rotary_pos_emb(x, cos, sin, unsqueeze_dim=2)``: x [B, L, H, D].""" |
| cos, sin = cos[:, :, None], sin[:, :, None] |
| return x * cos + _rotate_half(x) * sin |
|
|
|
|
| def _cat_kv(parts: Tensor | tuple[Tensor, ...], new: Tensor) -> Tensor: |
| if isinstance(parts, Tensor): |
| parts = (parts,) |
| return torch.cat([*(t.to(new.dtype) for t in parts), new], 2) |
|
|
|
|
| def _gemma_mask(mask: Tensor | tuple[Tensor, int], c: int) -> Tensor: |
| """A global layer's attention mask: given, or built from ``(kv_mask, s)`` inside the |
| checkpoint (chunked training read: incoming prefix mask, ``s`` earlier tokens of this |
| call, causal chunk) so no [B, c, P+s+c] mask is saved.""" |
| if isinstance(mask, Tensor): |
| return mask |
| kv_mask, s = mask |
| B = kv_mask.shape[0] |
| causal = torch.ones(c, c, dtype=torch.bool, device=kv_mask.device).tril() |
| return torch.cat([kv_mask[:, None, :].expand(-1, c, -1), |
| torch.ones(B, c, s, dtype=torch.bool, device=kv_mask.device), |
| causal[None].expand(B, -1, -1)], 2) |
|
|
|
|
| class _RowSelect(torch.autograd.Function): |
| """``x.index_select(0, rows)`` whose backward sums the rows' gradients with one GEMM |
| (``onehot(rows)^T @ grad``, fp32 accumulation) instead of an atomic scatter-add: fork |
| rows re-select their trunk's prefix K/V in every layer, and the scatter-add backward of |
| [rows, Hkv, P, D] was ~17% of a long-context training step on A100.""" |
|
|
| @staticmethod |
| def forward(ctx, x: Tensor, rows: Tensor) -> Tensor: |
| ctx.save_for_backward(rows) |
| ctx.n = x.shape[0] |
| return x.index_select(0, rows) |
|
|
| @staticmethod |
| def backward(ctx, grad: Tensor): |
| (rows,) = ctx.saved_tensors |
| onehot = F.one_hot(rows, ctx.n).to(grad.dtype).T |
| out = (onehot @ grad.reshape(grad.shape[0], -1)).reshape(ctx.n, *grad.shape[1:]) |
| return out, None |
|
|
|
|
| def _select_rows(x: Tensor, rows: Tensor) -> Tensor: |
| if torch.is_grad_enabled() and x.requires_grad: |
| return _RowSelect.apply(x, rows) |
| return x.index_select(0, rows) |
|
|
|
|
| def _gemma_sdpa(q: Tensor, k: Tensor, v: Tensor, mask: Tensor, scale: float, |
| window: int = 0) -> Tensor: |
| """SDPA of q [B, H, L, D] over keys [B, Hkv, K, D] with a bool mask [B, L, K]. |
| |
| ``window`` > 0 (sliding layer; keys = the T = window-1 tail then the L new keys, as |
| :meth:`GemmaBackbone._window_mask` lays them out): queries run in blocks of |
| ``window``, each against only the keys its window can reach (T + window), so a long |
| read costs O(L * 2 window) instead of O(L * (T + L)). Same math: every dropped key is |
| masked for every query of the block.""" |
| B, H, L, D = q.shape |
| T = k.shape[2] - L |
| if window <= 0 or L <= window or T != window - 1: |
| return F.scaled_dot_product_attention(q, k, v, attn_mask=mask[:, None], scale=scale, |
| enable_gqa=True) |
| W = window |
| nb = -(-L // W) |
| pad = nb * W - L |
| if pad: |
| q = F.pad(q, (0, 0, 0, pad)) |
| k, v = F.pad(k, (0, 0, 0, pad)), F.pad(v, (0, 0, 0, pad)) |
| mask = F.pad(mask, (0, pad, 0, pad)) |
| mask[:, L:, 0] = True |
| hkv = k.shape[1] |
| |
| kb = k.unfold(2, W + T, W).permute(0, 2, 1, 4, 3).reshape(B * nb, hkv, W + T, D) |
| vb = v.unfold(2, W + T, W).permute(0, 2, 1, 4, 3).reshape(B * nb, hkv, W + T, D) |
| qb = q.reshape(B, H, nb, W, D).transpose(1, 2).reshape(B * nb, H, W, D) |
| mc = mask.contiguous() |
| s0, s1, s2 = mc.stride() |
| mb = mc.as_strided((B, nb, W, W + T), (s0, W * (s1 + s2), s1, s2)).reshape(B * nb, W, W + T) |
| o = F.scaled_dot_product_attention(qb, kb, vb, attn_mask=mb[:, None], scale=scale, |
| enable_gqa=True) |
| o = o.reshape(B, nb, H, W, D).transpose(1, 2).reshape(B, H, nb * W, D) |
| return o[:, :, :L] |
|
|
|
|
| def _gemma_rest(layer, h: Tensor, attn: Tensor, ple: Tensor | None, sel=None) -> Tensor: |
| """Everything after the attention output projection: post-attention norm + residual, |
| GeGLU MLP, the per-layer input, the layer scalar (trunk/training path).""" |
| h = h + _gnorm(layer.post_attention_layernorm, attn) |
| m = _expert(layer, sel) |
| x = _gnorm(layer.pre_feedforward_layernorm, h) |
| y = _lin(m.down_proj, _op("geglu", x)(_lin(m.gate_proj, x, sel), _lin(m.up_proj, x, sel)), |
| sel) |
| h = h + _gnorm(layer.post_feedforward_layernorm, y) |
| if ple is not None: |
| g = _op("geglu", h)(layer.per_layer_input_gate(h), ple) |
| h = h + _gnorm(layer.post_per_layer_input_norm, layer.per_layer_projection(g)) |
| return h * layer.layer_scalar |
|
|
|
|
| def _gemma_layer(layer, h, ple, cos, sin, kp, vp, mask, rows=None, sel=None, kn=None, |
| vn=None, window=0): |
| """One Gemma 4 layer (trunk / training path). ``kp``/``vp`` [B', Hkv, P, D] (or a tuple |
| of parts concatenated here): the incoming keys of this layer's K/V slot (a global |
| layer's whole prefix, a sliding layer's tail); ``rows``: as in :func:`_linear_layer` |
| (prefix rows selected inside the checkpoint). A K/V-sharing layer gets the source |
| layer's new K/V of this call as ``kn``/``vn`` and computes only q. ``mask`` [B, L, P+L] |
| (or ``(kv_mask, s)``, :func:`_gemma_mask`); ``window``: a sliding layer's window (its |
| keys = tail + new keys, :func:`_gemma_sdpa` blocks long reads). Returns (h, k, v of |
| this call's tokens).""" |
| a = layer.self_attn |
| B, L, _ = h.shape |
| if rows is not None: |
| kp = tuple(_select_rows(t, rows) for t in (kp if isinstance(kp, tuple) else (kp,))) |
| vp = tuple(_select_rows(t, rows) for t in (vp if isinstance(vp, tuple) else (vp,))) |
| xn = _gnorm(layer.input_layernorm, h) |
| q = _gnorm(a.q_norm, _lin(a.q_proj, xn, sel).view(B, L, -1, a.head_dim)) |
| q = _rope_g(q, cos, sin).transpose(1, 2) |
| if kn is None: |
| k = _gnorm(a.k_norm, _lin(a.k_proj, xn, sel).view(B, L, -1, a.head_dim)) |
| k = _rope_g(k, cos, sin).transpose(1, 2) |
| v = _gnorm(a.v_norm, _lin(a.v_proj, xn, sel).view(B, L, -1, a.head_dim)).transpose(1, 2) |
| else: |
| k, v = kn, vn |
| keys, values = _cat_kv(kp, k), _cat_kv(vp, v) |
| o = _gemma_sdpa(q, keys, values, _gemma_mask(mask, L), a.scaling, window) |
| o = o.transpose(1, 2).reshape(B, L, -1) |
| return _gemma_rest(layer, h, _lin(a.o_proj, o, sel), ple, sel), k, v |
|
|
|
|
| def _gemma_fork_layer(layer, h, ple, cos, sin, fs: ForkState, slot: int, owner: bool, |
| window: int, sel=None): |
| """Fork path of one Gemma 4 layer (batch-invariant, like :func:`_fork_attn_layer`). An |
| owner layer writes the rows' new K/V into slot ``slot``'s own buffers; a K/V-sharing |
| layer reads them (written by its source layer earlier in the same pass). A sliding |
| layer (``window`` > 0) attends its prefix tail (``fs.wpoff``/``fs.wplen``) and its own |
| keys within the window; a global layer the whole prefix by reference.""" |
| from deg.kernels import fork_attention |
|
|
| a = layer.self_attn |
| B, L, _ = h.shape |
| xn = _gfnorm(layer.input_layernorm, h) |
| if owner: |
| pq, pk_, pv_ = _bi_fused(a, ("q_proj", "k_proj", "v_proj"), xn) |
| pk_, pv_ = _flora(a.k_proj, xn, pk_, sel), _flora(a.v_proj, xn, pv_, sel) |
| else: |
| pq = _bi(a.q_proj, xn) |
| pq = _flora(a.q_proj, xn, pq, sel) |
| q = _rope_g(_gfnorm(a.q_norm, pq.reshape(B, L, -1, a.head_dim)), cos, sin) |
| ok, ov = fs.own_k[slot], fs.own_v[slot] |
| if owner: |
| k = _rope_g(_gfnorm(a.k_norm, pk_.reshape(B, L, -1, a.head_dim)), cos, sin) |
| v = _gfnorm(a.v_norm, pv_.reshape(B, L, -1, a.head_dim)) |
| idx = (fs.own_len[:, None] + torch.arange(L, device=h.device)[None]) |
| idx = idx[:, None, :, None].expand(-1, k.shape[2], -1, k.shape[-1]) |
| ok.scatter_(2, idx, k.transpose(1, 2).to(ok.dtype)) |
| ov.scatter_(2, idx, v.transpose(1, 2).to(ov.dtype)) |
| if window > 0: |
| assert fs.wpoff is not None and fs.wplen is not None |
| poff, plen = fs.wpoff, fs.wplen |
| else: |
| poff, plen = fs.poff, fs.plen |
| o = fork_attention(q.contiguous(), fs.pk[slot], fs.pv[slot], ok, ov, fs.own_len, poff, |
| plen, a.scaling, window).reshape(B, L, -1) |
| h = h + _gfnorm(layer.post_attention_layernorm, _flora(a.o_proj, o, _bi(a.o_proj, o), sel)) |
| m = _expert(layer, sel) |
| x = _gfnorm(layer.pre_feedforward_layernorm, h) |
| gate, up = _bi_fused(m, ("gate_proj", "up_proj"), x) |
| gate, up = _flora(m.gate_proj, x, gate, sel), _flora(m.up_proj, x, up, sel) |
| y = _fop("geglu", x)(gate, up) |
| h = h + _gfnorm(layer.post_feedforward_layernorm, _flora(m.down_proj, y, _bi(m.down_proj, y), |
| sel)) |
| if ple is not None: |
| g = _fop("geglu", h)(_bi(layer.per_layer_input_gate, h), ple) |
| h = h + _gfnorm(layer.post_per_layer_input_norm, _bi(layer.per_layer_projection, g)) |
| return h * layer.layer_scalar |
|
|
|
|
| |
|
|
|
|
| @dataclass(frozen=True) |
| class Family: |
| """The layer functions of one architecture behind :class:`Backbone` (SPEC.md §12.8). |
| Every family is a hybrid of state-carrying layers (``rec_type``) and attention |
| layers; :class:`State` / :class:`ForkState` hold, per state layer, (recurrent state, |
| conv tail) and, per attention layer, K/V, so the trunk/fork/chunk/decode machinery is |
| shared and only these functions differ. |
| |
| * ``rec_layer(layer, h, valid, lengths, rec, tail, rows, sel) -> (h, rec, tail)`` |
| * ``attn_qkv(layer, h, cos, sin, sel) -> (q, k, v [B, H, L, D] rotary applied, gate)`` |
| * ``attn_out(layer, h, o [B, L, Hq*D], gate, sel) -> h`` (residual + MLP) |
| * ``fork_rec_layer(layer, h, valid, lengths, fs, li, keep, sel) -> h`` and |
| ``fork_attn_layer(layer, h, cos, sin, fs, ai, sel) -> h``: batch-invariant fork path |
| * ``rec_shapes(layer) -> (recurrent state shape, conv tail shape)`` per row |
| * ``final_norm``: attribute of the text model; ``fork_norm(module, h)`` its |
| batch-invariant form; ``mrope``: rotary takes [3, B, L] positions.""" |
|
|
| name: str |
| rec_type: str |
| rec_layer: Any |
| attn_qkv: Any |
| attn_out: Any |
| fork_rec_layer: Any |
| fork_attn_layer: Any |
| rec_shapes: Any |
| final_norm: str |
| fork_norm: Any |
| mrope: bool |
| lora_targets: dict[str, tuple[str, ...]] |
|
|
|
|
| def _qwen_rec_shapes(layer) -> tuple[tuple[int, ...], tuple[int, ...]]: |
| at = layer.linear_attn |
| return ((at.num_v_heads, at.head_k_dim, at.head_v_dim), |
| (at.conv_kernel_size - 1, at.conv_dim)) |
|
|
|
|
| def _lfm_rec_shapes(layer) -> tuple[tuple[int, ...], tuple[int, ...]]: |
| cv = layer.conv |
| return (0,), (cv.conv_kernel_size - 1, cv.conv.weight.shape[0]) |
|
|
|
|
| QWEN35 = Family("qwen3_5", "linear_attention", _linear_layer, _attn_qkv, _attn_out, |
| _fork_linear_layer, _fork_attn_layer, _qwen_rec_shapes, "norm", _fnorm, True, |
| LORA_TARGETS) |
| LFM2 = Family("lfm2", "conv", _lfm_conv_layer, _lfm_attn_qkv, _lfm_attn_out, |
| _lfm_fork_conv_layer, _lfm_fork_attn_layer, _lfm_rec_shapes, "embedding_norm", |
| _lfnorm, False, |
| {"conv": ("in_proj", "out_proj"), |
| "self_attn": ("q_proj", "k_proj", "v_proj", "out_proj"), |
| "feed_forward": ("w1", "w3", "w2")}) |
|
|
|
|
| |
| |
| GEMMA4 = Family("gemma4", "", None, None, None, None, None, None, "norm", _gfnorm, False, |
| {"self_attn": ("q_proj", "k_proj", "v_proj", "o_proj"), |
| "mlp": ("gate_proj", "up_proj", "down_proj")}) |
|
|
|
|
| def family_of(lm: nn.Module) -> Family: |
| kind = lm.config.model_type |
| if kind == "lfm2": |
| return LFM2 |
| if kind == "gemma4_text": |
| return GEMMA4 |
| if kind.startswith("qwen3_5"): |
| return QWEN35 |
| raise ValueError(f"no deg backbone family for model_type {kind!r}") |
|
|
|
|
| class Backbone(nn.Module): |
| """Functional hybrid text stack (Qwen3.5 or LFM2.5, :class:`Family`) over explicit |
| prefix state (no HF cache objects, so gradients flow from forks into the shared |
| trunk). ``Backbone(lm)`` of a Gemma 4 model is a :class:`GemmaBackbone`.""" |
|
|
| def __new__(cls, lm: nn.Module | None = None, *args: Any, **kwargs: Any) -> Backbone: |
| if cls is Backbone and lm is not None and family_of(lm) is GEMMA4: |
| return super().__new__(GemmaBackbone) |
| return super().__new__(cls) |
|
|
| def __init__(self, lm: nn.Module) -> None: |
| super().__init__() |
| self.lm = lm |
| text = lm.model |
| self.text = getattr(text, "language_model", text) |
| self.config = self.text.config |
| self.fam = family_of(lm) |
| |
| self.types = ["rec" if t == self.fam.rec_type else "attn" |
| for t in self.config.layer_types] |
| self.grad_ckpt = False |
|
|
| @property |
| def hidden(self) -> int: |
| return self.config.hidden_size |
|
|
| def rope(self, h: Tensor, pos: Tensor) -> tuple[Tensor, Tensor]: |
| """cos/sin for positions ``pos`` [B, L].""" |
| return self.text.rotary_emb(h, pos[None].expand(3, -1, -1) if self.fam.mrope else pos) |
|
|
| def final_norm(self, h: Tensor) -> Tensor: |
| return getattr(self.text, self.fam.final_norm)(h) |
|
|
| def fork_final_norm(self, h: Tensor) -> Tensor: |
| return self.fam.fork_norm(getattr(self.text, self.fam.final_norm), h) |
|
|
| def rec_buffers(self, layer, batch: int, device, dtype) -> tuple[Tensor, Tensor]: |
| """Zero (recurrent state fp32, conv tail) of one state layer for ``batch`` rows.""" |
| rec, tail = self.fam.rec_shapes(layer) |
| return (torch.zeros(batch, *rec, device=device, dtype=torch.float32), |
| torch.zeros(batch, *tail, device=device, dtype=dtype)) |
|
|
| def empty_state(self, batch: int, device, dtype) -> State: |
| lin, kv = [], [] |
| for layer, kind in zip(self.text.layers, self.types): |
| if kind == "rec": |
| lin.append(self.rec_buffers(layer, batch, device, dtype)) |
| else: |
| a = layer.self_attn |
| shape = (batch, self.config.num_key_value_heads, 0, a.head_dim) |
| kv.append((torch.zeros(shape, device=device, dtype=dtype), |
| torch.zeros(shape, device=device, dtype=dtype))) |
| return State(lin, kv, torch.zeros(batch, 0, dtype=torch.bool, device=device), |
| torch.zeros(batch, dtype=torch.long, device=device)) |
|
|
| def forward(self, ids: Tensor, lengths: Tensor, state: State | None = None, |
| consume: bool = False, rows: Tensor | None = None, sel: Tensor | None = None |
| ) -> tuple[Tensor, State | None]: |
| """ids [B,L] right-padded, lengths [B] real tokens -> (final-norm hidden [B,L,C], |
| state after each row's real tokens). Inputs longer than :data:`DENSE_PREFILL_MAX` |
| stream through :meth:`_forward_chunked` (no O(L^2) tensor, no length cap). |
| |
| ``consume``: the caller gives up ``state``; each attention layer's prefix K/V is |
| released (``state.kv[i] = None``) as soon as its extended copy exists, so |
| extending a long context peaks at its K/V plus one layer, not twice its K/V. |
| |
| ``rows`` (training forks): row b continues row ``rows[b]`` of ``state`` (the |
| un-selected trunk state). Each layer selects its prefix rows inside the gradient |
| checkpoint, so the saved activations hold no per-row copy of the prefix (K/V of |
| every attention layer, recurrent state of every linear layer): memory per fork |
| row is independent of the context length except during one layer's recompute. |
| Returns ``(hidden, None)``: no continued state is built. |
| |
| ``sel`` [B, n*r]: per-row LoRA adapter selector (None = adapters off).""" |
| if rows is not None and ids.shape[1] > DENSE_PREFILL_MAX: |
| assert state is not None |
| state, rows = state.select(rows), None |
| if ids.shape[1] > DENSE_PREFILL_MAX: |
| if self.grad_ckpt and torch.is_grad_enabled(): |
| return self._forward_chunked_ckpt(ids, lengths, state, sel) |
| return self._forward_chunked(ids, lengths, state, consume, sel) |
| B, L = ids.shape |
| device = ids.device |
| if profile._ACTIVE.get() is not None: |
| profile.count(forwards=1, forward_rows=B, forward_tokens=B * L, |
| forward_real_tokens=int(lengths.sum())) |
| h = self.text.embed_tokens(ids) |
| if state is None: |
| state = self.empty_state(B, device, h.dtype) |
| kv_mask, plens = state.kv_mask, state.lengths |
| if rows is not None: |
| kv_mask, plens = kv_mask.index_select(0, rows), plens.index_select(0, rows) |
| valid = torch.arange(L, device=device)[None] < lengths[:, None] |
| pos = plens[:, None] + torch.arange(L, device=device)[None] |
| cos, sin = self.rope(h, pos) |
| causal = torch.ones(L, L, dtype=torch.bool, device=device).tril() |
| mask = torch.cat([kv_mask[:, None, :].expand(-1, L, -1), |
| causal[None].expand(B, -1, -1)], 2) |
| validf = valid.to(torch.float32) |
| lin_out, kv_out = [], [] |
| li = ai = 0 |
| for layer, kind in zip(self.text.layers, self.types): |
| if kind == "rec": |
| rec, tail = state.lin[li] |
| args = (layer, h, validf, lengths, rec, tail, rows, sel) |
| if self.grad_ckpt and torch.is_grad_enabled(): |
| h, rec, tail = checkpoint(self.fam.rec_layer, *args, use_reentrant=False) |
| else: |
| h, rec, tail = self.fam.rec_layer(*args) |
| if rows is None: |
| lin_out.append((rec, tail)) |
| del rec, tail |
| li += 1 |
| else: |
| kp, vp = state.kv[ai] |
| args = (self.fam, layer, h, cos, sin, kp, vp, mask, rows, sel) |
| if self.grad_ckpt and torch.is_grad_enabled(): |
| h, k, v = checkpoint(_attn_layer, *args, use_reentrant=False) |
| else: |
| h, k, v = _attn_layer(*args) |
| if rows is None: |
| kv_out.append((torch.cat([kp.to(k.dtype), k], 2), |
| torch.cat([vp.to(v.dtype), v], 2))) |
| del k, v |
| if consume: |
| del kp, vp, args |
| state.kv[ai] = None |
| ai += 1 |
| h = self.final_norm(h) |
| if rows is not None: |
| return h, None |
| new = State(lin_out, kv_out, torch.cat([state.kv_mask, valid], 1), |
| state.lengths + lengths) |
| return h, new |
|
|
| def _forward_chunked(self, ids: Tensor, lengths: Tensor, state: State | None, |
| consume: bool = False, sel: Tensor | None = None, |
| soft: tuple[Tensor, Tensor] | None = None) -> tuple[Tensor, State]: |
| """Chunked prefill, same contract as :meth:`forward`: the input is read |
| :data:`PREFILL_CHUNK` tokens at a time, carrying each linear layer's recurrent |
| state and conv tail and appending each attention layer's K/V. Rows are |
| right-padded, so a real query's earlier positions in this call are all real: |
| attention needs only the incoming prefix mask plus causality inside the chunk. |
| |
| Peak memory is O(L) (K/V, hidden) plus O(chunk) work. On CUDA without autograd, |
| attention runs the flash-style fork kernel (prefix K/V read in place from one |
| preallocated buffer per layer, causal own keys) whenever the incoming prefix is |
| unpadded; otherwise (CPU, autograd, padded incoming prefix) SDPA with an explicit |
| [B, chunk, P+chunk] mask: O(chunk * (P+chunk)), never L x L.""" |
| from deg.kernels import fork_attention |
|
|
| B, L = ids.shape |
| device = ids.device |
| C = PREFILL_CHUNK |
| if profile._ACTIVE.get() is not None: |
| profile.count(forwards=1, forward_rows=B, forward_tokens=B * L, |
| forward_real_tokens=int(lengths.sum())) |
| if state is None: |
| state = self.empty_state(B, device, self.text.embed_tokens.weight.dtype) |
| P0 = state.kv_mask.shape[1] |
| fast = (ids.is_cuda and not torch.is_grad_enabled() |
| and bool(state.kv_mask.all())) |
| lin = list(state.lin) |
| if fast: |
| bufs = [] |
| for i in range(len(state.kv)): |
| kp, vp = state.kv[i] |
| shape = (kp.shape[1], B, P0 + L, kp.shape[3]) |
| kb = torch.empty(shape, device=device, dtype=self.text.embed_tokens.weight.dtype) |
| vb = torch.empty_like(kb) |
| kb[:, :, :P0] = kp.transpose(0, 1) |
| vb[:, :, :P0] = vp.transpose(0, 1) |
| bufs.append((kb, vb)) |
| if consume: |
| del kp, vp |
| state.kv[i] = None |
| poff = torch.arange(B, device=device) * (P0 + L) |
| qstart = torch.zeros(B, dtype=torch.long, device=device) |
| else: |
| kv = list(state.kv) |
| outs = [] |
| for s in range(0, L, C): |
| e = min(L, s + C) |
| c = e - s |
| lens = (lengths - s).clamp(0, c) |
| h = self.text.embed_tokens(ids[:, s:e]) |
| valid = torch.arange(c, device=device)[None] < lens[:, None] |
| pos = state.lengths[:, None] + s + torch.arange(c, device=device)[None] |
| cos, sin = self.rope(h, pos) |
| validf = valid.to(torch.float32) |
| if not fast: |
| causal = torch.ones(c, c, dtype=torch.bool, device=device).tril() |
| mask = torch.cat([state.kv_mask[:, None, :].expand(-1, c, -1), |
| torch.ones(B, c, s, dtype=torch.bool, device=device), |
| causal[None].expand(B, -1, -1)], 2) |
| li = ai = 0 |
| for layer, kind in zip(self.text.layers, self.types): |
| if kind == "rec": |
| rec, tail = lin[li] |
| h, rec, tail = self.fam.rec_layer(layer, h, validf, lens, rec, tail, None, |
| sel) |
| lin[li] = (rec, tail) |
| li += 1 |
| continue |
| a = layer.self_attn |
| q, k, v, gate = self.fam.attn_qkv(layer, h, cos, sin, sel) |
| if fast: |
| kb, vb = bufs[ai] |
| o = fork_attention(q.transpose(1, 2).contiguous(), |
| kb.view(kb.shape[0], -1, kb.shape[3]), |
| vb.view(vb.shape[0], -1, vb.shape[3]), |
| k.contiguous(), v.contiguous(), qstart, poff, |
| torch.full((B,), P0 + s, device=device), a.scaling) |
| kb[:, :, P0 + s:P0 + e] = k.transpose(0, 1) |
| vb[:, :, P0 + s:P0 + e] = v.transpose(0, 1) |
| else: |
| kp, vp = kv[ai] |
| keys = torch.cat([kp.to(k.dtype), k], 2) |
| values = torch.cat([vp.to(v.dtype), v], 2) |
| o = F.scaled_dot_product_attention(q, keys, values, attn_mask=mask[:, None], |
| scale=a.scaling, enable_gqa=True) |
| o = o.transpose(1, 2) |
| kv[ai] = (keys, values) |
| if consume and s == 0: |
| state.kv[ai] = None |
| h = self.fam.attn_out(layer, h, o.reshape(B, c, -1), gate, sel) |
| ai += 1 |
| outs.append(self.final_norm(h)) |
| if fast: |
| kv = [(kb.transpose(0, 1), vb.transpose(0, 1)) for kb, vb in bufs] |
| valid = torch.arange(L, device=device)[None] < lengths[:, None] |
| new = State(lin, kv, torch.cat([state.kv_mask, valid], 1), state.lengths + lengths) |
| return torch.cat(outs, 1), new |
|
|
| def _forward_chunked_ckpt(self, ids: Tensor, lengths: Tensor, state: State | None, |
| sel: Tensor | None = None) -> tuple[Tensor, State]: |
| """Training read of an input longer than :data:`DENSE_PREFILL_MAX` (autograd + |
| ``grad_ckpt``): :data:`PREFILL_CHUNK` tokens at a time like :meth:`_forward_chunked`, |
| every layer of every chunk gradient-checkpointed. A checkpoint saves only its |
| inputs (the chunk's hidden, the carried recurrent state / conv tail, and the |
| earlier chunks' K/V as a tuple of the tensors that exist anyway); the attention |
| mask is built inside the checkpoint, so nothing O(L^2) is saved and one layer's |
| recompute is O(chunk * L). The returned state's K/V concatenates the chunks.""" |
| B, L = ids.shape |
| device = ids.device |
| C = PREFILL_CHUNK |
| if state is None: |
| state = self.empty_state(B, device, self.text.embed_tokens.weight.dtype) |
| lin = list(state.lin) |
| kparts: list[list[Tensor]] = [[kp] for kp, _ in state.kv] |
| vparts: list[list[Tensor]] = [[vp] for _, vp in state.kv] |
| outs = [] |
| for s in range(0, L, C): |
| e = min(L, s + C) |
| c = e - s |
| lens = (lengths - s).clamp(0, c) |
| h = self.text.embed_tokens(ids[:, s:e]) |
| valid = torch.arange(c, device=device)[None] < lens[:, None] |
| pos = state.lengths[:, None] + s + torch.arange(c, device=device)[None] |
| cos, sin = self.rope(h, pos) |
| validf = valid.to(torch.float32) |
| li = ai = 0 |
| for layer, kind in zip(self.text.layers, self.types): |
| if kind == "rec": |
| rec, tail = lin[li] |
| h, rec, tail = checkpoint(self.fam.rec_layer, layer, h, validf, lens, rec, |
| tail, None, sel, use_reentrant=False) |
| lin[li] = (rec, tail) |
| li += 1 |
| continue |
| h, k, v = checkpoint(_chunk_attn_layer, self.fam, layer, h, cos, sin, |
| tuple(kparts[ai]), tuple(vparts[ai]), state.kv_mask, s, sel, |
| use_reentrant=False) |
| kparts[ai].append(k) |
| vparts[ai].append(v) |
| ai += 1 |
| outs.append(self.final_norm(h)) |
| kv = [(torch.cat(kp, 2), torch.cat(vp, 2)) for kp, vp in zip(kparts, vparts)] |
| valid = torch.arange(L, device=device)[None] < lengths[:, None] |
| new = State(lin, kv, torch.cat([state.kv_mask, valid], 1), state.lengths + lengths) |
| return torch.cat(outs, 1), new |
|
|
| def fork(self, prefixes: State | list[State], prow: list[int] | int, cap: int |
| ) -> ForkState: |
| """Fresh fork rows; row b continues ``prefixes[prow[b]]`` (batch-1, unpadded |
| states). ``prow`` as an int = that many rows of a single prefix. Own K/V capacity |
| ``cap``. A single prefix is read in place; several are concatenated once.""" |
| if isinstance(prefixes, State): |
| prefixes = [prefixes] |
| if isinstance(prow, int): |
| prow = [0] * prow |
| for p in prefixes: |
| assert p.lengths.shape[0] == 1 and bool(p.kv_mask.all()) |
| rows = len(prow) |
| device = prefixes[0].lengths.device |
| if device.type == "cpu": |
| cap += CPU_BUCKET |
| own_k, own_v = [], [] |
| for k, _ in prefixes[0].kv: |
| shape = (rows, k.shape[1], cap, k.shape[3]) |
| own_k.append(torch.zeros(shape, device=device, dtype=k.dtype)) |
| own_v.append(torch.zeros(shape, device=device, dtype=k.dtype)) |
| lens = [int(p.kv[0][0].shape[2]) if p.kv else int(p.lengths[0]) for p in prefixes] |
| offs = [sum(lens[:i]) for i in range(len(lens))] |
| if len(prefixes) == 1: |
| pk = [k[0] for k, _ in prefixes[0].kv] |
| pv = [v[0] for _, v in prefixes[0].kv] |
| prec = [r.contiguous() for r, _ in prefixes[0].lin] |
| else: |
| pk = [torch.cat([p.kv[i][0][0] for p in prefixes], 1) |
| for i in range(len(prefixes[0].kv))] |
| pv = [torch.cat([p.kv[i][1][0] for p in prefixes], 1) |
| for i in range(len(prefixes[0].kv))] |
| prec = [torch.cat([p.lin[i][0] for p in prefixes]).contiguous() |
| for i in range(len(prefixes[0].lin))] |
| idx = torch.tensor(prow, dtype=torch.long, device=device) |
| tails = [torch.cat([p.lin[i][1] for p in prefixes]).index_select(0, idx) |
| for i in range(len(prefixes[0].lin))] |
| return ForkState(prec, pk, pv, torch.tensor(offs, device=device)[idx], |
| torch.tensor(lens, device=device)[idx], idx, |
| [None] * len(prefixes[0].lin), tails, own_k, own_v, |
| torch.zeros(rows, dtype=torch.long, device=device)) |
|
|
| def fork_forward(self, ids: Tensor, lengths: Tensor, fs: ForkState, keep: bool = False, |
| decode: bool = False, sel: Tensor | None = None) -> Tensor: |
| """See :meth:`_fork_forward`; ``decode`` (one token per row) selects the small-M |
| GEMM config for every projection of this pass; ``sel`` [B, n*r] selects each |
| row's LoRA adapter (None = off).""" |
| global _SMALL_GEMM |
| prev, _SMALL_GEMM = _SMALL_GEMM, decode |
| try: |
| return self._fork_forward(ids, lengths, fs, keep, sel) |
| finally: |
| _SMALL_GEMM = prev |
|
|
| def _fork_forward(self, ids: Tensor, lengths: Tensor, fs: ForkState, keep: bool = False, |
| sel: Tensor | None = None) -> Tensor: |
| """Run fork rows ids [B,L] (right-padded, ``lengths`` real tokens) on top of the |
| shared prefix and their own earlier tokens -> final-norm hidden [B,L,C]. |
| |
| Batch-invariant: every op is per row (bi_linear GEMMs, fork kernels, elementwise), |
| so a row's hidden states are bitwise the same in any batch. ``keep`` stores the |
| rows' new recurrent state, conv tail and own length (for continued decoding).""" |
| B, L = ids.shape |
| if not ids.is_cuda and (B > 1 or (L > 1 and L % CPU_BUCKET)): |
| return self._fork_forward_rows(ids, lengths, fs, keep, sel) |
| device = ids.device |
| if profile._ACTIVE.get() is not None and not torch.cuda.is_current_stream_capturing(): |
| profile.count(forwards=1, forward_rows=B, forward_tokens=B * L) |
| h = self.text.embed_tokens(ids) |
| valid = (torch.arange(L, device=device)[None] < lengths[:, None]).to(torch.float32) |
| pos = fs.plen[:, None] + fs.own_len[:, None] + torch.arange(L, device=device)[None] |
| cos, sin = self.rope(h, pos) |
| li = ai = 0 |
| for layer, kind in zip(self.text.layers, self.types): |
| if kind == "rec": |
| h = self.fam.fork_rec_layer(layer, h, valid, lengths, fs, li, keep, sel) |
| li += 1 |
| else: |
| h = self.fam.fork_attn_layer(layer, h, cos, sin, fs, ai, sel) |
| ai += 1 |
| if keep: |
| if fs.static: |
| fs.own_len.add_(lengths) |
| else: |
| fs.own_len = fs.own_len + lengths |
| fs.fresh = False |
| return self.fork_final_norm(h) |
|
|
| def _fork_forward_rows(self, ids: Tensor, lengths: Tensor, fs: ForkState, keep: bool, |
| sel: Tensor | None = None) -> Tensor: |
| """CPU reference: torch's CPU elementwise kernels switch between SIMD body and |
| scalar tail by tensor size, so a row's numerics would depend on the batch. Run |
| each row alone at a width padded to :data:`CPU_BUCKET` (decode steps: width 1).""" |
| B, L = ids.shape |
| W = 1 if L == 1 else -(-L // CPU_BUCKET) * CPU_BUCKET |
| outs, recs, tails, lens = [], [], [], [] |
| for b in range(B): |
| sub = ForkState(fs.prefix_rec, fs.pk, fs.pv, fs.poff[b:b + 1], fs.plen[b:b + 1], |
| fs.prow[b:b + 1], |
| [r[b:b + 1] if r is not None else None for r in fs.rec], |
| [t[b:b + 1] for t in fs.tail], [k[b:b + 1] for k in fs.own_k], |
| [v[b:b + 1] for v in fs.own_v], fs.own_len[b:b + 1], fs.fresh, |
| wpoff=None if fs.wpoff is None else fs.wpoff[b:b + 1], |
| wplen=None if fs.wplen is None else fs.wplen[b:b + 1]) |
| row = torch.full((1, W), self.lm.config.pad_token_id or 0, dtype=ids.dtype) |
| row[0, :L] = ids[b] |
| outs.append(self._fork_forward(row, lengths[b:b + 1], sub, keep, |
| None if sel is None else sel[b:b + 1])[:, :L]) |
| recs.append(sub.rec) |
| tails.append(sub.tail) |
| lens.append(sub.own_len) |
| if keep: |
| fs.rec = [torch.cat([r[i] for r in recs]) for i in range(len(fs.rec))] |
| fs.tail = [torch.cat([t[i] for t in tails]) for i in range(len(fs.tail))] |
| fs.own_len = torch.cat(lens) |
| fs.fresh = False |
| return torch.cat(outs) |
|
|
| def lm_logits(self, hidden: Tensor) -> Tensor: |
| """LM head through the batch-invariant GEMM (a row's logits never depend on the |
| other rows decoding with it). The Triton GEMM has no backward: when autograd needs |
| a gradient (training) the head is a plain differentiable ``F.linear``.""" |
| weight = self.lm.lm_head.weight |
| if torch.is_grad_enabled() and (hidden.requires_grad or weight.requires_grad): |
| return F.linear(hidden, weight) |
| from deg.kernels import bi_linear_small |
|
|
| return bi_linear_small(hidden, weight) |
|
|
|
|
| class GemmaBackbone(Backbone): |
| """Gemma 4 text stack (:data:`GEMMA4`) on the same State / ForkState contract. |
| |
| K/V slots: one per layer that computes K/V (layers ``0 .. n - num_kv_shared_layers - 1``; |
| E2B: 15 = 12 sliding + 3 global); ``kv_slot[i]`` maps every layer to its slot (a |
| K/V-sharing layer to the last earlier owner of its type). ``State.kv[slot]`` holds, per |
| row, a global slot's whole prefix K/V [B, Hkv, P, D] (``kv_mask`` marks real positions) |
| and a sliding slot's **tail**: the K/V of the row's last ``T = sliding_window - 1`` real |
| positions [B, Hkv, T, D], right-aligned at the row's length (slot j = position |
| ``lengths - T + j``; slots before position 0 are zeros and masked). A query never sees |
| more than T earlier positions in a sliding layer, so the tail is exact and a fork's |
| sliding state is bounded (T per layer, whatever the context length). ``State.lin`` is |
| empty (no recurrent layer). |
| |
| PLE: :meth:`ple_inputs` gathers the frozen per-layer token table (on the device, or in |
| host memory after :meth:`ple_to_host`) and adds the projected embedding.""" |
|
|
| def __init__(self, lm: nn.Module) -> None: |
| super().__init__(lm) |
| cfg = self.config |
| if getattr(cfg, "enable_moe_block", False): |
| raise ValueError("Gemma 4 MoE blocks are not supported") |
| if getattr(cfg, "attention_k_eq_v", False): |
| raise ValueError("Gemma 4 attention_k_eq_v (v = k) is not supported") |
| if getattr(cfg, "use_bidirectional_attention", None): |
| raise ValueError("Gemma 4 bidirectional attention is not supported") |
| n = cfg.num_hidden_layers |
| first_shared = n - int(getattr(cfg, "num_kv_shared_layers", 0) or 0) |
| self.kinds = ["sliding" if t == "sliding_attention" else "full" |
| for t in cfg.layer_types] |
| self.window = int(cfg.sliding_window) |
| self.tail = self.window - 1 |
| self.owner = [i < first_shared for i in range(n)] |
| slot_of_layer: dict[int, int] = {} |
| self.slot_kinds: list[str] = [] |
| self.slot_layers: list[int] = [] |
| last_of_kind: dict[str, int] = {} |
| for i in range(n): |
| if self.owner[i]: |
| slot_of_layer[i] = len(self.slot_kinds) |
| self.slot_kinds.append(self.kinds[i]) |
| self.slot_layers.append(i) |
| last_of_kind[self.kinds[i]] = slot_of_layer[i] |
| else: |
| slot_of_layer[i] = last_of_kind[self.kinds[i]] |
| self.kv_slot = [slot_of_layer[i] for i in range(n)] |
| |
| self.last_use = [max(i for i in range(n) if self.kv_slot[i] == k) |
| for k in range(len(self.slot_kinds))] |
| self.ple_dim = int(getattr(cfg, "hidden_size_per_layer_input", 0) or 0) |
| if self.ple_dim: |
| self.text.embed_tokens_per_layer.weight.requires_grad_(False) |
| self.softcap = getattr(cfg, "final_logit_softcapping", None) |
| self.types = ["attn"] * n |
|
|
| @property |
| def ple_on_host(self) -> bool: |
| t = getattr(self.text, "embed_tokens_per_layer", None) |
| return t is not None and t.weight.device != self.text.embed_tokens.weight.device |
|
|
| def ple_to_host(self) -> None: |
| """Keep the frozen PLE table (262144 x 35*256 bf16 = 4.4 GiB on E2B) in pinned host |
| memory: each pass gathers its tokens' rows on the host and copies them in. Call |
| after the model is on its device; never trained.""" |
| t = self.text.embed_tokens_per_layer |
| t.weight.requires_grad_(False) |
| if t.weight.device.type == "cpu": |
| return |
| w = t.weight.detach().cpu() |
| t.weight = nn.Parameter(w.pin_memory() if torch.cuda.is_available() else w, |
| requires_grad=False) |
|
|
| def ple_rows(self, ids: Tensor) -> Tensor: |
| """Unscaled PLE table rows of ``ids`` [..] -> [.., n_layers * ple_dim] on ids' device |
| (gathered on the host when the table lives there).""" |
| w = self.text.embed_tokens_per_layer.weight |
| if w.device == ids.device: |
| return F.embedding(ids, w) |
| rows = F.embedding(ids.cpu(), w) |
| if ids.is_cuda: |
| rows = rows.pin_memory() |
| return rows.to(ids.device, non_blocking=True) |
|
|
| def ple_inputs(self, ids: Tensor, h: Tensor, fork: bool = False, |
| rows: Tensor | None = None) -> Tensor | None: |
| """Per-layer inputs [B, L, n_layers, ple_dim] (HF ``get_per_layer_inputs`` + |
| ``project_per_layer_inputs``) for token ids [B, L] whose scaled embeddings are |
| ``h``; ``fork``: batch-invariant GEMM/norm (fork path); ``rows``: precomputed |
| :meth:`ple_rows` (a decode graph's static buffer).""" |
| if not self.ple_dim: |
| return None |
| text = self.text |
| B, L = ids.shape |
| nl = self.config.num_hidden_layers |
| emb = text.embed_tokens_per_layer |
| tok = (self.ple_rows(ids) if rows is None else rows) * emb.embed_scale.to( |
| device=h.device, dtype=emb.weight.dtype) |
| proj_mod = text.per_layer_model_projection |
| proj = _bi(proj_mod, h) if fork else proj_mod(h) |
| proj = (proj * text.per_layer_model_projection_scale).reshape(B, L, nl, self.ple_dim) |
| proj = (_gfnorm if fork else _gnorm)(text.per_layer_projection_norm, proj) |
| return (proj + tok.reshape(B, L, nl, self.ple_dim).to(proj.dtype)) \ |
| * text.per_layer_input_scale |
|
|
| def ropes(self, h: Tensor, pos: Tensor) -> dict[str, tuple[Tensor, Tensor]]: |
| """cos/sin [B, L, D] per layer kind for positions ``pos`` [B, L].""" |
| out = {} |
| for kind, name in (("sliding", "sliding_attention"), ("full", "full_attention")): |
| if kind in self.kinds: |
| out[kind] = self.text.rotary_emb(h, pos, name) |
| return out |
|
|
| def rope(self, h: Tensor, pos: Tensor) -> tuple[Tensor, Tensor]: |
| raise TypeError("Gemma 4 has one rotary table per layer kind: use ropes()") |
|
|
| def _slot_shape(self, slot: int) -> tuple[int, int]: |
| a = self.text.layers[self.slot_layers[slot]].self_attn |
| return a.k_proj.out_features // a.head_dim, a.head_dim |
|
|
| def empty_state(self, batch: int, device, dtype) -> State: |
| kv = [] |
| for slot, kind in enumerate(self.slot_kinds): |
| hkv, d = self._slot_shape(slot) |
| n = self.tail if kind == "sliding" else 0 |
| kv.append((torch.zeros(batch, hkv, n, d, device=device, dtype=dtype), |
| torch.zeros(batch, hkv, n, d, device=device, dtype=dtype))) |
| return State([], kv, torch.zeros(batch, 0, dtype=torch.bool, device=device), |
| torch.zeros(batch, dtype=torch.long, device=device)) |
|
|
| def _window_mask(self, plens: Tensor, L: int) -> Tensor: |
| """Sliding layers' mask [B, L, T+L] over (tail, this call's tokens): tail slot j |
| (position ``plens - T + j``) is visible to query i iff it exists and ``j >= i``; |
| own key j iff ``i - window < j <= i``.""" |
| device = plens.device |
| T = self.tail |
| i = torch.arange(L, device=device)[:, None] |
| j = torch.arange(T, device=device)[None] |
| tail_ok = (j >= i)[None] & (j[None] >= (T - plens)[:, None, None]) |
| jj = torch.arange(L, device=device)[None] |
| own = (jj <= i) & (i - jj < self.window) |
| return torch.cat([tail_ok, own[None].expand(plens.shape[0], -1, -1)], 2) |
|
|
| def _tail_of(self, ext: Tensor, lengths: Tensor) -> Tensor: |
| """The last T real positions of ext [B, Hkv, T+L, D] (tail + this call's tokens, |
| right-padded): gathered at each row's length, like a conv tail.""" |
| T = self.tail |
| idx = lengths[:, None] + torch.arange(T, device=ext.device)[None] |
| return ext.gather(2, idx[:, None, :, None].expand(-1, ext.shape[1], -1, ext.shape[3])) |
|
|
| def forward(self, ids: Tensor, lengths: Tensor, state: State | None = None, |
| consume: bool = False, rows: Tensor | None = None, sel: Tensor | None = None, |
| soft: tuple[Tensor, Tensor] | None = None) -> tuple[Tensor, State | None]: |
| """Same contract as :meth:`Backbone.forward`. Inputs longer than |
| :data:`DENSE_PREFILL_MAX` are read :data:`PREFILL_CHUNK` tokens at a time (the |
| state, sliding tails included, is carried across chunks). |
| |
| ``soft`` = (mask [B, L] bool, embeds [N, C]): multimodal soft tokens (HF Gemma 4 |
| placeholders): the N masked positions (row-major) take ``embeds`` in place of their |
| token embeddings, and their PLE token part is the pad token's (as HF does).""" |
| if rows is not None and ids.shape[1] > DENSE_PREFILL_MAX: |
| assert state is not None |
| state, rows = state.select(rows), None |
| if soft is not None and (rows is not None or soft[0].shape != ids.shape |
| or int(soft[0].sum()) != soft[1].shape[0]): |
| raise ValueError("soft tokens: mask [B, L] over ids with one embedding per " |
| "masked position (not on fork rows)") |
| if ids.shape[1] > DENSE_PREFILL_MAX: |
| return self._forward_chunked(ids, lengths, state, consume, sel, soft) |
| if profile._ACTIVE.get() is not None: |
| profile.count(forwards=1, forward_rows=ids.shape[0], |
| forward_tokens=ids.numel(), forward_real_tokens=int(lengths.sum())) |
| return self._read(ids, lengths, state, consume, rows, sel, soft=soft) |
|
|
| def embed(self, ids: Tensor, soft: tuple[Tensor, Tensor] | None = None, |
| fork: bool = False) -> tuple[Tensor, Tensor | None]: |
| """(scaled token embeddings [B, L, C] with soft tokens merged in, PLE inputs).""" |
| if soft is None: |
| h = self.text.embed_tokens(ids) |
| return h, self.ple_inputs(ids, h, fork=fork) |
| mask, feats = soft |
| tids = ids.masked_fill(mask, int(self.config.pad_token_id or 0)) |
| h = self.text.embed_tokens(tids) |
| h = h.masked_scatter(mask[..., None], feats.to(h.dtype)) |
| return h, self.ple_inputs(tids, h, fork=fork) |
|
|
| def _read(self, ids: Tensor, lengths: Tensor, state: State | None, consume: bool, |
| rows: Tensor | None, sel: Tensor | None, |
| parts: list[tuple[list[Tensor], list[Tensor]]] | None = None, s: int = 0, |
| soft: tuple[Tensor, Tensor] | None = None) -> tuple[Tensor, State | None]: |
| """One dense read of ids [B, L] on top of ``state``. ``parts`` (chunked training |
| read): per global slot, the incoming prefix K/V and the earlier chunks' K/V of this |
| call (concatenated inside each checkpoint, so no O(L^2/C) copies are saved; the new |
| chunk's K/V is appended); ``s`` = tokens of this call before ids. ``state`` then |
| carries the base ``kv_mask``, the carried sliding tails and each row's real tokens |
| so far (``lengths``); its global entries are returned unchanged.""" |
| B, L = ids.shape |
| device = ids.device |
| h, ple = self.embed(ids, soft) |
| if state is None: |
| state = self.empty_state(B, device, h.dtype) |
| kv_mask, plens = state.kv_mask, state.lengths |
| if rows is not None: |
| kv_mask, plens = kv_mask.index_select(0, rows), plens.index_select(0, rows) |
| pos = plens[:, None] + torch.arange(L, device=device)[None] |
| ropes = self.ropes(h, pos) |
| ckpt = self.grad_ckpt and torch.is_grad_enabled() |
| masks: dict[str, Any] = {"sliding": self._window_mask(plens, L)} |
| if parts is not None: |
| masks["full"] = (kv_mask, s) |
| else: |
| causal = torch.ones(L, L, dtype=torch.bool, device=device).tril() |
| masks["full"] = torch.cat([kv_mask[:, None, :].expand(-1, L, -1), |
| causal[None].expand(B, -1, -1)], 2) |
| new_kv: list[tuple[Tensor, Tensor] | None] = [None] * len(self.slot_kinds) |
| cur: dict[int, tuple[Tensor, Tensor]] = {} |
| for i, layer in enumerate(self.text.layers): |
| slot, kind = self.kv_slot[i], self.kinds[i] |
| kp, vp = state.kv[slot] |
| if parts is not None and kind == "full": |
| kp, vp = tuple(parts[slot][0]), tuple(parts[slot][1]) |
| kn, vn = (None, None) if self.owner[i] else cur[slot] |
| args = (layer, h, None if ple is None else ple[:, :, i], *ropes[kind], kp, vp, |
| masks[kind], rows, sel, kn, vn, self.window if kind == "sliding" else 0) |
| if ckpt: |
| h, k, v = checkpoint(_gemma_layer, *args, use_reentrant=False) |
| else: |
| h, k, v = _gemma_layer(*args) |
| del args |
| if self.owner[i]: |
| cur[slot] = (k, v) |
| if rows is None: |
| if kind == "sliding": |
| new_kv[slot] = (self._tail_of(_cat_kv(kp, k), lengths), |
| self._tail_of(_cat_kv(vp, v), lengths)) |
| elif parts is None: |
| new_kv[slot] = (_cat_kv(kp, k), _cat_kv(vp, v)) |
| del k, v |
| if consume and i == self.last_use[slot]: |
| del kp, vp |
| state.kv[slot] = None |
| h = self.final_norm(h) |
| if rows is not None: |
| return h, None |
| if parts is not None: |
| for slot, kind in enumerate(self.slot_kinds): |
| if kind == "full": |
| parts[slot][0].append(cur[slot][0]) |
| parts[slot][1].append(cur[slot][1]) |
| new_kv[slot] = state.kv[slot] |
| valid = torch.arange(L, device=device)[None] < lengths[:, None] |
| new = State([], new_kv, torch.cat([state.kv_mask, valid], 1), |
| state.lengths + lengths) |
| return h, new |
|
|
| def _forward_chunked(self, ids: Tensor, lengths: Tensor, state: State | None, |
| consume: bool = False, sel: Tensor | None = None, |
| soft: tuple[Tensor, Tensor] | None = None) -> tuple[Tensor, State]: |
| """Chunked read: :data:`PREFILL_CHUNK` tokens at a time, each chunk a dense read on |
| top of the state left by the previous one (sliding tails carried, global K/V |
| appended; rows are right-padded, so a row's real tokens stay contiguous). With |
| autograd + ``grad_ckpt``, the global K/V of the earlier chunks are passed to the |
| checkpoints as parts (never re-concatenated outside them); the returned state |
| concatenates them once.""" |
| B, L = ids.shape |
| device = ids.device |
| C = PREFILL_CHUNK |
| if profile._ACTIVE.get() is not None: |
| profile.count(forwards=1, forward_rows=B, forward_tokens=B * L, |
| forward_real_tokens=int(lengths.sum())) |
| if state is None: |
| state = self.empty_state(B, device, self.text.embed_tokens.weight.dtype) |
| train = self.grad_ckpt and torch.is_grad_enabled() |
| outs = [] |
| index = None |
| if soft is not None: |
| index = soft[0].flatten().cumsum(0).reshape(soft[0].shape) - 1 |
|
|
| def soft_of(s0: int, e: int) -> tuple[Tensor, Tensor] | None: |
| if soft is None or index is None: |
| return None |
| m = soft[0][:, s0:e] |
| return m, soft[1][index[:, s0:e][m]] |
|
|
| if not train: |
| for s0 in range(0, L, C): |
| e = min(L, s0 + C) |
| h, nxt = self._read(ids[:, s0:e], (lengths - s0).clamp(0, e - s0), state, |
| consume or s0 > 0, None, sel, soft=soft_of(s0, e)) |
| assert nxt is not None |
| state = nxt |
| outs.append(h) |
| return torch.cat(outs, 1), state |
| parts = [([kp], [vp]) if kind == "full" else ([], []) |
| for (kp, vp), kind in zip(state.kv, self.slot_kinds)] |
| cur = State([], list(state.kv), state.kv_mask, state.lengths) |
| for s0 in range(0, L, C): |
| e = min(L, s0 + C) |
| lens = (lengths - s0).clamp(0, e - s0) |
| |
| h, nxt = self._read(ids[:, s0:e], lens, cur, False, None, sel, parts, s0, |
| soft_of(s0, e)) |
| assert nxt is not None |
| |
| cur = State([], nxt.kv, state.kv_mask, state.lengths + lengths.clamp(max=e)) |
| outs.append(h) |
| kv = [(torch.cat(parts[k][0], 2), torch.cat(parts[k][1], 2)) if kind == "full" |
| else cur.kv[k] for k, kind in enumerate(self.slot_kinds)] |
| valid = torch.arange(L, device=device)[None] < lengths[:, None] |
| return torch.cat(outs, 1), State([], kv, torch.cat([state.kv_mask, valid], 1), |
| state.lengths + lengths) |
|
|
| def _forward_chunked_ckpt(self, ids: Tensor, lengths: Tensor, state: State | None, |
| sel: Tensor | None = None) -> tuple[Tensor, State]: |
| return self._forward_chunked(ids, lengths, state, False, sel) |
|
|
|
|
| def fork(self, prefixes: State | list[State], prow: list[int] | int, cap: int |
| ) -> ForkState: |
| """As :meth:`Backbone.fork`: global slots read the prefixes' K/V by reference |
| (several concatenated once); sliding slots read their tails (T per prefix, |
| concatenated: row b's live tail is ``wplen[b] = min(len, T)`` positions at |
| ``wpoff[b]``).""" |
| if isinstance(prefixes, State): |
| prefixes = [prefixes] |
| if isinstance(prow, int): |
| prow = [0] * prow |
| for p in prefixes: |
| assert p.lengths.shape[0] == 1 and bool(p.kv_mask.all()) |
| rows = len(prow) |
| device = prefixes[0].lengths.device |
| if device.type == "cpu": |
| cap += CPU_BUCKET |
| T = self.tail |
| own_k, own_v, pk, pv = [], [], [], [] |
| for slot in range(len(self.slot_kinds)): |
| k0 = prefixes[0].kv[slot][0] |
| shape = (rows, k0.shape[1], cap, k0.shape[3]) |
| own_k.append(torch.zeros(shape, device=device, dtype=k0.dtype)) |
| own_v.append(torch.zeros(shape, device=device, dtype=k0.dtype)) |
| if len(prefixes) == 1: |
| pk.append(k0[0]) |
| pv.append(prefixes[0].kv[slot][1][0]) |
| else: |
| pk.append(torch.cat([p.kv[slot][0][0] for p in prefixes], 1)) |
| pv.append(torch.cat([p.kv[slot][1][0] for p in prefixes], 1)) |
| lens = [int(p.lengths[0]) for p in prefixes] |
| offs = [sum(lens[:i]) for i in range(len(lens))] |
| live = [min(n, T) for n in lens] |
| woff = [i * T + T - n for i, n in enumerate(live)] |
| idx = torch.tensor(prow, dtype=torch.long, device=device) |
|
|
| def at(xs: list[int]) -> Tensor: |
| return torch.tensor(xs, dtype=torch.long, device=device)[idx] |
|
|
| return ForkState([], pk, pv, at(offs), at(lens), idx, [], [], own_k, own_v, |
| torch.zeros(rows, dtype=torch.long, device=device), |
| wpoff=at(woff), wplen=at(live)) |
|
|
| def _fork_forward(self, ids: Tensor, lengths: Tensor, fs: ForkState, keep: bool = False, |
| sel: Tensor | None = None) -> Tensor: |
| """:meth:`Backbone._fork_forward` over the K/V slots (sliding layers on their |
| prefix tails, K/V-sharing layers on their source's buffers). ``fs.ple_rows``: the |
| PLE table rows of ``ids`` gathered beforehand (a decode graph's static buffer when |
| the table lives in host memory).""" |
| B, L = ids.shape |
| if not ids.is_cuda and (B > 1 or (L > 1 and L % CPU_BUCKET)): |
| return self._fork_forward_rows(ids, lengths, fs, keep, sel) |
| device = ids.device |
| if profile._ACTIVE.get() is not None and not torch.cuda.is_current_stream_capturing(): |
| profile.count(forwards=1, forward_rows=B, forward_tokens=B * L) |
| h = self.text.embed_tokens(ids) |
| ple = self.ple_inputs(ids, h, fork=True, rows=fs.ple_rows) |
| pos = fs.plen[:, None] + fs.own_len[:, None] + torch.arange(L, device=device)[None] |
| ropes = self.ropes(h, pos) |
| for i, layer in enumerate(self.text.layers): |
| kind = self.kinds[i] |
| h = _gemma_fork_layer(layer, h, None if ple is None else ple[:, :, i], *ropes[kind], |
| fs, self.kv_slot[i], self.owner[i], |
| self.window if kind == "sliding" else 0, sel) |
| if keep: |
| if fs.static: |
| fs.own_len.add_(lengths) |
| else: |
| fs.own_len = fs.own_len + lengths |
| fs.fresh = False |
| return self.fork_final_norm(h) |
|
|
| def lm_logits(self, hidden: Tensor) -> Tensor: |
| """:meth:`Backbone.lm_logits` with Gemma's final logit soft-cap |
| (``cap * tanh(logits / cap)``, fp32; argmax unchanged).""" |
| logits = super().lm_logits(hidden) |
| if not self.softcap: |
| return logits |
| return torch.tanh(logits.float() / self.softcap) * self.softcap |
|
|
| def static_fork_state(self, rows: int, pcap: int, cap: int) -> ForkState: |
| """CUDA-graph buffers (:class:`DecodeGraph`): global slots' prefixes up to ``pcap`` |
| tokens in total, sliding slots' tails for up to ``rows`` prefixes, own K/V of |
| ``cap`` tokens per row.""" |
| dev = self.text.embed_tokens.weight.device |
| dtype = self.text.embed_tokens.weight.dtype |
| pk, pv, ok, ov = [], [], [], [] |
| for slot, kind in enumerate(self.slot_kinds): |
| hkv, d = self._slot_shape(slot) |
| n = rows * self.tail if kind == "sliding" else pcap |
| pk.append(torch.zeros(hkv, n, d, device=dev, dtype=dtype)) |
| pv.append(torch.zeros(hkv, n, d, device=dev, dtype=dtype)) |
| ok.append(torch.zeros(rows, hkv, cap, d, device=dev, dtype=dtype)) |
| ov.append(torch.zeros(rows, hkv, cap, d, device=dev, dtype=dtype)) |
| z = torch.zeros(rows, dtype=torch.long, device=dev) |
| return ForkState([], pk, pv, z.clone(), z.clone(), z.clone(), [], [], ok, ov, |
| z.clone(), fresh=False, static=True, wpoff=z.clone(), wplen=z.clone()) |
|
|
| def load_static(self, fs: ForkState, prefixes: list[State], prow: list[int]) -> None: |
| """Copy batch-1 prefixes into :meth:`static_fork_state` buffers; row b (< |
| len(prow)) continues ``prefixes[prow[b]]``, spare rows idle on prefix 0.""" |
| rows = fs.own_len.shape[0] |
| T = self.tail |
| lens = [int(p.lengths[0]) for p in prefixes] |
| assert len(prefixes) <= rows and len(prow) <= rows |
| offs = [sum(lens[:i]) for i in range(len(lens))] |
| for pi, (p, o, n) in enumerate(zip(prefixes, offs, lens)): |
| for slot, kind in enumerate(self.slot_kinds): |
| k, v = p.kv[slot] |
| if kind == "sliding": |
| fs.pk[slot][:, pi * T:(pi + 1) * T].copy_(k[0]) |
| fs.pv[slot][:, pi * T:(pi + 1) * T].copy_(v[0]) |
| else: |
| assert o + n <= fs.pk[slot].shape[1] |
| fs.pk[slot][:, o:o + n].copy_(k[0]) |
| fs.pv[slot][:, o:o + n].copy_(v[0]) |
| dev = fs.own_len.device |
| idx = torch.tensor(list(prow) + [0] * (rows - len(prow)), dtype=torch.long, device=dev) |
| live = [min(n, T) for n in lens] |
| assert fs.wpoff is not None and fs.wplen is not None |
| fs.poff.copy_(torch.tensor(offs, device=dev)[idx]) |
| fs.plen.copy_(torch.tensor(lens, device=dev)[idx]) |
| fs.prow.copy_(idx) |
| fs.wpoff.copy_(torch.tensor([i * T + T - n for i, n in enumerate(live)], device=dev)[idx]) |
| fs.wplen.copy_(torch.tensor(live, device=dev)[idx]) |
| fs.own_len.zero_() |
| fs.fresh = True |
|
|
|
|
|
|
| class MediaTowers(nn.Module): |
| """Gemma 4's frozen vision and audio encoders with their embedders (HF |
| ``Gemma4Model.vision_tower`` / ``audio_tower`` / ``embed_vision`` / ``embed_audio``), |
| loaded from the base snapshot on demand (:func:`load_media_towers`; optional: a text |
| deployment never loads them). Never trained or saved by deg: a checkpoint's media |
| path is its base snapshot's (``meta["base"]``). Outputs are soft-token embeddings in |
| the language model's space for :class:`Media` / :meth:`Renderer.context`.""" |
|
|
| def __init__(self, config: Any) -> None: |
| super().__init__() |
| from transformers import AutoModel |
| from transformers.models.gemma4.modeling_gemma4 import Gemma4MultimodalEmbedder |
|
|
| self.config = config |
| text = config.text_config |
| self.vision_tower = (AutoModel.from_config(config.vision_config) |
| if config.vision_config is not None else None) |
| self.audio_tower = (AutoModel.from_config(config.audio_config) |
| if config.audio_config is not None else None) |
| self.embed_vision = (Gemma4MultimodalEmbedder(config.vision_config, text) |
| if config.vision_config is not None else None) |
| self.embed_audio = (Gemma4MultimodalEmbedder(config.audio_config, text) |
| if config.audio_config is not None else None) |
|
|
| @torch.no_grad() |
| def image_embeds(self, pixel_values: Tensor, position_ids: Tensor) -> list[Tensor]: |
| """HF ``get_image_features``: processor outputs (``pixel_values`` [I, patches, |
| 3*p*p], ``image_position_ids`` [I, patches, 2]) -> per image [n_i, hidden].""" |
| assert self.vision_tower is not None and self.embed_vision is not None |
| out = self.vision_tower(pixel_values=pixel_values, pixel_position_ids=position_ids) |
| pooled = self.embed_vision(inputs_embeds=out.last_hidden_state) |
| k = self.config.vision_config.pooling_kernel_size |
| output_length = pixel_values.shape[-2] // (k * k) |
| k_squared = int((position_ids.shape[1] // output_length) ** 0.5) ** 2 |
| sizes = ((position_ids != -1).all(-1).sum(-1) // k_squared).tolist() |
| return list(torch.split(pooled.reshape(-1, pooled.shape[-1]), sizes)) |
|
|
| @torch.no_grad() |
| def audio_embeds(self, features: Tensor, mask: Tensor) -> list[Tensor]: |
| """HF ``get_audio_features``: feature-extractor outputs (``input_features`` [A, T, |
| mels], ``input_features_mask`` [A, T]) -> per clip its valid soft tokens.""" |
| assert self.audio_tower is not None and self.embed_audio is not None |
| out = self.audio_tower(features, mask, return_dict=True) |
| pooled = self.embed_audio(inputs_embeds=out.last_hidden_state) |
| valid = out.attention_mask |
| return [pooled[i][valid[i]] for i in range(pooled.shape[0])] |
|
|
|
|
| MEDIA_PREFIXES = ("model.vision_tower.", "model.audio_tower.", "model.embed_vision.", |
| "model.embed_audio.") |
|
|
|
|
| def load_media_towers(snapshot: Path | str | None = None, *, backbone: str | None = None, |
| device: str | torch.device = "cpu", |
| dtype: torch.dtype = torch.bfloat16) -> MediaTowers: |
| """The frozen vision/audio towers of a Gemma 4 snapshot (default: the backbone's base |
| snapshot), read from its safetensors without loading the text model.""" |
| import transformers |
| from safetensors import safe_open |
|
|
| snap = Path(snapshot) if snapshot is not None else base_snapshot(backbone or |
| "gemma-4-e2b-it") |
| cfg = transformers.AutoConfig.from_pretrained(str(snap)) |
| towers = MediaTowers(cfg) |
| state: dict[str, Tensor] = {} |
| for f in sorted(snap.glob("*.safetensors")): |
| with safe_open(str(f), "pt") as fh: |
| for k in fh.keys(): |
| if k.startswith(MEDIA_PREFIXES): |
| state[k[len("model."):]] = fh.get_tensor(k) |
| missing, unexpected = towers.load_state_dict(state, strict=False) |
| if unexpected or [k for k in missing if not k.endswith(("inv_freq",))]: |
| raise RuntimeError(f"media towers mismatch: {(list(unexpected) + missing)[:5]}") |
| towers.to(device=device, dtype=dtype) |
| towers.requires_grad_(False) |
| towers.eval() |
| return towers |
|
|
|
|
| class DecodeGraph: |
| """One generate decode step (``rows`` fork rows x 1 token, LM head, argmax) captured |
| as a CUDA graph over static buffers: prefix K/V for up to ``pcap`` tokens in total |
| (row offsets/lengths are device tensors), per-row recurrent state / conv tail / own |
| K/V (``cap`` own tokens). :meth:`load` copies the prefixes in; the header pass runs |
| eagerly on the same buffers; every decode step is then one ``replay``. The step is |
| the same batch-invariant fork forward as the eager path, so replay changes speed, |
| not semantics. ``sel``: the fixed route of every step (a deg-MoE's |
| :class:`WriteRoute` for ``rows`` rows; None = the backbone as is).""" |
|
|
| def __init__(self, backbone: Backbone, rows: int, pcap: int, cap: int, |
| sel: WriteRoute | None = None) -> None: |
| self.backbone, self.rows, self.pcap, self.cap = backbone, rows, pcap, cap |
| self.sel = sel |
| dev = backbone.text.embed_tokens.weight.device |
| dtype = backbone.text.embed_tokens.weight.dtype |
| self.ids = torch.zeros(rows, 1, dtype=torch.long, device=dev) |
| self.lens = torch.zeros(rows, dtype=torch.long, device=dev) |
| self.graph: torch.cuda.CUDAGraph | None = None |
| self.logits: Tensor | None = None |
| self.top1: Tensor | None = None |
| |
| |
| self.ple_rows: Tensor | None = None |
| if isinstance(backbone, GemmaBackbone): |
| self.fs = backbone.static_fork_state(rows, pcap, cap) |
| if backbone.ple_on_host: |
| w = backbone.text.embed_tokens_per_layer.weight |
| self.ple_rows = torch.zeros(rows, 1, w.shape[1], device=dev, dtype=w.dtype) |
| return |
| rec, tail, pk, pv, ok, ov = [], [], [], [], [], [] |
| for layer, kind in zip(backbone.text.layers, backbone.types): |
| if kind == "rec": |
| r, t = backbone.rec_buffers(layer, rows, dev, dtype) |
| rec.append(r) |
| tail.append(t) |
| else: |
| hkv, d = backbone.config.num_key_value_heads, layer.self_attn.head_dim |
| pk.append(torch.zeros(hkv, pcap, d, device=dev, dtype=dtype)) |
| pv.append(torch.zeros(hkv, pcap, d, device=dev, dtype=dtype)) |
| ok.append(torch.zeros(rows, hkv, cap, d, device=dev, dtype=dtype)) |
| ov.append(torch.zeros(rows, hkv, cap, d, device=dev, dtype=dtype)) |
| zeros = torch.zeros(rows, dtype=torch.long, device=dev) |
| self.fs = ForkState([r[:1] for r in rec], pk, pv, zeros.clone(), zeros.clone(), |
| zeros.clone(), rec, tail, ok, ov, zeros.clone(), fresh=False, |
| static=True) |
|
|
| def _step(self) -> None: |
| self.fs.ple_rows = self.ple_rows |
| try: |
| hid = self.backbone.fork_forward(self.ids, self.lens, self.fs, keep=True, |
| decode=True, sel=self.sel) |
| finally: |
| self.fs.ple_rows = None |
| self.logits = self.backbone.lm_logits(hid[:, 0]).float() |
| self.top1 = self.logits.argmax(-1) |
|
|
| def capture(self) -> None: |
| self.fs.fresh = False |
| side = torch.cuda.Stream() |
| side.wait_stream(torch.cuda.current_stream()) |
| with torch.cuda.stream(side): |
| for _ in range(3): |
| self._step() |
| torch.cuda.current_stream().wait_stream(side) |
| self.graph = torch.cuda.CUDAGraph() |
| with torch.cuda.graph(self.graph): |
| self._step() |
|
|
| def load(self, prefixes: list[State], prow: list[int]) -> ForkState: |
| """Point the static rows at new batch-1 prefixes: row b (< len(prow)) continues |
| ``prefixes[prow[b]]``; spare rows idle on prefix 0. Fresh rows, own_len 0.""" |
| lens = [int(p.lengths[0]) for p in prefixes] |
| assert sum(lens) <= self.pcap and len(prow) <= self.rows |
| if isinstance(self.backbone, GemmaBackbone): |
| self.backbone.load_static(self.fs, prefixes, prow) |
| return self.fs |
| offs = [sum(lens[:i]) for i in range(len(lens))] |
| for p, o, n in zip(prefixes, offs, lens): |
| for i, (k, v) in enumerate(p.kv): |
| self.fs.pk[i][:, o:o + n].copy_(k[0]) |
| self.fs.pv[i][:, o:o + n].copy_(v[0]) |
| rows = list(prow) + [0] * (self.rows - len(prow)) |
| dev = self.fs.own_len.device |
| idx = torch.tensor(rows, dtype=torch.long, device=dev) |
| self.fs.poff.copy_(torch.tensor(offs, device=dev)[idx]) |
| self.fs.plen.copy_(torch.tensor(lens, device=dev)[idx]) |
| self.fs.prow.copy_(idx) |
| self.fs.prefix_rec = [torch.cat([p.lin[i][0] for p in prefixes]).contiguous() |
| for i in range(len(prefixes[0].lin))] |
| for i, t in enumerate(self.fs.tail): |
| t.copy_(torch.cat([p.lin[i][1] for p in prefixes]).index_select(0, idx)) |
| self.fs.own_len.zero_() |
| self.fs.fresh = True |
| return self.fs |
|
|
| def step(self, ids: Tensor, lens: Tensor) -> tuple[Tensor, Tensor]: |
| """One decode step for every row (``lens`` 0 = idle row) -> (logits, argmax).""" |
| if self.ple_rows is not None: |
| assert isinstance(self.backbone, GemmaBackbone) |
| self.ple_rows.copy_(self.backbone.ple_rows(ids.cpu()), non_blocking=True) |
| self.ids.copy_(ids, non_blocking=True) |
| self.lens.copy_(lens, non_blocking=True) |
| assert self.graph is not None and self.logits is not None and self.top1 is not None |
| self.fs.fresh = False |
| self.graph.replay() |
| return self.logits, self.top1 |
|
|
|
|
| def _pow2(n: int, lo: int) -> int: |
| out = lo |
| while out < n: |
| out *= 2 |
| return out |
|
|
|
|
| |
|
|
|
|
| class PointerHead(nn.Module): |
| def __init__(self, hidden: int, dim: int = HEAD_DIM) -> None: |
| super().__init__() |
| self.query = nn.Linear(hidden, dim) |
| self.key = nn.Linear(hidden, dim) |
| self.prior = nn.Linear(hidden, 1) |
| self.scale = 1.0 / math.sqrt(dim) |
|
|
| def forward(self, answer: Tensor, options: Tensor) -> Tensor: |
| """answer [..., C], options [..., O, C] -> logits [..., O].""" |
| q = self.query(answer).unsqueeze(-2) |
| return ((self.key(options) * q).sum(-1) * self.scale |
| + self.prior(options).squeeze(-1)) |
|
|
|
|
| class SpanHead(nn.Module): |
| def __init__(self, hidden: int, dim: int = HEAD_DIM) -> None: |
| super().__init__() |
| self.q_start, self.k_start = nn.Linear(hidden, dim), nn.Linear(hidden, dim) |
| self.q_end, self.k_end = nn.Linear(hidden, dim), nn.Linear(hidden, dim) |
| self.null_start, self.null_end = nn.Linear(hidden, 1), nn.Linear(hidden, 1) |
| self.scale = 1.0 / math.sqrt(dim) |
|
|
| def forward(self, answer: Tensor, context: Tensor) -> tuple[Tensor, Tensor]: |
| """answer [n, C], context [T, C] -> start/end logits [n, T+1] (last = null).""" |
| s = (self.q_start(answer) @ self.k_start(context).T) * self.scale |
| e = (self.q_end(answer) @ self.k_end(context).T) * self.scale |
| return (torch.cat([s, self.null_start(answer)], -1), |
| torch.cat([e, self.null_end(answer)], -1)) |
|
|
|
|
| class AnchorHead(nn.Module): |
| """Item anchors: one independent "an item of this array starts here" logit per context |
| token, ``q(a) . k(c_t) / sqrt(d) + u . c_t + b`` (the prior ``u . c_t + b`` starts at |
| :data:`ANCHOR_INIT_BIAS`).""" |
|
|
| def __init__(self, hidden: int, dim: int = HEAD_DIM) -> None: |
| super().__init__() |
| self.query = nn.Linear(hidden, dim) |
| self.key = nn.Linear(hidden, dim) |
| self.prior = nn.Linear(hidden, 1) |
| with torch.no_grad(): |
| self.prior.weight.mul_(0.1) |
| self.prior.bias.fill_(ANCHOR_INIT_BIAS) |
| self.scale = 1.0 / math.sqrt(dim) |
|
|
| def forward(self, answer: Tensor, context: Tensor) -> Tensor: |
| """answer [n, C], context [T, C] -> logits [n, T].""" |
| return ((self.query(answer) @ self.key(context).T) * self.scale |
| + self.prior(context)[:, 0][None]) |
|
|
|
|
| def rel_bucket(d: Tensor) -> Tensor: |
| """Bucket of a token offset ``t - anchor`` (:data:`REL_EDGES`; 0 .. len(REL_EDGES)).""" |
| return torch.bucketize(d, torch.tensor(REL_EDGES, device=d.device), right=True) |
|
|
|
|
| def span_rel_bias(table: Tensor, anchor: int, hi: int, r0: int, r1: int |
| ) -> tuple[Tensor, Tensor]: |
| """Anchored span bias for context tokens [r0, r1): ``table[which, inside, bucket]`` |
| with ``inside = anchor <= t < hi`` (the item's region) -> (start bias, end bias).""" |
| t = torch.arange(r0, r1, device=table.device) |
| b = rel_bucket(t - anchor) |
| inside = ((t >= anchor) & (t < hi)).long() |
| return table[0, inside, b], table[1, inside, b] |
|
|
|
|
| def region_bounds(region: tuple[int, int], T: int) -> tuple[int, int]: |
| lo, hi = region |
| return max(0, lo), (T if hi < 0 else min(hi, T)) |
|
|
|
|
| class Heads(nn.Module): |
| def __init__(self, hidden: int) -> None: |
| super().__init__() |
| self.choice = PointerHead(hidden) |
| self.score = PointerHead(hidden) |
| self.noul = nn.Linear(hidden, 1) |
| self.span = SpanHead(hidden) |
| self.count = nn.Linear(hidden, MAX_COUNT + 1) |
| self.anchor = AnchorHead(hidden) |
| |
| self.span_bias = nn.Parameter(torch.zeros(2, 2, len(REL_EDGES) + 1)) |
| self.register_buffer("temperature", torch.ones(len(KINDS))) |
|
|
| def temp(self, kind: str) -> Tensor: |
| return self.temperature[KINDS.index(kind)] |
|
|
|
|
| class _Picked: |
| """The rows of one fork's hidden states [Lq, C] that the heads read (answer and |
| option positions), indexable by original position; the rest is dropped.""" |
|
|
| __slots__ = ("rows", "at") |
|
|
| def __init__(self, hidden: Tensor, positions: tuple[int, ...]) -> None: |
| keep = sorted(set(positions)) |
| self.at = {p: k for k, p in enumerate(keep)} |
| self.rows = hidden[keep] |
|
|
| def __getitem__(self, p: int) -> Tensor: |
| return self.rows[self.at[p]] |
|
|
|
|
| def fork_token_budget(device: torch.device) -> int: |
| """Padded fork tokens per batched fork pass: :data:`FORK_TOKENS`, lowered on CUDA to |
| what fits in :data:`FORK_MEM_FRACTION` of the memory free right now (device free + |
| cached-but-unused), so a fork pass next to a big context (or another process on the |
| GPU) is planned to fit instead of OOM-and-split. Batch composition never changes a |
| row's outputs (batch-invariant kernels), only speed.""" |
| if device.type != "cuda": |
| return FORK_TOKENS |
| free, _ = torch.cuda.mem_get_info(device) |
| free += torch.cuda.memory_reserved(device) - torch.cuda.memory_allocated(device) |
| return max(256, min(FORK_TOKENS, int(free * FORK_MEM_FRACTION) // FORK_BYTES_PER_TOKEN)) |
|
|
|
|
| @dataclass |
| class QueryOutput: |
| """Raw (uncalibrated) logits for one query.""" |
|
|
| kind: str |
| logits: Tensor | None = None |
| start: Tensor | None = None |
| end: Tensor | None = None |
|
|
|
|
| @dataclass |
| class DegModel: |
| backbone: Backbone |
| heads: Heads |
| renderer: Renderer |
| meta: dict[str, Any] = field(default_factory=dict) |
|
|
| graphs: dict[tuple[int, int, int], DecodeGraph] = field(default_factory=dict) |
| |
| |
| base: Backbone | None = None |
| generate_with: str = "self" |
| _gen_view: DegModel | None = None |
| |
| |
| lora: dict[str, Any] | None = None |
| |
| |
| moe: dict[str, Any] | None = None |
|
|
| def gen_route(self, rows: int) -> WriteRoute | None: |
| """``sel`` of x-generate calls (KNOWN VALUES prefill, generate header, decode): |
| an MoE's W route for ``rows`` rows; None for every other layout.""" |
| if self.moe is None: |
| return None |
| if "experts" in self.moe: |
| return TaskRoute("g") |
| r = int(self.moe["writer_rank"]) |
| return WriteRoute(torch.ones(rows, r, dtype=torch.float32, device=self.device) |
| if r else None) |
|
|
| def adapter_of(self, kind: str) -> str | None: |
| """The adapter a fork kind runs with (None = adapter off, the base model).""" |
| if self.lora is None: |
| return None |
| names = self.lora["adapters"] |
| if names == ["shared"]: |
| return "shared" if kind in ADAPTER_OF_KIND else None |
| name = ADAPTER_OF_KIND.get(kind) |
| return name if name in names else None |
|
|
| def _sel(self, names: list[str | None]) -> Tensor | None: |
| if self.lora is None or all(n is None for n in names): |
| return None |
| r, order = int(self.lora["rank"]), self.lora["adapters"] |
| sel = torch.zeros(len(names), r * len(order), dtype=torch.float32) |
| for i, n in enumerate(names): |
| if n is not None: |
| k = order.index(n) |
| sel[i, k * r:(k + 1) * r] = 1.0 |
| return sel.to(self.device) |
|
|
| def lora_sel(self, kinds: list[str]) -> Tensor | TaskRoute | None: |
| """Per-row adapter selector for fork rows of these kinds (None: all off); for a |
| D/E/G task MoE the route of the call (every kind must share it).""" |
| if self.moe is not None and "experts" in self.moe: |
| routes = {self.fork_route(k) for k in kinds} |
| assert len(routes) == 1, f"one expert per fork call, got {routes}" |
| return TaskRoute(routes.pop()) |
| return self._sel([self.adapter_of(k) for k in kinds]) |
|
|
| def fork_route(self, kind: str) -> str | None: |
| """The expert a fork kind runs in a D/E/G task MoE (None: not one); fork batches |
| are grouped by it (:meth:`answer_many`, ``deg.train.forward_losses``).""" |
| if self.moe is None or "experts" not in self.moe: |
| return None |
| return TASK_EXPERT_OF_KIND[kind] |
|
|
| def context_sel(self, rows: int) -> Tensor | None: |
| """Selector of the context (trunk) read: None when the context is read by the base |
| weights (``lora["context"] == "base"``: one shared prefill for every fork, |
| generate included).""" |
| if self.lora is None or self.lora["context"] == "base": |
| return None |
| return self._sel([self.lora["context"]] * rows) |
|
|
| def generator(self) -> tuple[DegModel, bool, bool]: |
| """How x-generate fields decode: ``(model, trunk_shared, base_layout)``. |
| ``model`` decodes (adapter-off rows); ``trunk_shared``: the request's trunk state |
| can be extended by the KNOWN VALUES block (else the context is read again by |
| ``model``); ``base_layout``: the base model writes (base header + assistant turn). |
| |
| * full fine-tune, generate_with self: (self, True, False) -- deg's own LM |
| (``gen_layout == "chat"``: (self, True, True), its own weights in the chat layout); |
| * full fine-tune, generate_with base: (base copy, False, True) -- two-copy hybrid; |
| * LoRA, context base: (self, True, True) -- one prefill, generation = base; |
| * LoRA, context adapter: (self, False, True) -- adapter-off re-read of the context; |
| * deg-MoE: (self, True, True) -- one prefill (R), the W experts write in the base |
| layout (calls routed by :meth:`gen_route`). |
| """ |
| if self.moe is not None: |
| return self, True, True |
| if self.lora is not None: |
| return self, self.lora["context"] == "base", True |
| gm = self.generate_model() |
| if gm is self: |
| return self, True, self.gen_layout == "chat" |
| return gm, False, True |
|
|
| @property |
| def gen_layout(self) -> str: |
| """``meta["gen_layout"]`` (:data:`GEN_LAYOUTS`; default ``deg``).""" |
| layout = self.meta.get("gen_layout", "deg") |
| assert layout in GEN_LAYOUTS, layout |
| return layout |
|
|
| @property |
| def device(self) -> torch.device: |
| return next(self.heads.parameters()).device |
|
|
| def generate_model(self) -> DegModel: |
| """The model that decodes x-generate fields: ``self``, or (hybrid) a view with the |
| frozen base backbone, the same tokenizer and its own decode-graph cache.""" |
| if self.generate_with in ("self", "moe"): |
| return self |
| assert self.generate_with == "base", self.generate_with |
| return self.base_view() |
|
|
| def base_view(self) -> DegModel: |
| """This model with the frozen base backbone (:func:`attach_base`) in place of its |
| own: same heads, tokenizer and meta, its own decode-graph cache.""" |
| if self.base is None: |
| raise RuntimeError("no base backbone: attach_base(model) first") |
| if self._gen_view is None or self._gen_view.backbone is not self.base: |
| self._gen_view = DegModel(self.base, self.heads, self.renderer, self.meta) |
| return self._gen_view |
|
|
| def decode_graph(self, rows: int, prefix_len: int, cap: int) -> DecodeGraph: |
| """Cached decode graph for at least ``rows`` rows (1, 2, 4, ...), a prefix of |
| ``prefix_len`` tokens (pow2 >= 2048) and ``cap`` own tokens (pow2 >= 512).""" |
| key = (_pow2(rows, 1), _pow2(prefix_len, 2048), _pow2(cap, 512)) |
| if key not in self.graphs: |
| graph = DecodeGraph(self.backbone, *key, sel=self.gen_route(key[0])) |
| graph.capture() |
| self.graphs[key] = graph |
| return self.graphs[key] |
|
|
| def modules(self) -> list[nn.Module]: |
| return [self.backbone, self.heads] |
|
|
| @property |
| def array_mode(self) -> str: |
| """How array items are addressed (:data:`ARRAY_MODES`): ``anchor`` (position |
| anchors, no count cap) or ``index`` ("ITEM i of" + the capped count head; run-1).""" |
| mode = self.meta.get("array_mode", "index") |
| assert mode in ARRAY_MODES, mode |
| return mode |
|
|
| def array_mode_for(self, doc_tokens: int) -> str: |
| """The array mode of one document of ``doc_tokens`` context tokens: an ``anchor`` |
| model routes documents shorter than ``meta["anchor_min_tokens"]`` (default 0 = |
| never) to ``index`` (SPEC.md §12.6 Routing); inference and training share it.""" |
| mode = self.array_mode |
| if mode == "anchor" and doc_tokens < int(self.meta.get("anchor_min_tokens", 0)): |
| return "index" |
| return mode |
|
|
| |
|
|
| def apply_heads(self, eq: EncodedQuery, hidden: Tensor, context_hidden: Tensor | None |
| ) -> QueryOutput: |
| """hidden [Lq, C] of the fork (fp32 heads).""" |
| a = hidden[eq.answer_position].float() |
| if eq.kind in ("choice", "score"): |
| opts = hidden[list(eq.option_positions)].float() |
| head = self.heads.choice if eq.kind == "choice" else self.heads.score |
| return QueryOutput(eq.kind, logits=head(a, opts)) |
| if eq.kind == "noul": |
| lg = self.heads.noul(a)[0] |
| return QueryOutput("noul", logits=torch.stack([lg, torch.zeros_like(lg)])) |
| if eq.kind == "count": |
| return QueryOutput("count", logits=self.heads.count(a)) |
| assert context_hidden is not None |
| T = context_hidden.shape[0] |
| if eq.kind == "anchor": |
| lg = self.heads.anchor(a[None], context_hidden.float())[0] |
| lo, hi = region_bounds(eq.region, T) |
| off = torch.ones(T, dtype=torch.bool, device=lg.device) |
| off[lo:hi] = False |
| return QueryOutput("anchor", logits=lg.masked_fill(off, float("-inf"))) |
| assert eq.kind == "extract" |
| start, end = self.heads.span(a[None], context_hidden.float()) |
| start, end = start[0], end[0] |
| if eq.anchor >= 0: |
| bs, be = span_rel_bias(self.heads.span_bias, eq.anchor, |
| region_bounds(eq.region, T)[1], 0, T) |
| start = torch.cat([start[:T] + bs, start[T:]]) |
| end = torch.cat([end[:T] + be, end[T:]]) |
| return QueryOutput("extract", start=start, end=end) |
|
|
| def run_forks(self, state: State, row_of: Tensor, seqs: list[tuple[int, ...]], |
| pad_to: int | None = None, kinds: list[str] | None = None) -> list[Tensor]: |
| """Run forks (row_of[i] = trunk row of seqs[i]) as one right-padded batch -> per |
| fork hidden [len_i, C]. ``kinds`` (per fork) select the LoRA adapters; a kind with |
| no adapter (``generate``) runs adapter-off (= the base model).""" |
| device = self.device |
| width = pad_to or max(len(s) for s in seqs) |
| ids = torch.full((len(seqs), width), self.renderer.pad_id, dtype=torch.long, |
| device=device) |
| for i, s in enumerate(seqs): |
| ids[i, :len(s)] = torch.tensor(s, dtype=torch.long, device=device) |
| lengths = torch.tensor([len(s) for s in seqs], dtype=torch.long, device=device) |
| sel = self.lora_sel(kinds) if kinds is not None else None |
| hidden, _ = self.backbone(ids, lengths, state, rows=row_of, sel=sel) |
| return [hidden[i, :len(s)] for i, s in enumerate(seqs)] |
|
|
| def trunk(self, contexts: list[EncodedContext]) -> tuple[Tensor, State]: |
| device = self.device |
| width = max(len(c.ids) for c in contexts) |
| ids = torch.full((len(contexts), width), self.renderer.pad_id, dtype=torch.long, |
| device=device) |
| for i, c in enumerate(contexts): |
| ids[i, :len(c.ids)] = torch.tensor(c.ids, dtype=torch.long, device=device) |
| lengths = torch.tensor([len(c.ids) for c in contexts], dtype=torch.long, device=device) |
| if not any(c.media for c in contexts): |
| return self.backbone(ids, lengths, sel=self.context_sel(len(contexts))) |
| if not isinstance(self.backbone, GemmaBackbone): |
| raise ValueError(f"{self.meta.get('backbone')} reads text contexts only") |
| mask = torch.zeros(ids.shape, dtype=torch.bool, device=device) |
| feats = [] |
| for i, c in enumerate(contexts): |
| for start, emb in sorted(c.media, key=lambda m: m[0]): |
| mask[i, start:start + emb.shape[0]] = True |
| feats.append(emb.to(device)) |
| return self.backbone(ids, lengths, sel=self.context_sel(len(contexts)), |
| soft=(mask, torch.cat(feats))) |
|
|
| def answer_batched(self, state: State, trunk_hidden: Tensor, ctx: EncodedContext, |
| queries: list[EncodedQuery]) -> list[QueryOutput]: |
| """Inference path for one context: :meth:`answer_many` with a single trunk.""" |
| return self.answer_many([(state, trunk_hidden, ctx)], [(0, q) for q in queries]) |
|
|
| def answer_many(self, trunks: list[tuple[State, Tensor, EncodedContext]], |
| queries: list[tuple[int, EncodedQuery]]) -> list[QueryOutput]: |
| """Every query (``(trunk index, query)``; trunks may belong to different requests) |
| is a fork row of ONE batch (chunked only by :data:`FORK_TOKENS` padded tokens for |
| memory; a batch that hits a CUDA OOM is split in half and retried) on top of its |
| own trunk, then the heads run batched per trunk. All ops are batch-invariant |
| (:mod:`deg.kernels`), so a query's raw outputs are bitwise identical in any batch |
| composition (alone, reordered, inside a bigger schema, next to other requests).""" |
| if not queries: |
| return [] |
| device = self.device |
| budget = fork_token_budget(device) |
| route = [self.fork_route(q.kind) or "" for _, q in queries] |
| order = sorted(range(len(queries)), key=lambda i: (route[i], len(queries[i][1].ids))) |
| chunks: list[list[int]] = [[]] |
| for i in order: |
| width = len(queries[i][1].ids) |
| if chunks[-1] and (width * (len(chunks[-1]) + 1) > budget |
| or route[chunks[-1][0]] != route[i]): |
| chunks.append([]) |
| chunks[-1].append(i) |
| hidden_of: dict[int, _Picked] = {} |
| chunks.reverse() |
| while chunks: |
| chunk = chunks.pop() |
| try: |
| self._fork_chunk(trunks, queries, chunk, hidden_of) |
| except torch.OutOfMemoryError: |
| if len(chunk) == 1: |
| raise |
| split = True |
| else: |
| split = False |
| if split: |
| if device.type == "cuda": |
| torch.cuda.empty_cache() |
| chunks += [chunk[len(chunk) // 2:], chunk[:len(chunk) // 2]] |
| return self._heads_per_trunk(trunks, queries, hidden_of) |
|
|
| def _fork_chunk(self, trunks: list[tuple[State, Tensor, EncodedContext]], |
| queries: list[tuple[int, EncodedQuery]], chunk: list[int], |
| hidden_of: dict[int, _Picked]) -> None: |
| """Fork rows ``chunk`` as one batch -> ``hidden_of[i]``: only the positions the |
| heads read (answer + option tokens), so a big schema's fork hidden states are not |
| all held until the heads run.""" |
| device = self.device |
| seqs = [queries[i][1].ids for i in chunk] |
| width = max(len(x) for x in seqs) |
| ids = torch.full((len(seqs), width), self.renderer.pad_id, dtype=torch.long) |
| for r, x in enumerate(seqs): |
| ids[r, :len(x)] = torch.tensor(x, dtype=torch.long) |
| lengths = torch.tensor([len(x) for x in seqs], dtype=torch.long) |
| used = sorted({queries[i][0] for i in chunk}) |
| fs = self.backbone.fork([trunks[t][0] for t in used], |
| [used.index(queries[i][0]) for i in chunk], width) |
| sel = self.lora_sel([queries[i][1].kind for i in chunk]) |
| hid = self.backbone.fork_forward(ids.to(device), lengths.to(device), fs, sel=sel) |
| for r, i in enumerate(chunk): |
| q = queries[i][1] |
| hidden_of[i] = _Picked(hid[r], (q.answer_position, *q.option_positions)) |
| del hid, fs |
|
|
| def _heads_per_trunk(self, trunks: list[tuple[State, Tensor, EncodedContext]], |
| queries: list[tuple[int, EncodedQuery]], |
| hidden_of: dict[int, _Picked]) -> list[QueryOutput]: |
| out: list[QueryOutput | None] = [None] * len(queries) |
| for t, (_, trunk_hidden, ctx) in enumerate(trunks): |
| idx = [i for i, (tt, _) in enumerate(queries) if tt == t] |
| if not idx: |
| continue |
| ctx_hidden = trunk_hidden[0, ctx.ctx_start: ctx.ctx_start + len(ctx.offsets)] |
| for i, o in zip(idx, self.batched_heads([queries[i][1] for i in idx], |
| [hidden_of[i] for i in idx], ctx_hidden)): |
| out[i] = o |
| return [o for o in out if o is not None] |
|
|
| def batched_heads(self, queries: list[EncodedQuery], hidden: list[Any], |
| ctx_hidden: Tensor) -> list[QueryOutput]: |
| """All heads over all forks at once, fp32, through the batch-invariant GEMM: each |
| fork's logits depend only on its own hidden states (and the context).""" |
| from deg.kernels import bi_linear |
|
|
| hd = self.heads |
|
|
| def lin(m: nn.Linear, x: Tensor) -> Tensor: |
| return bi_linear(x, m.weight, m.bias) |
|
|
| out: list[QueryOutput | None] = [None] * len(queries) |
| ans = torch.stack([h[q.answer_position] for q, h in zip(queries, hidden)]).float() |
| by_kind: dict[str, list[int]] = {} |
| for i, q in enumerate(queries): |
| by_kind.setdefault(q.kind, []).append(i) |
| for kind, idx in by_kind.items(): |
| a = ans[idx] |
| if kind in ("choice", "score"): |
| head = hd.choice if kind == "choice" else hd.score |
| owner = [j for j, i in enumerate(idx) for _ in queries[i].option_positions] |
| opts = torch.stack([hidden[i][p] for i in idx |
| for p in queries[i].option_positions]).float() |
| qa = lin(head.query, a) |
| ko = lin(head.key, opts) |
| dots = bi_linear(ko, qa) |
| own = torch.tensor(owner, device=dots.device) |
| logits = (dots.gather(1, own[:, None])[:, 0] * head.scale |
| + lin(head.prior, opts)[:, 0]) |
| n0 = 0 |
| for i in idx: |
| n = len(queries[i].option_positions) |
| out[i] = QueryOutput(kind, logits=logits[n0:n0 + n]) |
| n0 += n |
| elif kind == "noul": |
| lg = lin(hd.noul, a)[:, 0] |
| for j, i in enumerate(idx): |
| out[i] = QueryOutput("noul", logits=torch.stack( |
| [lg[j], torch.zeros_like(lg[j])])) |
| elif kind == "count": |
| lg = lin(hd.count, a) |
| for j, i in enumerate(idx): |
| out[i] = QueryOutput("count", logits=lg[j]) |
| elif kind == "anchor": |
| for i, lg in zip(idx, self._anchor_logits_blocked( |
| a, ctx_hidden, [queries[i].region for i in idx])): |
| out[i] = QueryOutput("anchor", logits=lg) |
| else: |
| assert kind == "extract" |
| sp = hd.span |
| T = ctx_hidden.shape[0] |
| locs = [(queries[i].anchor, region_bounds(queries[i].region, T)[1]) |
| for i in idx] |
| if T <= HEAD_ROWS and len(idx) * (T + 1) * 8 <= HEAD_LOGIT_BYTES: |
| ctx = ctx_hidden.float() |
| start = torch.cat([bi_linear(lin(sp.q_start, a), lin(sp.k_start, ctx)) |
| * sp.scale, lin(sp.null_start, a)], -1) |
| end = torch.cat([bi_linear(lin(sp.q_end, a), lin(sp.k_end, ctx)) |
| * sp.scale, lin(sp.null_end, a)], -1) |
| for j, i in enumerate(idx): |
| st, en = start[j], end[j] |
| if locs[j][0] >= 0: |
| bs, be = span_rel_bias(hd.span_bias, *locs[j], 0, T) |
| st = torch.cat([st[:T] + bs, st[T:]]) |
| en = torch.cat([en[:T] + be, en[T:]]) |
| out[i] = QueryOutput("extract", start=st, end=en) |
| continue |
| for i, (st, en) in zip(idx, self._span_logits_blocked(a, ctx_hidden, locs)): |
| out[i] = QueryOutput("extract", start=st, end=en) |
| return [o for o in out if o is not None] |
|
|
| def _anchor_logits_blocked(self, a: Tensor, ctx_hidden: Tensor, |
| regions: list[tuple[int, int]]) -> list[Tensor]: |
| """Anchor logits [T] per answer row of ``a`` (off-region tokens = -inf), in answer |
| micro-batches of at most :data:`HEAD_LOGIT_BYTES` and :data:`HEAD_ROWS`-row context |
| blocks (keys re-projected per block); finished rows move to host memory on CUDA. |
| Every element is a fixed-config dot product of its own row and token, so the values |
| do not depend on the batch, the block size, or the other rows.""" |
| from deg.kernels import bi_linear |
|
|
| an = self.heads.anchor |
| T = ctx_hidden.shape[0] |
|
|
| def lin(m: nn.Linear, x: Tensor) -> Tensor: |
| return bi_linear(x, m.weight, m.bias) |
|
|
| qa = lin(an.query, a) |
| rows = max(1, HEAD_LOGIT_BYTES // (4 * max(1, T))) |
| res: list[Tensor] = [] |
| for q0 in range(0, a.shape[0], rows): |
| q1 = min(a.shape[0], q0 + rows) |
| lg = torch.empty(q1 - q0, T, device=a.device, dtype=torch.float32) |
| for r0 in range(0, T, HEAD_ROWS): |
| r1 = min(T, r0 + HEAD_ROWS) |
| ctx = ctx_hidden[r0:r1].float() |
| lg[:, r0:r1] = (bi_linear(qa[q0:q1], lin(an.key, ctx)) * an.scale |
| + lin(an.prior, ctx)[:, 0][None]) |
| del ctx |
| for j in range(q1 - q0): |
| lo, hi = region_bounds(regions[q0 + j], T) |
| lg[j, :lo] = float("-inf") |
| lg[j, hi:] = float("-inf") |
| if a.is_cuda: |
| lg = lg.cpu() |
| res += [lg[j] for j in range(q1 - q0)] |
| return res |
|
|
| def _span_logits_blocked(self, a: Tensor, ctx_hidden: Tensor, |
| locs: list[tuple[int, int]] | None = None |
| ) -> list[tuple[Tensor, Tensor]]: |
| """Span start/end logits [T+1] per answer row of ``a`` for a long context or a big |
| schema, never holding more than :data:`HEAD_LOGIT_BYTES` of logits or a |
| :data:`HEAD_ROWS`-row block of fp32 context on the device: answer rows run in |
| micro-batches, context keys are projected per row block (recomputed per |
| micro-batch: no [T, HEAD_DIM] key copy), and finished logits move to host memory. |
| Every element is the same fixed-config dot product as the one-shot path |
| (:func:`deg.kernels.bi_linear` is row- and column-invariant), so the values are |
| bitwise those of the one-shot path.""" |
| from deg.kernels import bi_linear |
|
|
| sp = self.heads.span |
| T = ctx_hidden.shape[0] |
|
|
| def lin(m: nn.Linear, x: Tensor) -> Tensor: |
| return bi_linear(x, m.weight, m.bias) |
|
|
| qs, qe = lin(sp.q_start, a), lin(sp.q_end, a) |
| ns, ne = lin(sp.null_start, a), lin(sp.null_end, a) |
| rows = max(1, HEAD_LOGIT_BYTES // (8 * (T + 1))) |
| host = a.is_cuda |
| res: list[tuple[Tensor, Tensor]] = [] |
| for q0 in range(0, a.shape[0], rows): |
| q1 = min(a.shape[0], q0 + rows) |
| start = torch.empty(q1 - q0, T + 1, device=a.device, dtype=torch.float32) |
| end = torch.empty_like(start) |
| for r0 in range(0, T, HEAD_ROWS): |
| r1 = min(T, r0 + HEAD_ROWS) |
| ctx = ctx_hidden[r0:r1].float() |
| start[:, r0:r1] = bi_linear(qs[q0:q1], lin(sp.k_start, ctx)) * sp.scale |
| end[:, r0:r1] = bi_linear(qe[q0:q1], lin(sp.k_end, ctx)) * sp.scale |
| del ctx |
| for j in range(q1 - q0): |
| if locs is not None and locs[q0 + j][0] >= 0: |
| bs, be = span_rel_bias(self.heads.span_bias, *locs[q0 + j], r0, r1) |
| start[j, r0:r1] += bs |
| end[j, r0:r1] += be |
| start[:, T:] = ns[q0:q1] |
| end[:, T:] = ne[q0:q1] |
| if host: |
| start, end = start.cpu(), end.cpu() |
| res += [(start[j], end[j]) for j in range(q1 - q0)] |
| return res |
|
|
| |
|
|
| def save(self, path: Path, extra: dict[str, Any] | None = None) -> None: |
| from safetensors.torch import save_file |
|
|
| path.mkdir(parents=True, exist_ok=True) |
| lm_state = {k: v.detach().to(torch.bfloat16).contiguous() |
| for k, v in self.backbone.lm.state_dict().items() |
| if k != "lm_head.weight" and ".lora_" not in k |
| and not _is_moe_key(k) |
| and not k.endswith(PLE_TABLE_KEY)} |
| save_file(lm_state, str(path / "backbone.safetensors")) |
| if self.moe is not None: |
| save_file({k: (v.detach().float() if ".lora_" in k else |
| v.detach().to(torch.bfloat16)).contiguous() |
| for k, v in moe_state(self).items()}, str(path / "moe.safetensors")) |
| if self.lora is not None: |
| save_file({k: v.detach().float().contiguous() for k, v in lora_state(self).items()}, |
| str(path / "adapters.safetensors")) |
| save_file({k: v.detach().float().contiguous() for k, v in self.heads.state_dict().items()}, |
| str(path / "heads.safetensors")) |
| backbone = self.meta.get("backbone", DEFAULT_BACKBONE) |
| meta = {**self.meta, **(extra or {}), "lora": self.lora, "moe": self.moe, |
| "kinds": list(KINDS), "max_count": MAX_COUNT, |
| "array_mode": self.array_mode, "backbone": backbone, |
| "base": str(self.meta.get("base", base_snapshot(backbone)))} |
| (path / "deg.json").write_text(json.dumps(meta, indent=2)) |
|
|
|
|
| |
| |
| PLE_TABLE_KEY = "embed_tokens_per_layer.weight" |
| GENERATE_WITH = ("self", "base", "moe") |
| |
| |
| |
| |
| GEN_LAYOUTS = ("deg", "chat") |
| |
| |
| GENERATOR_PREFIX = "base+lora:" |
| GENERATOR_FORMAT = "deg-generator-lora-v1" |
|
|
|
|
| def parse_generate(spec: str | None) -> tuple[str | None, Path | None]: |
| """``(mode, generator dir)`` of a ``generate`` option: ``"self"``, ``"base"`` or |
| ``"base+lora:<dir>"`` (mode ``"base"`` with that generator).""" |
| if spec is not None and spec.startswith(GENERATOR_PREFIX): |
| path = Path(spec[len(GENERATOR_PREFIX):]) |
| if not (path / "generator.json").is_file(): |
| raise ValueError(f"no generator.json in {path}") |
| return "base", path |
| if spec is not None and spec not in GENERATE_WITH: |
| raise ValueError(f"generate must be one of {GENERATE_WITH} or " |
| f"{GENERATOR_PREFIX}<dir>: {spec!r}") |
| return spec, None |
|
|
|
|
| def merge_generator(lm: nn.Module, path: Path | str, backbone: str | None = None |
| ) -> dict[str, Any]: |
| """Merge a GENERATE LoRA (``generator.json`` + ``lora.safetensors``, keys |
| ``<module>.lora_A`` [r, in] / ``<module>.lora_B`` [out, r] named as in ``lm``) into |
| ``lm``'s weights in place: ``W += alpha / rank * B @ A`` in fp32 on the CPU, then cast |
| to ``W``'s dtype (deterministic). ``backbone`` given: the generator must be for that |
| base (``generator.json`` ``backbone``; older generators are qwen3.5-0.8b), else |
| ValueError. Returns the generator config.""" |
| from safetensors.torch import load_file |
|
|
| path = Path(path) |
| cfg = json.loads((path / "generator.json").read_text()) |
| if cfg.get("format") != GENERATOR_FORMAT: |
| raise ValueError(f"{path}: not a {GENERATOR_FORMAT} generator") |
| gen_backbone = cfg.get("backbone", DEFAULT_BACKBONE) |
| want = backbone_spec(backbone).name if backbone is not None else gen_backbone |
| if want != gen_backbone: |
| raise ValueError(f"{path}: a generator for {gen_backbone}, not {want}") |
| ad = load_file(str(path / "lora.safetensors")) |
| scale = float(cfg["alpha"]) / int(cfg["rank"]) |
| names = sorted({k.rsplit(".lora_", 1)[0] for k in ad}) |
| if not names or len(ad) != 2 * len(names): |
| raise ValueError(f"{path}: lora.safetensors needs lora_A and lora_B per module") |
| for name in names: |
| w = lm.get_submodule(name).weight |
| delta = ad[f"{name}.lora_B"].float() @ ad[f"{name}.lora_A"].float() |
| if delta.shape != w.shape: |
| raise ValueError(f"{name}: delta {tuple(delta.shape)} vs {tuple(w.shape)}") |
| with torch.no_grad(): |
| w.copy_((w.detach().float().cpu() + scale * delta).to(w.dtype).to(w.device)) |
| return cfg |
|
|
|
|
| def add_lora(model: DegModel, rank: int, alpha: float | None = None, |
| adapters: list[str] | tuple[str, ...] = ("shared",), context: str = "base", |
| seed: int = 0) -> list[nn.Parameter]: |
| """Add named LoRA adapters (rank ``rank``, scale ``alpha / rank``; default alpha = |
| 2 rank) to every :data:`LORA_TARGETS` projection of ``model.backbone``; B starts at 0, |
| so the adapted model starts as the base. ``adapters``: ``["shared"]`` (one adapter for |
| every decide/extract fork) or a subset of ``decide``/``extract`` (by fork kind, |
| :data:`ADAPTER_OF_KIND`). ``context``: ``"base"`` (the context is read once by the |
| base weights; adapters only on fork tokens) or the adapter that also reads the |
| context. Returns the new parameters (the trainable backbone set).""" |
| names = list(adapters) |
| if names != ["shared"] and not (names and set(names) <= {"decide", "extract"} |
| and len(set(names)) == len(names)): |
| raise ValueError(f"adapters must be ['shared'] or from decide,extract: {names}") |
| if context != "base" and context not in names: |
| raise ValueError(f"lora context must be 'base' or one of {names}: {context!r}") |
| if rank < 1: |
| raise ValueError(f"lora rank must be >= 1: {rank}") |
| alpha = float(alpha if alpha is not None else 2 * rank) |
| gen = torch.Generator().manual_seed(seed) |
| n = rank * len(names) |
| params = [] |
| for layer in model.backbone.text.layers: |
| for owner_name, targets in model.backbone.fam.lora_targets.items(): |
| owner = getattr(layer, owner_name, None) |
| if owner is None: |
| continue |
| for t in targets: |
| lin = getattr(owner, t, None) |
| if lin is None: |
| continue |
| w = lin.weight |
| a = torch.randn(n, w.shape[1], generator=gen) / math.sqrt(w.shape[1]) |
| lin.lora_A = nn.Parameter(a.to(w.device, w.dtype)) |
| lin.lora_B = nn.Parameter(torch.zeros(w.shape[0], n, device=w.device, |
| dtype=w.dtype)) |
| lin.lora_scale = alpha / rank |
| params += [lin.lora_A, lin.lora_B] |
| model.lora = {"rank": int(rank), "alpha": alpha, "adapters": names, "context": context} |
| model.meta["lora"] = dict(model.lora) |
| return params |
|
|
|
|
| MOE_FFN = "mlp" |
| WRITER_TARGETS = {"linear_attn": ("in_proj_qkv", "in_proj_z", "out_proj"), |
| "self_attn": ("q_proj", "k_proj", "v_proj", "o_proj")} |
|
|
|
|
| def _is_moe_key(key: str) -> bool: |
| """A key of an MoE expert beside the dense FFN: W (``mlp_w``) or a D/E/G expert |
| (``mlp_d``, ``mlp_g``).""" |
| return re.search(rf"\.{MOE_FFN}_[a-z]\.", key) is not None |
|
|
|
|
| def add_moe(model: DegModel, base_lm: nn.Module | None = None, writer_rank: int = 32, |
| writer_alpha: float | None = None, seed: int = 0) -> list[nn.Parameter]: |
| """Make ``model`` a deg-MoE (see :class:`WriteRoute`): give every decoder layer a W |
| FFN expert ``mlp_w`` = a copy of ``base_lm``'s MLP of that layer (None: a copy of the |
| layer's own R MLP, a placeholder that :func:`load_model` overwrites from |
| ``moe.safetensors``) and, with ``writer_rank`` > 0, a writer LoRA (rank r, scale |
| ``alpha / r``, default alpha 2r, B = 0) on the attention / linear-attention projections |
| (:data:`WRITER_TARGETS`) that only W calls switch on. R (everything deg had) is not |
| touched, so decide/extract outputs are unchanged by construction. Returns the W |
| parameters (W experts + writer LoRA): the trainable set of :mod:`deg.moetrain`. |
| Qwen3.5 only; refuses a LoRA-layout model (its adapters use the same slots).""" |
| import copy |
|
|
| if model.moe is not None: |
| raise ValueError("already a deg-MoE") |
| if model.lora is not None: |
| raise ValueError("deg-MoE needs a full fine-tune checkpoint, not a LoRA one") |
| if model.backbone.fam is not QWEN35: |
| raise ValueError(f"deg-MoE supports Qwen3.5 backbones, not {model.backbone.fam.name}") |
| if writer_rank < 0: |
| raise ValueError(f"writer rank must be >= 0: {writer_rank}") |
| layers = model.backbone.text.layers |
| if base_lm is not None: |
| btext = getattr(base_lm.model, "language_model", base_lm.model) |
| if family_of(base_lm) is not QWEN35 or len(btext.layers) != len(layers): |
| raise ValueError("base model does not match the deg backbone") |
| src = [blayer.mlp for blayer in btext.layers] |
| else: |
| src = [layer.mlp for layer in layers] |
| w = layers[0].mlp.gate_proj.weight |
| params: list[nn.Parameter] = [] |
| for layer, m in zip(layers, src): |
| exp = copy.deepcopy(m) |
| exp.__dict__.pop("_deg_fused_gate_proj_up_proj", None) |
| exp.to(device=w.device, dtype=w.dtype) |
| exp.requires_grad_(False) |
| layer.mlp_w = exp |
| params += list(exp.parameters()) |
| alpha = float(writer_alpha if writer_alpha is not None else 2 * max(writer_rank, 1)) |
| if writer_rank: |
| gen = torch.Generator().manual_seed(seed) |
| for layer in layers: |
| for owner_name, targets in WRITER_TARGETS.items(): |
| owner = getattr(layer, owner_name, None) |
| if owner is None: |
| continue |
| for t in targets: |
| lin = getattr(owner, t) |
| pw = lin.weight |
| a = (torch.randn(writer_rank, pw.shape[1], generator=gen) |
| / math.sqrt(pw.shape[1])) |
| lin.lora_A = nn.Parameter(a.to(pw.device, pw.dtype), requires_grad=False) |
| lin.lora_B = nn.Parameter(torch.zeros(pw.shape[0], writer_rank, |
| device=pw.device, dtype=pw.dtype), |
| requires_grad=False) |
| lin.lora_scale = alpha / writer_rank |
| params += [lin.lora_A, lin.lora_B] |
| model.moe = {"writer_rank": int(writer_rank), "writer_alpha": alpha, |
| "writer_targets": {k: list(v) for k, v in WRITER_TARGETS.items()}} |
| model.meta["moe"] = dict(model.moe) |
| model.generate_with = "moe" |
| model.graphs.clear() |
| return params |
|
|
|
|
| TASK_EXPERTS = ("d", "g") |
|
|
|
|
| def add_task_moe(model: DegModel, experts: tuple[str, ...] | list[str] = TASK_EXPERTS |
| ) -> list[nn.Parameter]: |
| """Make ``model`` a D/E/G task MoE (:class:`TaskRoute`): every decoder layer gets FFN |
| experts ``mlp_<x>`` for x in ``experts``, copies of the layer's ``mlp`` (on a fresh |
| model: the base's FFN, so every expert starts as the base and the model's outputs are |
| unchanged), trainable. E (``mlp``) also reads the context. Routing: decide forks -> D, |
| extract / count / anchor forks and the context -> E, generate (KNOWN VALUES, header, |
| decoded tokens) -> G. Returns the new parameters. Refuses an MoE or LoRA layout.""" |
| import copy |
|
|
| if model.moe is not None or model.lora is not None: |
| raise ValueError("task MoE needs a dense full fine-tune layout") |
| if sorted(experts) != sorted(TASK_EXPERTS): |
| raise ValueError(f"task experts must be {TASK_EXPERTS}: {experts}") |
| params: list[nn.Parameter] = [] |
| for layer in model.backbone.text.layers: |
| m = getattr(layer, MOE_FFN) |
| for x in experts: |
| exp = copy.deepcopy(m) |
| for key in [k for k in exp.__dict__ if k.startswith("_deg_fused_")]: |
| exp.__dict__.pop(key) |
| setattr(layer, f"{MOE_FFN}_{x}", exp) |
| params += list(exp.parameters()) |
| model.moe = {"experts": list(experts), "reader": "e", |
| "routes": dict(TASK_EXPERT_OF_KIND)} |
| model.meta["moe"] = dict(model.moe) |
| model.generate_with = "moe" |
| model.graphs.clear() |
| return params |
|
|
|
|
| def moe_state(model: DegModel) -> dict[str, Tensor]: |
| """The W side of a deg-MoE: W experts + writer LoRA (keys as in the backbone lm).""" |
| return {k: v for k, v in model.backbone.lm.state_dict().items() |
| if _is_moe_key(k) or ".lora_" in k} |
|
|
|
|
| def moe_params(model: DegModel) -> dict[str, int]: |
| """Parameter counts: ``total`` in memory (tied embeddings once), ``r_ffn``/``w_ffn`` |
| (the experts), ``writer_lora``, and ``active`` per token (R or W path; a W token also |
| runs the writer LoRA).""" |
| seen: set[int] = set() |
| total = r = w = lora = 0 |
| for name, p in model.backbone.lm.named_parameters(): |
| if id(p) in seen: |
| continue |
| seen.add(id(p)) |
| n = p.numel() |
| total += n |
| if _is_moe_key(name): |
| w += n |
| elif ".lora_" in name: |
| lora += n |
| elif f".{MOE_FFN}." in name: |
| r += n |
| return {"total": total, "r_ffn": r, "w_ffn": w, "writer_lora": lora, |
| "active_reader": total - w - lora, "active_writer": total - r} |
|
|
|
|
| def task_moe_params(model: DegModel) -> dict[str, int]: |
| """D/E/G task MoE parameter counts (tied embeddings once): ``total`` in memory with |
| and without the PLE table (``ple``, frozen, host-able), one FFN copy per expert |
| (``ffn_e`` = ``ffn_d`` = ``ffn_g``), and ``active`` per token (one expert).""" |
| seen: set[int] = set() |
| out = {"total": 0, "ple": 0, "ffn_e": 0, "ffn_d": 0, "ffn_g": 0} |
| for name, p in model.backbone.lm.named_parameters(): |
| if id(p) in seen: |
| continue |
| seen.add(id(p)) |
| n = p.numel() |
| out["total"] += n |
| if name.endswith(PLE_TABLE_KEY): |
| out["ple"] += n |
| for x in ("d", "g"): |
| if f".{MOE_FFN}_{x}." in name: |
| out[f"ffn_{x}"] += n |
| if f".{MOE_FFN}." in name: |
| out["ffn_e"] += n |
| out["total_without_ple"] = out["total"] - out["ple"] |
| out["active_without_ple"] = out["total_without_ple"] - out["ffn_d"] - out["ffn_g"] |
| return out |
|
|
|
|
| def lora_state(model: DegModel) -> dict[str, Tensor]: |
| return {k: v for k, v in model.backbone.lm.state_dict().items() if ".lora_" in k} |
|
|
|
|
| def attach_base(model: DegModel, lm: nn.Module | None = None) -> Backbone: |
| """Load the frozen base model of ``model``'s backbone (``meta["backbone"]``: Qwen3.5 |
| or LFM2.5; ``lm`` overrides it; tests) beside ``model`` as ``model.base``. Its token |
| embeddings (= the tied LM head) are shared with ``model`` when they are equal (deg |
| trains with them frozen), so the copy costs only the body (~1 GB bf16 for the 0.8B, |
| ~2.1 GB for LFM2.5-1.2B); otherwise it keeps its own. Never trained: no grad, eval |
| mode, not in ``model.modules()`` (not saved, not optimised).""" |
| emb = model.backbone.text.embed_tokens.weight |
| if lm is None: |
| lm = load_base_lm(model.meta.get("backbone"), emb.dtype) |
| if family_of(lm) is not model.backbone.fam: |
| raise ValueError(f"base {lm.config.model_type} does not match the deg backbone " |
| f"{model.backbone.fam.name}") |
| text = getattr(lm.model, "language_model", lm.model) |
| own = text.embed_tokens.weight |
| if own.shape == emb.shape and torch.equal(own.detach().to("cpu", emb.dtype), |
| emb.detach().to("cpu")): |
| text.embed_tokens.weight = emb |
| lm.lm_head.weight = emb |
| shared = True |
| else: |
| shared = False |
| ple = getattr(model.backbone.text, "embed_tokens_per_layer", None) |
| if ple is not None: |
| bple = text.embed_tokens_per_layer |
| if bple.weight.shape != ple.weight.shape or not torch.equal( |
| bple.weight.detach().to("cpu", ple.weight.dtype), ple.weight.detach().cpu()): |
| raise ValueError("base PLE table differs from the deg model's (frozen in " |
| "deg training): not this model's base") |
| text.embed_tokens_per_layer = None |
| lm.to(device=emb.device, dtype=emb.dtype) |
| if ple is not None: |
| text.embed_tokens_per_layer = ple |
| lm.requires_grad_(False) |
| lm.eval() |
| model.base = Backbone(lm) |
| model.base.shared_embeddings = shared |
| model._gen_view = None |
| return model.base |
|
|
|
|
| def load_model(path: Path | str | None = None, *, device: str | torch.device = "cuda", |
| dtype: torch.dtype = torch.bfloat16, lm: nn.Module | None = None, |
| tokenizer: Any = None, generate: str | None = None, |
| base_lm: nn.Module | None = None, backbone: str | None = None, |
| ple_host: bool | None = None) -> DegModel: |
| """:func:`_load_model`; ``ple_host`` (default: env ``DEG_PLE_HOST=1``): a Gemma 4 |
| model keeps its frozen PLE table in host memory (:meth:`GemmaBackbone.ple_to_host`; |
| ~2.3B params off the GPU, ~4.4 GiB on E2B).""" |
| model = _load_model(path, device=device, dtype=dtype, lm=lm, tokenizer=tokenizer, |
| generate=generate, base_lm=base_lm, backbone=backbone) |
| if ple_host is None: |
| ple_host = os.environ.get("DEG_PLE_HOST", "") == "1" |
| if ple_host and isinstance(model.backbone, GemmaBackbone): |
| model.backbone.ple_to_host() |
| if model.base is not None: |
| model.base.text.embed_tokens_per_layer = model.backbone.text.embed_tokens_per_layer |
| return model |
|
|
|
|
| def _load_model(path: Path | str | None = None, *, device: str | torch.device = "cuda", |
| dtype: torch.dtype = torch.bfloat16, lm: nn.Module | None = None, |
| tokenizer: Any = None, generate: str | None = None, |
| base_lm: nn.Module | None = None, backbone: str | None = None) -> DegModel: |
| """``path=None`` -> fresh heads on the base model; else a saved deg checkpoint. |
| ``lm``/``tokenizer`` override the base (tests use a tiny random Qwen3.5 / LFM2). |
| |
| ``backbone`` (:data:`BACKBONES`; default: the checkpoint's ``meta["backbone"]``, else |
| ``qwen3.5-0.8b``) picks the base model; a checkpoint refuses a different one. With |
| ``lm`` given, the backbone is ``lm``'s architecture. |
| |
| ``generate`` (default: the checkpoint's ``meta["generate_backbone"]``, else |
| ``"self"``): ``"base"`` decodes x-generate fields with the frozen base weights |
| (:func:`attach_base`; ``base_lm`` overrides them) while decide/extract use deg's; |
| ``"base+lora:<dir>"`` is that hybrid with a GENERATE LoRA merged into the base copy |
| (:func:`merge_generator`).""" |
| from safetensors.torch import load_file |
|
|
| generate, generator = parse_generate(generate) |
| saved: dict[str, Any] | None = None |
| if path is not None: |
| path = Path(path) |
| saved = json.loads((path / "deg.json").read_text()) |
| |
| ckpt_backbone = saved.get("backbone", DEFAULT_BACKBONE) |
| if backbone is not None and backbone_spec(backbone).name != ckpt_backbone: |
| raise ValueError(f"checkpoint {path} is a {ckpt_backbone} model, not {backbone}") |
| backbone = ckpt_backbone |
| spec = spec_of_lm(lm, backbone) if lm is not None else backbone_spec(backbone) |
| base = base_snapshot(spec.name) if lm is None or tokenizer is None else None |
| if lm is None: |
| lm = load_base_lm(spec.name, dtype) |
| tokenizer = tokenizer or load_tokenizer(base) |
| meta: dict[str, Any] = {"base": str(base) if base else "custom", "backbone": spec.name} |
| backbone_mod = Backbone(lm) |
| heads = Heads(backbone_mod.hidden) |
| if path is not None: |
| assert saved is not None |
| state = load_file(str(path / "backbone.safetensors")) |
| missing, unexpected = lm.load_state_dict(state, strict=False) |
| bad = [k for k in missing if k != "lm_head.weight" |
| and not k.endswith(PLE_TABLE_KEY)] + list(unexpected) |
| if bad: |
| raise RuntimeError(f"checkpoint mismatch: {bad[:5]}") |
| missing, unexpected = heads.load_state_dict(load_file(str(path / "heads.safetensors")), |
| strict=False) |
| new_heads = [k for k in missing if k.startswith(("anchor.", "span_bias"))] |
| if unexpected or len(new_heads) != len(missing): |
| raise RuntimeError(f"heads mismatch: {(list(unexpected) + missing)[:5]}") |
| |
| meta["array_mode"] = "index" if new_heads else "anchor" |
| meta.update(saved) |
| meta["backbone"] = spec.name |
| else: |
| meta["array_mode"] = "anchor" |
| lm.to(device=device, dtype=dtype) |
| heads.to(device=device, dtype=torch.float32) |
| lm.eval() |
| heads.eval() |
| model = DegModel(backbone_mod, heads, Renderer(tokenizer, spec), meta) |
| if meta.get("moe"): |
| assert path is not None |
| cfg = meta["moe"] |
| if "experts" in cfg: |
| add_task_moe(model, cfg["experts"]) |
| else: |
| add_moe(model, None, cfg["writer_rank"], cfg["writer_alpha"]) |
| ws = load_file(str(path / "moe.safetensors")) |
| missing, unexpected = lm.load_state_dict(ws, strict=False) |
| lost = [k for k in missing if _is_moe_key(k) or ".lora_" in k] |
| if unexpected or lost: |
| raise RuntimeError(f"moe mismatch: {(list(unexpected) + lost)[:5]}") |
| lm.to(device=device, dtype=dtype) |
| lm.eval() |
| if generate not in (None, "moe") or generator is not None: |
| raise ValueError("a deg-MoE checkpoint generates with its W experts (moe)") |
| return model |
| if meta.get("lora"): |
| cfg = meta["lora"] |
| add_lora(model, cfg["rank"], cfg["alpha"], cfg["adapters"], cfg["context"]) |
| assert path is not None |
| ad = load_file(str(path / "adapters.safetensors")) |
| missing, unexpected = lm.load_state_dict(ad, strict=False) |
| if unexpected or any(".lora_" in k for k in missing): |
| raise RuntimeError(f"adapters mismatch: {(list(unexpected) + missing)[:5]}") |
| lm.to(device=device, dtype=dtype) |
| lm.eval() |
| if generate == "self" or generator is not None: |
| raise ValueError("a LoRA checkpoint always generates with the base (adapter off)") |
| return model |
| mode = generate or meta.get("generate_backbone", "self") |
| if mode == "moe": |
| raise ValueError("generate moe needs a deg-MoE checkpoint (deg.moetrain)") |
| if mode not in GENERATE_WITH: |
| raise ValueError(f"generate must be one of {GENERATE_WITH}: {mode!r}") |
| if mode == "base": |
| if generator is not None: |
| if base_lm is None: |
| base_lm = load_base_lm(spec.name, dtype) |
| cfg = merge_generator(base_lm, generator, spec.name) |
| model.meta["generator"] = {"path": str(generator), "rank": cfg["rank"], |
| "alpha": cfg["alpha"]} |
| attach_base(model, base_lm) |
| model.generate_with = "base" |
| return model |
|
|
