File size: 10,526 Bytes
287ca83 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 | """The decision model around a Hugging Face causal LM: two-way reading and a code-row readout.
**Two-way reading.** With ``attention="two_way"`` the full-attention layers see the whole prompt
in both directions, while the linear-attention layers (Qwen3.5's Gated DeltaNet) and any
sliding-window layers keep the model's own masks.
**Readout.** The hidden state at each question's readout position is scored against one row per
option code: the LM head's own rows of the code tokens (``readout="lm_head"``), or a separate
``[n_codes, d]`` matrix (``readout="head"``). The logits are normalised over the question's
options only.
"""
from __future__ import annotations
from dataclasses import dataclass
from typing import Any, Sequence
import torch
from torch import nn
from glance_runtime.layout import Example
def _core(lm: nn.Module) -> nn.Module:
"""The decoder stack (the module that builds attention masks); PEFT wrappers are unwrapped."""
if hasattr(lm, "get_base_model"):
lm = lm.get_base_model()
for attr in ("model", "language_model", "transformer"):
m = getattr(lm, attr, None)
if isinstance(m, nn.Module):
inner = getattr(m, "language_model", None)
return inner if isinstance(inner, nn.Module) else m
return lm
def _multimodal(lm: nn.Module) -> nn.Module | None:
"""The module holding the vision encoder and M-RoPE indexing (Qwen3.5's ``model``), if any."""
if hasattr(lm, "get_base_model"):
lm = lm.get_base_model()
for m in (getattr(lm, "model", None), lm):
if isinstance(m, nn.Module) and hasattr(m, "get_image_features") and hasattr(m, "get_rope_index") \
and hasattr(m, "get_placeholder_mask"):
return m
return None
def _text_config(lm: nn.Module):
cfg = lm.config
return getattr(cfg, "text_config", None) or cfg
def two_way_masks(lm: nn.Module, inputs_embeds: torch.Tensor, attention_mask: torch.Tensor) -> Any:
"""Bidirectional masks for full-attention layers; the model's own masks for the others."""
from transformers import masking_utils as mu
cfg = _core(lm).config
kw = {"config": cfg, "inputs_embeds": inputs_embeds, "attention_mask": attention_mask, "past_key_values": None}
full = mu.create_bidirectional_mask(**kw, allow_is_bidirectional_skip=False)
layer_types = getattr(cfg, "layer_types", None)
if not layer_types:
return full
masks: dict[str, Any] = {}
pos = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device)[None].expand(inputs_embeds.shape[0], -1)
for t in set(layer_types):
if t == "full_attention":
masks[t] = full
elif t == "linear_attention":
masks[t] = mu.create_recurrent_attention_mask(**kw)
elif t == "sliding_attention":
masks[t] = mu.create_sliding_window_causal_mask(**kw, position_ids=pos)
else:
masks[t] = mu.create_causal_mask(**kw, position_ids=pos)
return masks
def varlen_masks(lm: nn.Module, inputs_embeds: torch.Tensor, attention_mask: torch.Tensor, causal: bool) -> dict:
"""Masks for ``attn_impl="varlen"``: packing indices for full-attention layers (two-way or causal),
the model's own masks for the others."""
from transformers import masking_utils as mu
from glance_runtime.attention import VarlenMask
cfg = _core(lm).config
kw = {"config": cfg, "inputs_embeds": inputs_embeds, "attention_mask": attention_mask, "past_key_values": None}
pos = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device)[None].expand(inputs_embeds.shape[0], -1)
masks: dict[str, Any] = {}
for t in set(getattr(cfg, "layer_types", None) or ["full_attention"]):
if t == "full_attention":
masks[t] = VarlenMask.from_padding(attention_mask, causal=causal)
elif t == "linear_attention":
masks[t] = mu.create_recurrent_attention_mask(**kw)
else:
raise ValueError(f"attn_impl varlen does not handle {t!r} layers")
return masks
def set_full_attention_causal(lm: nn.Module, causal: bool) -> int:
"""Flip ``is_causal`` on full-attention modules (SDPA and flash paths read it). Returns how many."""
n = 0
for name, mod in lm.named_modules():
cls = type(mod).__name__
if cls.endswith("Attention") and "Vision" not in cls and hasattr(mod, "is_causal") and \
getattr(mod, "layer_type", "full_attention") != "sliding_attention":
mod.is_causal = causal
n += 1
return n
@dataclass
class Batch:
input_ids: torch.Tensor # (B, T)
attention_mask: torch.Tensor # (B, T)
readout_row: torch.Tensor # (R,) batch row of each readout
readout_pos: torch.Tensor # (R,) position of each readout
option_code: torch.Tensor # (N,) code index of each (readout, option) pair, flattened
option_readout: torch.Tensor # (N,) readout index of each pair
sizes: list[int] # options per readout
readouts: list # layout.Readout objects, in order
pixel_values: torch.Tensor | None = None # every image's patches, in example then image order
image_grid_thw: torch.Tensor | None = None # (n_images, 3)
def to(self, device: Any) -> "Batch":
"""Move the tensors the forward pass reads (indices stay where ``forward`` needs them)."""
self.input_ids = self.input_ids.to(device)
self.attention_mask = self.attention_mask.to(device)
if self.pixel_values is not None:
self.pixel_values = self.pixel_values.to(device)
self.image_grid_thw = self.image_grid_thw.to(device)
return self
def collate(examples: Sequence[Example], code_index: dict[str, int], pad_id: int = 0) -> Batch:
t = max(len(e.input_ids) for e in examples)
ids = torch.full((len(examples), t), pad_id, dtype=torch.long)
mask = torch.zeros((len(examples), t), dtype=torch.long)
rows, poss, codes, owners, sizes, ros = [], [], [], [], [], []
for b, e in enumerate(examples):
ids[b, : len(e.input_ids)] = torch.tensor(e.input_ids)
mask[b, : len(e.input_ids)] = 1
for ro in e.readouts:
r = len(ros)
rows.append(b)
poss.append(ro.position)
codes.extend(code_index[c] for c in ro.codes)
owners.extend([r] * len(ro.codes))
sizes.append(len(ro.codes))
ros.append(ro)
vision = [v for e in examples for v in e.vision]
pixels = torch.cat([v["pixel_values"] for v in vision]) if vision else None
grids = torch.cat([v["image_grid_thw"] for v in vision]) if vision else None
return Batch(ids, mask, torch.tensor(rows), torch.tensor(poss), torch.tensor(codes), torch.tensor(owners), sizes, ros,
pixels, grids)
class DecisionModel(nn.Module):
def __init__(self, lm: nn.Module, code_token_ids: Sequence[int], codes: Sequence[str], *,
attention: str = "two_way", readout: str = "lm_head", attn_impl: str = "sdpa"):
super().__init__()
if attention not in ("two_way", "causal"):
raise ValueError("attention must be two_way or causal")
if attn_impl not in ("sdpa", "varlen"):
raise ValueError("attn_impl must be sdpa or varlen")
if readout not in ("lm_head", "head"):
raise ValueError("readout must be lm_head or head")
self.lm = lm
self.attention = attention
self.readout = readout
self.codes = list(codes)
self.code_index = {c: i for i, c in enumerate(self.codes)}
self.register_buffer("code_token_ids", torch.tensor(list(code_token_ids), dtype=torch.long), persistent=False)
rows = self.lm_head_weight()[self.code_token_ids].detach().clone()
self.head = nn.Parameter(rows) if readout == "head" else None
self.image_encoder: Any = None # a glance_runtime.vision.ImageEncoder for vision bases
set_full_attention_causal(lm, attention == "causal")
self.attn_impl = attn_impl
if attn_impl == "varlen": # unpadded full attention on the text model only (glance_runtime.attention)
from glance_runtime import attention as att
att.use(_text_config(lm))
def lm_head_weight(self) -> torch.Tensor:
return self.lm.get_output_embeddings().weight
def code_rows(self) -> torch.Tensor:
return self.head if self.head is not None else self.lm_head_weight()[self.code_token_ids]
def hidden(self, input_ids: torch.Tensor, attention_mask: torch.Tensor, pixel_values: torch.Tensor | None = None,
image_grid_thw: torch.Tensor | None = None) -> torch.Tensor:
core = _core(self.lm)
embeds = core.get_input_embeddings()(input_ids)
position_ids = None
if pixel_values is not None:
mm = _multimodal(self.lm)
if mm is None:
raise ValueError("pixel_values given to a model without a vision encoder")
feats = mm.get_image_features(pixel_values, image_grid_thw, return_dict=True).pooler_output
feats = torch.cat(list(feats), 0).to(embeds.device, embeds.dtype)
image_mask, _ = mm.get_placeholder_mask(input_ids, inputs_embeds=embeds, image_features=feats)
embeds = embeds.masked_scatter(image_mask, feats)
types = (input_ids == mm.config.image_token_id).int()
position_ids, _ = mm.get_rope_index(input_ids, mm_token_type_ids=types, image_grid_thw=image_grid_thw,
attention_mask=attention_mask)
mask: Any = attention_mask
if self.attn_impl == "varlen":
mask = varlen_masks(self.lm, embeds, attention_mask, causal=self.attention == "causal")
elif self.attention == "two_way":
mask = two_way_masks(self.lm, embeds, attention_mask)
out = core(inputs_embeds=embeds, attention_mask=mask, position_ids=position_ids, use_cache=False)
return out.last_hidden_state
def forward(self, batch: Batch) -> list[torch.Tensor]:
"""One logit vector per readout, over that question's options in display order."""
h = self.hidden(batch.input_ids, batch.attention_mask, batch.pixel_values, batch.image_grid_thw)
hr = h[batch.readout_row, batch.readout_pos] # (R, d)
w = self.code_rows()[batch.option_code] # (N, d)
flat = (hr[batch.option_readout] * w.to(hr.dtype)).sum(-1)
return list(torch.split(flat, batch.sizes))
|