Buckets:
| """ThinkDecide: S1 backbone + latent thinking loop over thought slots (A) + one-shot decision head (B). | |
| Readout at initialisation is exactly S1's: per answer mask, s = z_Yes - z_No with z = lm_head(final_norm(h)). | |
| choice / score: softmax over the options' s; noul: logits [0, s] (== S1's softmax over [z_No, z_Yes]). | |
| Thinking: pass 1 sees [MASK] embeddings in the thought slots; after pass k the slots' inputs become | |
| e_mask + g * W(LN(h_k)) (h_k = last-layer pre-norm hidden state, W = I and g = 0 at init). | |
| Head: one attention block, queries = option slots, keys = option + thought slots; zero-init output, added | |
| to the option slots' pre-norm hidden state before S1's final norm and Yes/No rows. | |
| """ | |
| import math | |
| import torch | |
| import torch.nn.functional as F | |
| from torch import nn | |
| from torch.utils.checkpoint import checkpoint | |
| class DecisionHead(nn.Module): | |
| def __init__(self, d, n_heads=8): | |
| super().__init__() | |
| self.n_heads = n_heads | |
| self.norm_q = nn.LayerNorm(d) | |
| self.norm_kv = nn.LayerNorm(d) | |
| self.q = nn.Linear(d, d, bias=False) | |
| self.k = nn.Linear(d, d, bias=False) | |
| self.v = nn.Linear(d, d, bias=False) | |
| self.o = nn.Linear(d, d, bias=False) | |
| nn.init.zeros_(self.o.weight) | |
| def forward(self, h_opt, opt_mask, h_kv, kv_mask): | |
| """h_opt [B,N,d], opt_mask [B,N] bool; h_kv [B,M,d], kv_mask [B,M] bool -> delta [B,N,d].""" | |
| B, N, d = h_opt.shape | |
| M = h_kv.shape[1] | |
| H = self.n_heads | |
| q = self.q(self.norm_q(h_opt)).view(B, N, H, d // H).transpose(1, 2) | |
| kv = self.norm_kv(h_kv) | |
| k = self.k(kv).view(B, M, H, d // H).transpose(1, 2) | |
| v = self.v(kv).view(B, M, H, d // H).transpose(1, 2) | |
| att = F.scaled_dot_product_attention(q, k, v, attn_mask=kv_mask[:, None, None, :]) | |
| return self.o(att.transpose(1, 2).reshape(B, N, d)) * opt_mask[..., None] | |
| class ThinkDecide(nn.Module): | |
| def __init__(self, base, yes_id, no_id, n_think=0, use_head=True, head_heads=8, fb_mode="hidden", | |
| gate_init=0.0, gate_fixed=False, mask_id=151669): | |
| super().__init__() | |
| self.base = base # A2DQwen3LMHeadModel (bidirectional MDLM) | |
| self.bb = base.model | |
| d = base.config.hidden_size | |
| self.yes_id, self.no_id = yes_id, no_id | |
| self.n_think = n_think | |
| self.use_head = use_head | |
| self.ckpt = False | |
| self.fb_mode, self.mask_id = fb_mode, mask_id | |
| if n_think: | |
| # fb_mode "hidden" (plan): x = e_mask + g * W(LN(h)). "soft": x = e_mask + g * (E^T softmax(z) - e_mask), | |
| # z = the slot's own token logits (mask excluded): self-conditioning on the model's soft guess. | |
| self.fb_norm = nn.LayerNorm(d) | |
| self.fb_proj = nn.Linear(d, d, bias=False) | |
| nn.init.eye_(self.fb_proj.weight) | |
| self.gate = nn.Parameter(torch.full((1,), float(gate_init)), requires_grad=not gate_fixed) | |
| if fb_mode == "soft": | |
| self.fb_norm.requires_grad_(False) | |
| self.fb_proj.requires_grad_(False) | |
| if use_head: | |
| self.head = DecisionHead(d, head_heads) | |
| def compile_layers(self): | |
| """Fuse each decoder layer's elementwise ops (RMSNorm, RoPE, SwiGLU, casts); shapes are dynamic.""" | |
| import torch._inductor.config as ic | |
| ic.triton.coalesce_tiling_analysis = False # torch 2.8 bug: sympy is_constant() crash in this heuristic | |
| for layer in self.bb.layers: | |
| layer.compile(dynamic=True) | |
| # ------------------------------------------------------------------ backbone | |
| def layers(self, x, attn, pos_emb, capture=()): | |
| caps = {} | |
| for i, layer in enumerate(self.bb.layers, 1): | |
| if self.ckpt and self.training: | |
| x = checkpoint(layer, x, attn, None, None, False, None, pos_emb, use_reentrant=False) | |
| else: | |
| x = layer(x, attention_mask=attn, position_embeddings=pos_emb) | |
| if i in capture: | |
| caps[i] = x | |
| return (x, caps) if capture else x # pre-norm last-layer hidden state (+ outputs of blocks in capture) | |
| def readout(self, h_pre_opt): | |
| """S1 readout on (possibly head-adjusted) pre-norm states [..., d] -> s = z_Yes - z_No (fp32).""" | |
| hn = self.bb.norm(h_pre_opt) | |
| w = self.base.lm_head.weight | |
| rows = torch.stack([w[self.yes_id], w[self.no_id]]) | |
| with torch.autocast("cuda", enabled=False): | |
| z = F.linear(hn.to(rows.dtype), rows).float() # same op/dtype path as S1's projection | |
| return z[..., 0] - z[..., 1] | |
| def decide(self, h, batch): | |
| """h [B,L,d] pre-norm -> option scores s [B,N] (padded entries undefined; use opt_mask).""" | |
| b = torch.arange(h.shape[0], device=h.device)[:, None] | |
| h_opt = h[b, batch["opt_idx"]] # [B,N,d] | |
| if self.use_head: | |
| if self.n_think: | |
| h_kv = torch.cat([h_opt, h[b, batch["think_idx"]]], 1) | |
| kv_mask = torch.cat([batch["opt_mask"], batch["think_mask"]], 1) | |
| else: | |
| h_kv, kv_mask = h_opt, batch["opt_mask"] | |
| h_opt = h_opt + self.head(h_opt, batch["opt_mask"], h_kv, kv_mask).to(h_opt.dtype) | |
| return self.readout(h_opt) | |
| def forward(self, batch, K=1, read_at=None, capture=()): | |
| """Run K passes; return list of option-score tensors for passes in read_at (default: [K]). | |
| With capture=(layer indices,), also return {layer: hidden at option slots} from the final pass.""" | |
| read_at = sorted(set(read_at or [K])) | |
| ids, attn = batch["ids"], batch["attn"] | |
| x = self.bb.embed_tokens(ids) | |
| pos = torch.arange(ids.shape[1], device=ids.device)[None] | |
| pos_emb = self.bb.rotary_emb(x, pos) | |
| mask4 = attn[:, None, None, :] | |
| outs = [] | |
| if self.n_think: | |
| b = torch.arange(ids.shape[0], device=ids.device)[:, None] | |
| tidx = batch["think_idx"] | |
| base_t = x[b, tidx] # [B,T,d] mask embeddings | |
| caps = {} | |
| for k in range(1, K + 1): | |
| if capture and k == K: | |
| h, caps = self.layers(x, mask4, pos_emb, capture) | |
| else: | |
| h = self.layers(x, mask4, pos_emb) | |
| if k in read_at: | |
| outs.append(self.decide(h, batch)) | |
| if k < K and self.n_think: | |
| if self.fb_mode == "soft": | |
| E = self.bb.embed_tokens.weight | |
| z = F.linear(self.bb.norm(h[b, tidx]), E).float() # tied lm_head | |
| z[..., self.mask_id] = float("-inf") | |
| fb = (z.softmax(-1).to(E.dtype) @ E) - base_t | |
| else: | |
| fb = self.fb_proj(self.fb_norm(h[b, tidx])) | |
| new_t = base_t + self.gate.to(fb.dtype) * fb | |
| x = x.index_put((b.expand_as(tidx), tidx), new_t.to(x.dtype)) | |
| if capture: | |
| bo = torch.arange(ids.shape[0], device=ids.device)[:, None] | |
| return outs, {i: v[bo, batch["opt_idx"]] for i, v in caps.items()} | |
| return outs | |
| def param_groups(self, lr_backbone, lr_new, wd): | |
| new, new_nodecay, bb_decay, bb_nodecay = [], [], [], [] | |
| for n, p in self.named_parameters(): | |
| if not p.requires_grad: | |
| continue | |
| is_new = not n.startswith("base.") | |
| nodecay = p.ndim < 2 or "norm" in n or n.endswith("gate") | |
| (new_nodecay if is_new and nodecay else new if is_new else | |
| bb_nodecay if nodecay else bb_decay).append(p) | |
| groups = [dict(params=bb_decay, lr=lr_backbone, weight_decay=wd, name="bb"), | |
| dict(params=bb_nodecay, lr=lr_backbone, weight_decay=0.0, name="bb_nd"), | |
| dict(params=new, lr=lr_new, weight_decay=wd, name="new"), | |
| dict(params=new_nodecay, lr=lr_new, weight_decay=0.0, name="new_nd")] | |
| return [g for g in groups if g["params"]] | |
| def collate(items, device, pad_id): | |
| """Right-pad ids; gather option and thought indices (padded with index 0, masked).""" | |
| B = len(items) | |
| L = max(len(x["ids"]) for x in items) | |
| N = max(len(x["pos"]) for x in items) | |
| T = max((len(x["think"]) for x in items), default=0) | |
| ids = torch.full((B, L), pad_id, dtype=torch.long) | |
| attn = torch.zeros((B, L), dtype=torch.bool) | |
| opt_idx = torch.zeros((B, N), dtype=torch.long) | |
| opt_mask = torch.zeros((B, N), dtype=torch.bool) | |
| think_idx = torch.zeros((B, max(T, 1)), dtype=torch.long) | |
| think_mask = torch.zeros((B, max(T, 1)), dtype=torch.bool) | |
| for i, x in enumerate(items): | |
| n = len(x["ids"]) | |
| ids[i, :n] = torch.tensor(x["ids"]) | |
| attn[i, :n] = True | |
| opt_idx[i, :len(x["pos"])] = torch.tensor(x["pos"]) | |
| opt_mask[i, :len(x["pos"])] = True | |
| if x["think"]: | |
| think_idx[i, :len(x["think"])] = torch.tensor(x["think"]) | |
| think_mask[i, :len(x["think"])] = True | |
| out = dict(ids=ids, attn=attn, opt_idx=opt_idx, opt_mask=opt_mask, think_idx=think_idx, think_mask=think_mask) | |
| out = {k: v.to(device, non_blocking=True) for k, v in out.items()} | |
| out["primitive"] = [x["primitive"] for x in items] | |
| out["n_opt"] = [len(x["pos"]) for x in items] | |
| return out | |
| def item_logits(s, batch): | |
| """Option scores [B,N] -> list of per-item logits (noul -> [0, s]).""" | |
| out = [] | |
| for i, (prim, n) in enumerate(zip(batch["primitive"], batch["n_opt"])): | |
| if prim == "noul": | |
| out.append(torch.stack([s.new_zeros(()), s[i, 0]])[None]) | |
| else: | |
| out.append(s[i, :n][None]) | |
| return out | |
| def load_backbone(path, dtype, device): | |
| from transformers import AutoModelForMaskedLM | |
| base = AutoModelForMaskedLM.from_pretrained(path, trust_remote_code=True, local_files_only=True, | |
| torch_dtype=dtype) | |
| base.config.use_cache = False | |
| return base.to(device) | |
Xet Storage Details
- Size:
- 10 kB
- Xet hash:
- 8203823aaa140f7b066018ae449d0a773120361aa808572fa33b8688222cd39c
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.