| """Groot 1 — Tree-Reduce Decision Transformer. |
| |
| Forward pass (one request, Q questions): |
| |
| chunks = split(state) # parallel leaves |
| H_c = Encoder(chunks) # one batched pass, shared weights |
| leaves = mean_pool(H_c) + seg_emb # one vector per chunk |
| nodes = TreeReduce(leaves) # fork/join merge, log depth |
| memory = [H_c tokens ; nodes] # state memory, built ONCE |
| for each option of each question (parallel, batched): |
| h_i = Encoder([CLS] head [SEP] [MASK] option_i [SEP])[MASK] # pointwise |
| for each question (parallel): |
| S = [q_token ; h_1 .. h_K] # a SET: identical positions |
| S = Fusion(S, memory) # set self-attn + cross-attn + GeGLU |
| logits = scorer(S[1:]) # one logit per option |
| p = softmax(logits / T[qtype]) |
| |
| The state is encoded once and every question reads it through cross-attention, |
| so the latency of a request grows only with the (small) question blocks. The |
| same dataflow is a fork/join tree, which is exactly what Bend parallelises |
| (see bend/groot.bend). |
| """ |
| from __future__ import annotations |
|
|
| import json |
| import math |
| from dataclasses import dataclass |
| from pathlib import Path |
| from typing import Any, Sequence |
|
|
| import torch |
| import torch.nn.functional as F |
| from safetensors.torch import load_file, save_file |
| from torch import nn |
|
|
| from .config import GrootConfig, QTYPES |
| from .serialize import EncodedQuestion, EncodedState |
|
|
| NEG = -1e4 |
|
|
|
|
| |
| def rope_cos_sin(position_ids: torch.Tensor, dim: int, theta: float, dtype: torch.dtype): |
| """cos/sin tables for rotary embeddings; position_ids (B, L) -> (B, L, dim).""" |
| inv = 1.0 / (theta ** (torch.arange(0, dim, 2, device=position_ids.device, dtype=torch.float32) / dim)) |
| freqs = position_ids[..., None].float() * inv |
| emb = torch.cat((freqs, freqs), dim=-1) |
| return emb.cos().to(dtype), emb.sin().to(dtype) |
|
|
|
|
| def rotate_half(x: torch.Tensor) -> torch.Tensor: |
| x1, x2 = x.chunk(2, dim=-1) |
| return torch.cat((-x2, x1), dim=-1) |
|
|
|
|
| def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: |
| |
| return x * cos[:, None] + rotate_half(x) * sin[:, None] |
|
|
|
|
| |
| class SelfAttention(nn.Module): |
| def __init__(self, d: int, heads: int, dropout: float): |
| super().__init__() |
| self.h, self.dh = heads, d // heads |
| self.Wqkv = nn.Linear(d, 3 * d, bias=False) |
| self.Wo = nn.Linear(d, d, bias=False) |
| self.dropout = dropout |
|
|
| def forward(self, x, cos, sin, mask): |
| B, L, d = x.shape |
| q, k, v = self.Wqkv(x).view(B, L, 3, self.h, self.dh).permute(2, 0, 3, 1, 4) |
| q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin) |
| out = F.scaled_dot_product_attention( |
| q, k, v, attn_mask=mask[:, None, None, :], dropout_p=self.dropout if self.training else 0.0 |
| ) |
| return self.Wo(out.transpose(1, 2).reshape(B, L, d)) |
|
|
|
|
| class CrossAttention(nn.Module): |
| def __init__(self, d: int, heads: int, dropout: float): |
| super().__init__() |
| self.h, self.dh = heads, d // heads |
| self.Wq = nn.Linear(d, d, bias=False) |
| self.Wkv = nn.Linear(d, 2 * d, bias=False) |
| self.Wo = nn.Linear(d, d, bias=False) |
| nn.init.zeros_(self.Wo.weight) |
| self.dropout = dropout |
|
|
| def forward(self, x, memory, memory_mask): |
| B, L, d = x.shape |
| M = memory.shape[1] |
| q = self.Wq(x).view(B, L, self.h, self.dh).transpose(1, 2) |
| k, v = self.Wkv(memory).view(B, M, 2, self.h, self.dh).permute(2, 0, 3, 1, 4) |
| out = F.scaled_dot_product_attention( |
| q, k, v, attn_mask=memory_mask[:, None, None, :], dropout_p=self.dropout if self.training else 0.0 |
| ) |
| return self.Wo(out.transpose(1, 2).reshape(B, L, d)) |
|
|
|
|
| class GeGLU(nn.Module): |
| def __init__(self, d: int, inter: int, dropout: float): |
| super().__init__() |
| self.Wi = nn.Linear(d, 2 * inter, bias=False) |
| self.Wo = nn.Linear(inter, d, bias=False) |
| self.drop = nn.Dropout(dropout) |
|
|
| def forward(self, x): |
| a, g = self.Wi(x).chunk(2, dim=-1) |
| return self.Wo(self.drop(F.gelu(a) * g)) |
|
|
|
|
| class FusionLayer(nn.Module): |
| """Pre-norm block: self-attention (RoPE) -> cross-attention to memory -> GeGLU.""" |
|
|
| def __init__(self, cfg: GrootConfig): |
| super().__init__() |
| d = cfg.hidden_size |
| self.attn_norm = nn.LayerNorm(d, eps=cfg.norm_eps, bias=False) |
| self.attn = SelfAttention(d, cfg.fusion_heads, cfg.dropout) |
| self.cross_norm = nn.LayerNorm(d, eps=cfg.norm_eps, bias=False) |
| self.cross = CrossAttention(d, cfg.fusion_heads, cfg.dropout) |
| self.mlp_norm = nn.LayerNorm(d, eps=cfg.norm_eps, bias=False) |
| self.mlp = GeGLU(d, cfg.fusion_intermediate, cfg.dropout) |
| self.drop = nn.Dropout(cfg.dropout) |
|
|
| def forward(self, x, cos, sin, mask, memory, memory_mask): |
| x = x + self.drop(self.attn(self.attn_norm(x), cos, sin, mask)) |
| x = x + self.drop(self.cross(self.cross_norm(x), memory, memory_mask)) |
| x = x + self.drop(self.mlp(self.mlp_norm(x))) |
| return x |
|
|
|
|
| |
| class TreeReducer(nn.Module): |
| """Balanced binary fork/join merge of chunk summaries. |
| |
| node = LN(l + r + W [l ; r ; l*r ; |l-r|]) + level_emb[level] |
| A node with a single valid child passes that child through unchanged. |
| Returns every internal node (multi-scale memory) and the root. |
| """ |
|
|
| def __init__(self, cfg: GrootConfig): |
| super().__init__() |
| d = cfg.hidden_size |
| self.merge = nn.Linear(4 * d, d) |
| if cfg.tree_merge_zero_init: |
| nn.init.zeros_(self.merge.weight) |
| nn.init.zeros_(self.merge.bias) |
| self.norm = nn.LayerNorm(d, eps=cfg.norm_eps, bias=False) |
| self.level_emb = nn.Embedding(cfg.tree_levels, d) |
| nn.init.normal_(self.level_emb.weight, std=0.02) |
| self.levels = cfg.tree_levels |
|
|
| def forward(self, leaves: torch.Tensor, mask: torch.Tensor): |
| B, m, d = leaves.shape |
| P = 1 << max(1, (m - 1).bit_length()) |
| x = F.pad(leaves, (0, 0, 0, P - m)) |
| valid = F.pad(mask, (0, P - m), value=False) |
| nodes, node_masks = [], [] |
| level = 0 |
| while x.shape[1] > 1: |
| l, r = x[:, 0::2], x[:, 1::2] |
| vl, vr = valid[:, 0::2], valid[:, 1::2] |
| merged = self.norm(l + r + self.merge(torch.cat((l, r, l * r, (l - r).abs()), dim=-1))) |
| merged = merged + self.level_emb.weight[min(level, self.levels - 1)] |
| both = (vl & vr)[..., None] |
| only_l = (vl & ~vr)[..., None] |
| x = torch.where(both, merged, torch.where(only_l, l, r)) |
| valid = vl | vr |
| keep = (vl & vr) |
| nodes.append(x) |
| node_masks.append(keep) |
| level += 1 |
| root = x[:, 0] |
| if nodes: |
| nodes_t = torch.cat(nodes, dim=1) |
| masks_t = torch.cat(node_masks, dim=1) |
| else: |
| nodes_t = leaves[:, :0] |
| masks_t = mask[:, :0] |
| return nodes_t, masks_t, root |
|
|
|
|
| |
| @dataclass |
| class Batch: |
| """Padded tensors for a batch of states and questions (one encoder sequence per question).""" |
|
|
| chunk_ids: torch.Tensor |
| chunk_mask: torch.Tensor |
| chunk_owner: torch.Tensor |
| chunk_index: torch.Tensor |
| q_ids: torch.Tensor |
| q_mask: torch.Tensor |
| q_pos: torch.Tensor |
| q_markers: torch.Tensor |
| q_block: torch.Tensor |
| q_state: torch.Tensor |
| q_type: torch.Tensor |
| q_count: torch.Tensor |
|
|
| def to(self, device): |
| return Batch(**{k: v.to(device) for k, v in self.__dict__.items()}) |
|
|
| @property |
| def n_states(self) -> int: |
| return int(self.chunk_owner.max()) + 1 |
|
|
| @property |
| def n_questions(self) -> int: |
| return int(self.q_type.shape[0]) |
|
|
|
|
| def _pad(seqs, L, value): |
| out = torch.full((len(seqs), L), value, dtype=torch.long) |
| for i, s in enumerate(seqs): |
| out[i, : len(s)] = torch.tensor(s, dtype=torch.long) |
| return out |
|
|
|
|
| def collate(states: Sequence[EncodedState], questions: Sequence[EncodedQuestion], pad_id: int, q_state: Sequence[int] | None = None) -> Batch: |
| """Build a Batch; `q_state[i]` is the state index of question i (default: question i -> state i).""" |
| if q_state is None: |
| q_state = list(range(len(questions))) |
| chunk_ids, chunk_owner, chunk_index = [], [], [] |
| for s, st in enumerate(states): |
| for i, c in enumerate(st.chunks): |
| chunk_ids.append(c) |
| chunk_owner.append(s) |
| chunk_index.append(i) |
| cids = _pad(chunk_ids, max(len(c) for c in chunk_ids), pad_id) |
| Lq = max(len(q.ids) for q in questions) |
| Lq = ((Lq + 7) // 8) * 8 |
| qids = _pad([q.ids for q in questions], Lq, pad_id) |
| qmask = torch.zeros_like(qids, dtype=torch.bool) |
| for i, q in enumerate(questions): |
| qmask[i, : len(q.ids)] = True |
| qpos = _pad([q.position_ids for q in questions], Lq, 0) |
| qblock = _pad([q.blocks if q.blocks is not None else [0] * len(q.ids) for q in questions], Lq, -1) |
| K = max(q.n_options for q in questions) |
| return Batch( |
| chunk_ids=cids, |
| chunk_mask=cids.ne(pad_id), |
| chunk_owner=torch.tensor(chunk_owner), |
| chunk_index=torch.tensor(chunk_index), |
| q_ids=qids, |
| q_mask=qmask, |
| q_pos=qpos, |
| q_markers=_pad([q.markers for q in questions], K, 0), |
| q_block=qblock, |
| q_state=torch.tensor(list(q_state)), |
| q_type=torch.tensor([q.qtype for q in questions]), |
| q_count=torch.tensor([q.n_options for q in questions]), |
| ) |
|
|
|
|
| class GrootDecisionModel(nn.Module): |
| def __init__(self, config: GrootConfig, encoder: nn.Module): |
| super().__init__() |
| self.cfg = config |
| self.encoder = encoder |
| d = config.hidden_size |
| if getattr(encoder.config, "hidden_size", d) != d: |
| raise ValueError("config.hidden_size must match the backbone hidden size") |
| if config.encoder_layers is not None and hasattr(encoder, "layers"): |
| encoder.layers = encoder.layers[: config.encoder_layers] |
| encoder.config.num_hidden_layers = config.encoder_layers |
| self.seg_emb = nn.Embedding(config.max_chunks, d) |
| nn.init.normal_(self.seg_emb.weight, std=0.02) |
| |
| |
| self.chunk_pos_scale = nn.Parameter(torch.tensor(0.02)) |
| self.qtype_emb = nn.Embedding(len(config.qtypes), d) |
| nn.init.zeros_(self.qtype_emb.weight) |
| self.level_emb = nn.Embedding(256, d) |
| nn.init.normal_(self.level_emb.weight, std=0.02) |
| self.tree = TreeReducer(config) |
| self.fusion = nn.ModuleList([FusionLayer(config) for _ in range(config.fusion_layers)]) |
| self.final_norm = nn.LayerNorm(d, eps=config.norm_eps, bias=False) |
| dh = config.scorer_hidden or d |
| self.scorer = nn.Sequential(nn.LayerNorm(d, eps=config.norm_eps), nn.Linear(d, dh), nn.GELU(), nn.Linear(dh, 1)) |
| self.conf_head = nn.Sequential(nn.Linear(2 * d + 4, config.conf_hidden), nn.GELU(), nn.Linear(config.conf_hidden, 1)) |
| self.register_buffer("temperature", torch.full((len(config.qtypes),), float(config.temperature_init))) |
| if config.fusion_init_from_encoder: |
| self._init_fusion_from_encoder() |
|
|
| |
| @classmethod |
| def from_backbone(cls, config: GrootConfig, backbone: str | None = None) -> "GrootDecisionModel": |
| from transformers import AutoModel |
|
|
| kw = {} |
| if config.global_attention: |
| from transformers import AutoConfig |
|
|
| n = AutoConfig.from_pretrained(backbone or config.backbone).num_hidden_layers |
| kw["layer_types"] = ["full_attention"] * n |
| encoder = AutoModel.from_pretrained(backbone or config.backbone, attn_implementation="sdpa", **kw) |
| encoder.config.reference_compile = False |
| return cls(config, encoder) |
|
|
| @classmethod |
| def tiny(cls, config: GrootConfig, vocab_size: int = 1024, layers: int = 2) -> "GrootDecisionModel": |
| """A randomly initialised miniature (for tests and the Bend parity harness).""" |
| from transformers import AutoConfig, AutoModel |
|
|
| enc_cfg = AutoConfig.for_model( |
| "modernbert", |
| vocab_size=vocab_size, |
| hidden_size=config.hidden_size, |
| intermediate_size=config.fusion_intermediate, |
| num_hidden_layers=layers, |
| num_attention_heads=config.fusion_heads, |
| max_position_embeddings=8192, |
| global_attn_every_n_layers=3, |
| local_attention=128, |
| global_rope_theta=config.rope_theta, |
| local_rope_theta=config.rope_theta, |
| pad_token_id=0, |
| cls_token_id=2, |
| sep_token_id=1, |
| mask_token_id=4, |
| **({"layer_types": ["full_attention"] * layers} if config.global_attention else {}), |
| ) |
| enc_cfg.reference_compile = False |
| encoder = AutoModel.from_config(enc_cfg, attn_implementation="sdpa") |
| cfg = GrootConfig.from_dict({**config.to_dict(), "encoder_layers": None}) |
| return cls(cfg, encoder) |
|
|
| def _init_fusion_from_encoder(self) -> None: |
| layers = getattr(self.encoder, "layers", None) |
| if not layers: |
| return |
| src = list(layers)[-len(self.fusion) :] |
| for fl, el in zip(self.fusion, src): |
| try: |
| fl.attn.Wqkv.weight.data.copy_(el.attn.Wqkv.weight.data) |
| fl.attn.Wo.weight.data.copy_(el.attn.Wo.weight.data) |
| fl.mlp.Wi.weight.data.copy_(el.mlp.Wi.weight.data) |
| fl.mlp.Wo.weight.data.copy_(el.mlp.Wo.weight.data) |
| if hasattr(el.attn_norm, "weight"): |
| fl.attn_norm.weight.data.copy_(el.attn_norm.weight.data) |
| fl.mlp_norm.weight.data.copy_(el.mlp_norm.weight.data) |
| except (AttributeError, RuntimeError): |
| break |
|
|
| @property |
| def device(self) -> torch.device: |
| return self.temperature.device |
|
|
| |
| def sync_chunk_pos_trainable(self) -> None: |
| """Freeze the parameters of the chunk-position scheme that is not in use (DDP needs every trainable param used).""" |
| sinus = getattr(self.cfg, "chunk_pos", "learned") == "sinusoidal" |
| self.seg_emb.weight.requires_grad_(not sinus) |
| self.chunk_pos_scale.requires_grad_(sinus) |
|
|
| def chunk_position(self, idx: torch.Tensor) -> torch.Tensor: |
| if getattr(self.cfg, "chunk_pos", "learned") != "sinusoidal": |
| return self.seg_emb(idx.clamp(max=self.seg_emb.num_embeddings - 1)) |
| d = self.cfg.hidden_size |
| inv = torch.exp(torch.arange(0, d, 2, device=idx.device, dtype=torch.float32) * (-torch.log(torch.tensor(10000.0)) / d)) |
| ang = idx.float()[:, None] * inv[None, :] |
| pe = torch.zeros(idx.shape[0], d, device=idx.device) |
| pe[:, 0::2], pe[:, 1::2] = torch.sin(ang), torch.cos(ang) |
| return (pe * self.chunk_pos_scale).to(self.seg_emb.weight.dtype) |
|
|
|
|
| def encode_chunks(self, chunk_ids, chunk_mask): |
| return self.encoder(input_ids=chunk_ids, attention_mask=chunk_mask.long()).last_hidden_state |
|
|
| def build_memory(self, batch: Batch, n_examples: int): |
| """State memory (tokens + tree nodes) for every example in the batch.""" |
| H = self.encode_chunks(batch.chunk_ids, batch.chunk_mask) |
| H = H + self.chunk_position(batch.chunk_index)[:, None, :] |
| C, Lc, d = H.shape |
| |
| counts = torch.bincount(batch.chunk_owner, minlength=n_examples) |
| m = int(counts.max()) |
| |
| denom = batch.chunk_mask.sum(-1, keepdim=True).clamp_min(1) |
| leaf = (H * batch.chunk_mask[..., None]).sum(1) / denom |
| leaves = H.new_zeros(n_examples, m, d) |
| leaf_mask = torch.zeros(n_examples, m, dtype=torch.bool, device=H.device) |
| leaves[batch.chunk_owner, batch.chunk_index] = leaf |
| leaf_mask[batch.chunk_owner, batch.chunk_index] = True |
| nodes, node_mask, root = self.tree(leaves, leaf_mask) |
| |
| tokens = H.new_zeros(n_examples, m, Lc, d) |
| token_mask = torch.zeros(n_examples, m, Lc, dtype=torch.bool, device=H.device) |
| tokens[batch.chunk_owner, batch.chunk_index] = H |
| token_mask[batch.chunk_owner, batch.chunk_index] = batch.chunk_mask |
| memory = torch.cat((tokens.view(n_examples, m * Lc, d), nodes), dim=1) |
| memory_mask = torch.cat((token_mask.view(n_examples, m * Lc), node_mask), dim=1) |
| return memory, memory_mask, root |
|
|
| |
| def encode_questions(self, batch: Batch): |
| """Early fusion: each question sequence (head + state head + option blocks) through the encoder.""" |
| blk = batch.q_block |
| |
| allowed = (blk[:, None, :] >= 0) & ((blk[:, None, :] == 0) | (blk[:, None, :] == blk[:, :, None])) |
| mask4d = allowed[:, None] |
| masks = {"full_attention": mask4d, "sliding_attention": mask4d} |
| Hq = self.encoder(input_ids=batch.q_ids, attention_mask=masks, position_ids=batch.q_pos).last_hidden_state |
| marker = Hq.gather(1, batch.q_markers[..., None].expand(-1, -1, Hq.shape[-1])) |
| return marker, Hq[:, 0] |
|
|
| def decide(self, batch: Batch, memory, memory_mask, root): |
| """Set-level fusion of the options of every question against the state memory.""" |
| marker, cls = self.encode_questions(batch) |
| Q, K, d = marker.shape |
| slots = torch.arange(K, device=marker.device) |
| valid = slots[None, :] < batch.q_count[:, None] |
| ordinal = (batch.q_type != 0).to(marker.dtype)[:, None, None] |
| marker = marker + ordinal * self.level_emb(slots.clamp(max=255))[None] |
| S = torch.cat((cls[:, None] + self.qtype_emb(batch.q_type)[:, None], marker), dim=1) |
| S_mask = torch.cat((torch.ones(Q, 1, dtype=torch.bool, device=marker.device), valid), dim=1) |
| pos = torch.ones(Q, 1 + K, dtype=torch.long, device=marker.device) |
| pos[:, 0] = 0 |
| cos, sin = rope_cos_sin(pos, d // self.cfg.fusion_heads, self.cfg.rope_theta, S.dtype) |
| mem, mem_mask, rt = memory[batch.q_state], memory_mask[batch.q_state], root[batch.q_state] |
| for layer in self.fusion: |
| S = layer(S, cos, sin, S_mask, mem, mem_mask) |
| S = self.final_norm(S) |
| logits = self.scorer(S[:, 1:]).squeeze(-1).float().masked_fill(~S_mask[:, 1:], NEG) |
| probs = (logits / self.temperature[batch.q_type][:, None]).softmax(-1) |
| top = probs.topk(min(2, probs.shape[1]), dim=-1).values |
| if top.shape[1] == 1: |
| top = torch.cat((top, torch.zeros_like(top)), dim=-1) |
| count = batch.q_count.clamp_min(2).float() |
| entropy = -(probs * probs.clamp_min(1e-9).log()).sum(-1) / count.log() |
| feats = torch.stack((top[:, 0], top[:, 0] - top[:, 1], entropy, count / 255.0), dim=-1) |
| conf_in = torch.cat((rt.float(), S[:, 0].float(), feats.detach()), dim=-1) |
| conf_logit = self.conf_head(conf_in).squeeze(-1) |
| return {"logits": logits, "probabilities": probs, "confidence_logit": conf_logit} |
|
|
| def forward(self, batch: Batch): |
| memory, memory_mask, root = self.build_memory(batch, batch.n_states) |
| return self.decide(batch, memory, memory_mask, root) |
|
|
| |
| def save_pretrained(self, directory: str | Path, tokenizer=None) -> None: |
| root = Path(directory) |
| root.mkdir(parents=True, exist_ok=True) |
| self.cfg.save(root) |
| self.encoder.config.save_pretrained(root / "encoder") |
| if tokenizer is not None: |
| tokenizer.save_pretrained(root / "tokenizer") |
| state = {k: v.detach().cpu().contiguous() for k, v in self.state_dict().items()} |
| save_file(state, str(root / "model.safetensors"), metadata={"format": "pt", "family": "groot", "version": "1"}) |
| (root / "calibration.json").write_text(json.dumps({"temperature": self.temperature.tolist()}, indent=2) + "\n") |
|
|
| @classmethod |
| def from_pretrained(cls, directory: str | Path, device: str | torch.device = "cpu") -> "GrootDecisionModel": |
| from transformers import AutoConfig, AutoModel |
|
|
| root = Path(directory) |
| cfg = GrootConfig.load(root) |
| enc_cfg = AutoConfig.from_pretrained(root / "encoder") |
| enc_cfg.reference_compile = False |
| encoder = AutoModel.from_config(enc_cfg, attn_implementation="sdpa") |
| cfg = GrootConfig.from_dict({**cfg.to_dict(), "encoder_layers": None, "fusion_init_from_encoder": False}) |
| model = cls(cfg, encoder) |
| missing, unexpected = model.load_state_dict(load_file(str(root / "model.safetensors")), strict=False) |
| if unexpected or set(missing) - {"chunk_pos_scale"}: |
| raise RuntimeError(f"checkpoint mismatch: missing={missing} unexpected={unexpected}") |
| calib = root / "calibration.json" |
| if calib.exists(): |
| temps = json.loads(calib.read_text()).get("temperature") |
| if temps: |
| model.temperature.copy_(torch.tensor(temps, dtype=model.temperature.dtype)) |
| return model.to(device).eval() |
|
|
| def count_parameters(self) -> dict[str, int]: |
| enc = sum(p.numel() for p in self.encoder.parameters()) |
| total = sum(p.numel() for p in self.parameters()) |
| emb = self.encoder.get_input_embeddings().weight.numel() |
| return {"total": total, "encoder": enc, "embeddings": emb, "head": total - enc} |
|
|
