"""BananaMind 2.1 Unified as an HF causal LM. IMPORTANT - what `.logits` contains ----------------------------------- Towers A and C are mixed in *probability* space, so the natural output of this model is a normalised log-probability vector, not a logit vector: log p = logaddexp(log a + log_softmax(logits_A), log(1-a) + log_softmax(logits_C)) `.logits` carries that log-probability vector directly. This is safe for every standard consumer, because `log_softmax` is the identity on an already normalised log-probability vector (its logsumexp is 0), and `softmax(log p) = p`. So loglikelihood scoring, `generate()`, and temperature-1 sampling all behave correctly. What is *not* meaningful is treating these numbers as unnormalised scores with an arbitrary additive offset - they are already calibrated. Tower B is the relay. It has no output head and never appears in the mixture; it exists only to carry signal between A and C, which have no other path to each other. Set `config.cut_bridges = True` to sever every bridge, which turns A and C into two ordinary standalone transformers. KV cache -------- Supported, and on by default. A relay model has no single residual stack, so the three towers share one flat cache index space - tower A first, then B, then C - sized by `config.num_hidden_layers`, which is the sum of the three tower depths. Each tower keeps its own entries; nothing is shared between them. Two properties of this architecture make that work without any extra machinery: * A bridge (`Edge`) is a per-channel gate on a linear map, so it mixes channels but never positions. Every bridge contribution for a newly arriving token is computable from that token's own tower states, so no bridge output has to be cached alongside the KV states. * RoPE enters the attention logits only through the relative offset between query and key, so shifting a whole sequence - which is what left padding does - leaves every attention score unchanged. Absolute positions taken from the cache length are therefore correct for padded batches too. Parameter names match the training module exactly, so a checkpoint transfers without any key rewriting. """ import math from typing import Optional import torch import torch.nn as nn import torch.nn.functional as F from transformers import PreTrainedModel from transformers.cache_utils import Cache, DynamicCache from transformers.generation import GenerationMixin from transformers.modeling_outputs import CausalLMOutputWithPast from .configuration_bananamind21unified import BananaMind21UnifiedConfig class RMSNorm(nn.Module): def __init__(self, dim, eps=1e-6): super().__init__() self.eps = eps self.weight = nn.Parameter(torch.ones(dim)) def forward(self, x): x_float = x.float() rms = torch.rsqrt(x_float.pow(2).mean(-1, keepdim=True) + self.eps) return (x_float * rms * self.weight.float()).type_as(x) def precompute_freqs_cis(head_dim, seq_len, theta=100000.0): inv = 1.0 / (theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim)) t = torch.arange(seq_len, dtype=torch.float32) return torch.polar(torch.ones_like(torch.outer(t, inv)), torch.outer(t, inv)) def apply_rotary_emb(q, k, freqs_cis): q_complex = torch.view_as_complex(q.float().reshape(*q.shape[:-1], -1, 2)) k_complex = torch.view_as_complex(k.float().reshape(*k.shape[:-1], -1, 2)) freqs_cis = freqs_cis.unsqueeze(0).unsqueeze(0) q_out = torch.view_as_real(q_complex * freqs_cis).flatten(-2) k_out = torch.view_as_real(k_complex * freqs_cis).flatten(-2) return q_out.type_as(q), k_out.type_as(k) class _Spec: def __init__(self, hidden_size, num_hidden_layers, num_attention_heads, num_key_value_heads, intermediate_size): self.hidden_size = hidden_size self.num_hidden_layers = num_hidden_layers self.num_attention_heads = num_attention_heads self.num_key_value_heads = num_key_value_heads self.intermediate_size = intermediate_size class TowerAttention(nn.Module): def __init__(self, spec, head_dim, rms_norm_eps, layer_idx=None): super().__init__() self.n_head = spec.num_attention_heads self.n_kv_heads = spec.num_key_value_heads self.head_dim = head_dim self.n_rep = self.n_head // self.n_kv_heads # flat index into the shared KV cache; see `num_hidden_layers` in the config self.layer_idx = layer_idx self.q_proj = nn.Linear(spec.hidden_size, self.n_head * head_dim, bias=False) self.k_proj = nn.Linear(spec.hidden_size, self.n_kv_heads * head_dim, bias=False) self.v_proj = nn.Linear(spec.hidden_size, self.n_kv_heads * head_dim, bias=False) self.o_proj = nn.Linear(self.n_head * head_dim, spec.hidden_size, bias=False) self.q_norm = RMSNorm(head_dim, eps=rms_norm_eps) self.k_norm = RMSNorm(head_dim, eps=rms_norm_eps) def _attn_mask(self, attention_mask, q_len, kv_len, past_len, device): """Bottom-right aligned causal mask, intersected with the padding mask. `F.scaled_dot_product_attention(is_causal=True)` aligns its mask to the *top left*, which is only correct when `q_len == kv_len`. With a cache the queries sit at positions `past_len .. past_len+q_len-1` while the keys start at 0, so the mask has to be built explicitly. Returns `(attn_mask, is_causal)` and prefers the cheap fast paths. """ if attention_mask is None: if kv_len == q_len: return None, True # prefill: top-left alignment is correct if q_len == 1: return None, False # one new token attends to the whole cache q_pos = torch.arange(q_len, device=device) + past_len k_pos = torch.arange(kv_len, device=device) mask = k_pos[None, :] <= q_pos[:, None] mask = mask[None, None, :, :] if attention_mask is not None: mask = mask & attention_mask.to(torch.bool)[:, None, None, :kv_len] # A fully-padded query row would be all-False, and SDPA turns an # all-masked row into NaN, which then leaks through the residual # stream into every later position. A token may always attend to # itself, which is causally legal and keeps every row non-empty. mask = mask | (k_pos[None, :] == q_pos[:, None])[None, None, :, :] return mask, False def forward(self, x, freqs_cis, attention_mask=None, past_key_values=None, past_len=0): bsz, q_len, _ = x.size() q = self.q_proj(x).view(bsz, q_len, self.n_head, self.head_dim).transpose(1, 2) k = self.k_proj(x).view(bsz, q_len, self.n_kv_heads, self.head_dim).transpose(1, 2) v = self.v_proj(x).view(bsz, q_len, self.n_kv_heads, self.head_dim).transpose(1, 2) q = self.q_norm(q) k = self.k_norm(k) # `freqs_cis` is already sliced to this step's absolute positions q, k = apply_rotary_emb(q, k, freqs_cis) # Cache the pre-GQA-expansion states: n_kv_heads is 1 or 2 here, so # storing the expanded copies would cost up to 4x the memory for nothing. if past_key_values is not None: k, v = past_key_values.update(k, v, self.layer_idx) kv_len = k.size(-2) k = k.unsqueeze(2).expand(bsz, self.n_kv_heads, self.n_rep, kv_len, self.head_dim) k = k.reshape(bsz, self.n_head, kv_len, self.head_dim).contiguous() v = v.unsqueeze(2).expand(bsz, self.n_kv_heads, self.n_rep, kv_len, self.head_dim) v = v.reshape(bsz, self.n_head, kv_len, self.head_dim).contiguous() attn_mask, is_causal = self._attn_mask( attention_mask, q_len, kv_len, past_len, x.device ) y = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, is_causal=is_causal) y = y.transpose(1, 2).contiguous().view(bsz, q_len, self.n_head * self.head_dim) return self.o_proj(y) class TowerSwiGLUMLP(nn.Module): def __init__(self, spec): super().__init__() self.w_gate = nn.Linear(spec.hidden_size, spec.intermediate_size, bias=False) self.w_up = nn.Linear(spec.hidden_size, spec.intermediate_size, bias=False) self.w_down = nn.Linear(spec.intermediate_size, spec.hidden_size, bias=False) def forward(self, x): return self.w_down(F.silu(self.w_gate(x)) * self.w_up(x)) class TowerBlock(nn.Module): def __init__(self, spec, head_dim, rms_norm_eps, layer_idx=None): super().__init__() self.ln_1 = RMSNorm(spec.hidden_size, eps=rms_norm_eps) self.attn = TowerAttention(spec, head_dim, rms_norm_eps, layer_idx=layer_idx) self.ln_2 = RMSNorm(spec.hidden_size, eps=rms_norm_eps) self.mlp = TowerSwiGLUMLP(spec) def forward(self, x, freqs_cis, attention_mask=None, past_key_values=None, past_len=0): x = x + self.attn( self.ln_1(x), freqs_cis, attention_mask=attention_mask, past_key_values=past_key_values, past_len=past_len, ) x = x + self.mlp(self.ln_2(x)) return x class Edge(nn.Module): def __init__(self, dim_in, dim_out, gate_init=0.01): super().__init__() self.w = nn.Linear(dim_in, dim_out, bias=False) self.g = nn.Parameter(torch.full((dim_out,), float(gate_init))) def forward(self, x): return self.g * self.w(x) class RelayMode: """Which towers run, whether the bridges carry, and which heads vote. The point of the 2.1 topology is that A and C have no direct path to each other, so "what is the relay worth?" is only answerable by ablation. Each mode below cuts the model somewhere different; every one of them runs on the same weights, with no retraining and no reloading. ``run_b=False`` with ``bridges=True`` is the interesting one: tower B's blocks are replaced by the identity, but the bridge wiring stays live, so A and C still exchange signal - through a channel that does no computation. That separates "the relay computes something" from "a channel exists at all". """ __slots__ = ("name", "run_a", "run_b", "run_c", "bridges", "head_a", "head_c", "doc") def __init__(self, name, run_a, run_b, run_c, bridges, head_a, head_c, doc): self.name = name self.run_a, self.run_b, self.run_c = run_a, run_b, run_c self.bridges = bridges self.head_a, self.head_c = head_a, head_c self.doc = doc @property def two_headed(self): return self.head_a and self.head_c RELAY_MODES = { m.name: m for m in ( RelayMode("full", True, True, True, True, True, True, "the model as trained: three towers, all bridges, both heads mixed"), RelayMode("cut_bridges", True, True, True, False, True, True, "all 12 bridges severed; A and C become standalone transformers, " "B is orphaned; the two heads are still mixed"), RelayMode("bypass_b", True, False, True, True, True, True, "B's 5 blocks are skipped but every bridge stays live: A and C " "still exchange, through a relay that does no computation"), RelayMode("ab_only", True, True, False, True, True, False, "tower C is switched off; A and B run with the A<->B bridges " "live; head A alone produces the distribution"), RelayMode("cb_only", False, True, True, True, False, True, "tower A is switched off; C and B run with the C<->B bridges " "live; head C alone produces the distribution"), RelayMode("a_only", True, False, False, False, True, False, "tower A alone, no bridges, head A alone"), RelayMode("c_only", False, False, True, False, False, True, "tower C alone, no bridges, head C alone"), ) } _SINGLE_TOWER_MODES = {"a": "a_only", "c": "c_only"} def resolve_relay_mode(relay_mode=None, use_single_tower=None, cut_bridges=False): """Fold the three user-facing switches into exactly one `RelayMode`. They overlap on purpose - `cut_bridges` predates the others and stays supported - so anything contradictory is rejected rather than silently resolved in an order nobody can guess. """ chosen = [] if relay_mode is not None: name = str(relay_mode).lower() if name not in RELAY_MODES: raise ValueError( f"unknown relay_mode {relay_mode!r}; expected one of " f"{sorted(RELAY_MODES)}" ) chosen.append(("relay_mode", name)) if use_single_tower is not None: tower = str(use_single_tower).lower() if tower not in _SINGLE_TOWER_MODES: raise ValueError( f"use_single_tower must be 'a' or 'c' (tower B has no output " f"head and cannot run alone), got {use_single_tower!r}" ) chosen.append(("use_single_tower", _SINGLE_TOWER_MODES[tower])) if cut_bridges: chosen.append(("cut_bridges", "cut_bridges")) names = {name for _, name in chosen} if len(names) > 1: detail = ", ".join(f"{src}={name!r}" for src, name in chosen) raise ValueError(f"conflicting relay settings: {detail}") return RELAY_MODES[names.pop()] if names else RELAY_MODES["full"] class BananaMind21UnifiedPreTrainedModel(PreTrainedModel): config_class = BananaMind21UnifiedConfig base_model_prefix = "model" supports_gradient_checkpointing = False def _init_weights(self, module): if isinstance(module, nn.Linear): torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) if module.bias is not None: torch.nn.init.zeros_(module.bias) elif isinstance(module, nn.Embedding): torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) class BananaMind21UnifiedForCausalLM(BananaMind21UnifiedPreTrainedModel, GenerationMixin): def __init__(self, config): super().__init__(config) self.config = config a = _Spec(config.hidden_size_a, config.num_hidden_layers_a, config.num_attention_heads_a, config.num_key_value_heads_a, config.intermediate_size_a) b = _Spec(config.hidden_size_b, config.num_hidden_layers_b, config.num_attention_heads_b, config.num_key_value_heads_b, config.intermediate_size_b) c = _Spec(config.hidden_size_c, config.num_hidden_layers_c, config.num_attention_heads_c, config.num_key_value_heads_c, config.intermediate_size_c) self.wte = nn.Embedding(config.vocab_size, config.embed_width) self.in_proj_a = nn.Linear(config.embed_width, a.hidden_size, bias=False) self.in_proj_b = nn.Linear(config.embed_width, b.hidden_size, bias=False) # Flat KV-cache index space over the three towers, in the order the # config documents: A first, then B, then C. Each tower keeps its own # cache entries; nothing is shared between them. self._cache_offset_a = 0 self._cache_offset_b = a.num_hidden_layers self._cache_offset_c = a.num_hidden_layers + b.num_hidden_layers self.blocks_a = nn.ModuleList([TowerBlock(a, config.head_dim, config.rms_norm_eps, layer_idx=self._cache_offset_a + i) for i in range(a.num_hidden_layers)]) self.blocks_b = nn.ModuleList([TowerBlock(b, config.head_dim, config.rms_norm_eps, layer_idx=self._cache_offset_b + i) for i in range(b.num_hidden_layers)]) self.blocks_c = nn.ModuleList([TowerBlock(c, config.head_dim, config.rms_norm_eps, layer_idx=self._cache_offset_c + i) for i in range(c.num_hidden_layers)]) n = len(config.a_read) gi = config.gate_init self.edges_a2b = nn.ModuleList([Edge(a.hidden_size, b.hidden_size, gi) for _ in range(n)]) self.edges_c2b = nn.ModuleList([Edge(c.hidden_size, b.hidden_size, gi) for _ in range(n)]) self.edges_b2a = nn.ModuleList([Edge(b.hidden_size, a.hidden_size, gi) for _ in range(n)]) self.edges_b2c = nn.ModuleList([Edge(b.hidden_size, c.hidden_size, gi) for _ in range(n)]) self.ln_f_a = RMSNorm(a.hidden_size, eps=config.rms_norm_eps) self.ln_f_c = RMSNorm(c.hidden_size, eps=config.rms_norm_eps) self.lm_head_a = nn.Linear(a.hidden_size, config.vocab_size, bias=False) self.lm_head_c = nn.Linear(c.hidden_size, config.vocab_size, bias=False) self.mix_head = nn.Linear(a.hidden_size + c.hidden_size, 1, bias=True) self._embd_scale = math.sqrt(config.embed_width) self._freqs_cis_cache = None # Validated once at construction so a bad mode fails on load, not on the # first forward pass halfway through an evaluation sweep. self._config_mode = resolve_relay_mode( getattr(config, "relay_mode", None), getattr(config, "use_single_tower", None), getattr(config, "cut_bridges", False), ) self.post_init() @property def relay_mode(self): """The `RelayMode` this model runs in unless a call overrides it.""" return self._config_mode def set_relay_mode(self, relay_mode=None, use_single_tower=None): """Switch ablation mode in place, for sweeping without reloading. Returns the resolved `RelayMode`. Any KV cache built under the previous mode is invalid afterwards - the towers it holds state for may no longer be the towers that run - so start a fresh cache after calling this. """ mode = resolve_relay_mode(relay_mode, use_single_tower) self._config_mode = mode self.config.relay_mode = mode.name self.config.use_single_tower = None self.config.cut_bridges = mode.name == "cut_bridges" return mode def get_input_embeddings(self): return self.wte def set_input_embeddings(self, value): self.wte = value def get_output_embeddings(self): return self.lm_head_c def resolve_mode(self, relay_mode=None, use_single_tower=None): """The `RelayMode` in force, per-call arguments overriding the config.""" if relay_mode is None and use_single_tower is None: return self._config_mode return resolve_relay_mode(relay_mode, use_single_tower) def prepare_inputs_for_generation( self, input_ids, past_key_values=None, attention_mask=None, use_cache=True, relay_mode=None, use_single_tower=None, **kwargs ): # With a warm cache only the tokens the towers have not seen yet need to # be forwarded. Every live tower advances over the same positions, so # one probed cache length describes all of them. mode = self.resolve_mode(relay_mode, use_single_tower) past_len = self._cache_len(past_key_values, mode) if past_len > 0: input_ids = input_ids[:, past_len:] out = { "input_ids": input_ids, "attention_mask": attention_mask, "past_key_values": past_key_values, "use_cache": use_cache, } # Only forward an explicit override; otherwise the config mode applies. if relay_mode is not None: out["relay_mode"] = relay_mode if use_single_tower is not None: out["use_single_tower"] = use_single_tower return out def _get_freqs_cis(self, seq_len, device): cache = self._freqs_cis_cache if cache is None or cache.device != device or cache.size(0) < seq_len: cache = precompute_freqs_cis( self.config.head_dim, seq_len, self.config.rope_theta ).to(device) self._freqs_cis_cache = cache return cache[:seq_len] @staticmethod def _advance(blocks, x, cursor, target, freqs_cis, pending, attention_mask, past_key_values=None, past_len=0, run_blocks=True): """Run `blocks[cursor:target]`, folding in bridge arrivals as they land. With `run_blocks=False` the blocks are replaced by the identity but the arrivals are still added and the cursor still advances, which is what turns tower B into a pass-through channel in `bypass_b`. """ while cursor < target: arriving = pending.pop(cursor + 1, None) if arriving is not None: x = x + arriving if run_blocks: x = blocks[cursor]( x, freqs_cis, attention_mask=attention_mask, past_key_values=past_key_values, past_len=past_len, ) cursor += 1 return x, cursor def _cache_probe_idx(self, mode): """Flat cache index of the first block that actually runs in `mode`. `Cache.get_seq_length()` defaults to layer 0, which is tower A's first block - but A does not run in every mode, and an untouched layer reports length 0 forever. Probing a layer that really runs keeps the cache length honest under every ablation. """ if mode.run_a: return self._cache_offset_a if mode.run_b: return self._cache_offset_b return self._cache_offset_c def _cache_len(self, past_key_values, mode): if past_key_values is None: return 0 return past_key_values.get_seq_length(self._cache_probe_idx(mode)) def hidden_states(self, input_ids, attention_mask=None, past_key_values=None, relay_mode=None, use_single_tower=None): """Run the towers and return `(h_a, h_c, mix_logit)`. A tower switched off by the active mode returns `None` in its slot, and `mix_logit` is `None` whenever only one head is live - there is nothing for the mixer to weigh. In the default `full` mode all three are always tensors, so existing callers are unaffected. """ cfg = self.config mode = self.resolve_mode(relay_mode, use_single_tower) _, seq_len = input_ids.size() # Read the cache length once, before any layer writes to it. past_len = self._cache_len(past_key_values, mode) # RoPE is applied at absolute positions, so a cached run has to skip the # `past_len` entries the earlier tokens already used. freqs_cis = self._get_freqs_cis(past_len + seq_len, input_ids.device)[past_len:] embedded = self.wte(input_ids) * self._embd_scale x_a = self.in_proj_a(embedded) if mode.run_a else None x_b = self.in_proj_b(embedded) if (mode.run_b or mode.bridges) else None x_c = embedded if mode.run_c else None pend_a, pend_b, pend_c = {}, {}, {} cur_a = cur_b = cur_c = 0 kv = dict(past_key_values=past_key_values, past_len=past_len) step_a = dict(kv, run_blocks=mode.run_a) step_b = dict(kv, run_blocks=mode.run_b) step_c = dict(kv, run_blocks=mode.run_c) # A bridge is only carried when both of its endpoints are live: in # `ab_only` there is no tower C to read from or write back to, so the # C-side edges are skipped rather than fed zeros. carry_a = mode.bridges and mode.run_a carry_c = mode.bridges and mode.run_c relay_live = mode.bridges and x_b is not None # An `Edge` is a per-channel gate on a linear map, so a bridge mixes # channels but never positions. Every bridge contribution for the new # tokens is therefore computable from the new tokens' own tower states, # and no bridge output needs to be cached alongside the KV states. for k in range(len(cfg.a_read)): if x_a is not None: x_a, cur_a = self._advance(self.blocks_a, x_a, cur_a, cfg.a_read[k], freqs_cis, pend_a, attention_mask, **step_a) if x_c is not None: x_c, cur_c = self._advance(self.blocks_c, x_c, cur_c, cfg.c_read[k], freqs_cis, pend_c, attention_mask, **step_c) if relay_live: into_b = None if carry_a: into_b = self.edges_a2b[k](x_a) if carry_c: contrib = self.edges_c2b[k](x_c) into_b = contrib if into_b is None else into_b + contrib if into_b is not None: pend_b[cfg.b_land[k]] = into_b if x_b is not None: x_b, cur_b = self._advance(self.blocks_b, x_b, cur_b, cfg.b_read[k], freqs_cis, pend_b, attention_mask, **step_b) if relay_live: if carry_a: pend_a[cfg.a_land[k]] = self.edges_b2a[k](x_b) if carry_c: pend_c[cfg.c_land[k]] = self.edges_b2c[k](x_b) if x_a is not None: x_a, _ = self._advance(self.blocks_a, x_a, cur_a, cfg.num_hidden_layers_a, freqs_cis, pend_a, attention_mask, **step_a) if x_c is not None: x_c, _ = self._advance(self.blocks_c, x_c, cur_c, cfg.num_hidden_layers_c, freqs_cis, pend_c, attention_mask, **step_c) if x_b is not None: x_b, _ = self._advance(self.blocks_b, x_b, cur_b, cfg.num_hidden_layers_b, freqs_cis, pend_b, attention_mask, **step_b) h_a = self.ln_f_a(x_a) if x_a is not None else None h_c = self.ln_f_c(x_c) if x_c is not None else None mix_logit = None if mode.two_headed: mix_logit = self.mix_head(torch.cat([h_a, h_c], dim=-1)).squeeze(-1) return h_a, h_c, mix_logit def forward( self, input_ids, attention_mask=None, labels=None, past_key_values: Optional[Cache] = None, use_cache: Optional[bool] = None, relay_mode=None, use_single_tower=None, **kwargs, ): mode = self.resolve_mode(relay_mode, use_single_tower) if use_cache is None: use_cache = getattr(self.config, "use_cache", True) if labels is not None: # A scoring or training pass consumes the whole sequence in one go # and throws the states away, so building 25 layers of cache is pure # overhead. An explicitly supplied cache is still honoured. use_cache = use_cache and past_key_values is not None if use_cache and past_key_values is None: past_key_values = DynamicCache(config=self.config) if not use_cache: past_key_values = None h_a, h_c, mix_logit = self.hidden_states( input_ids, attention_mask=attention_mask, past_key_values=past_key_values, relay_mode=mode.name, ) # With one head live there is nothing to mix, and `log_softmax` alone # already gives the normalised log-probability vector `.logits` promises. if mode.two_headed: log_p_a = F.log_softmax(self.lm_head_a(h_a).float(), dim=-1) log_p_c = F.log_softmax(self.lm_head_c(h_c).float(), dim=-1) log_alpha = F.logsigmoid(mix_logit).unsqueeze(-1) log_one_minus = F.logsigmoid(-mix_logit).unsqueeze(-1) log_p = torch.logaddexp(log_alpha + log_p_a, log_one_minus + log_p_c) elif mode.head_a: log_p = F.log_softmax(self.lm_head_a(h_a).float(), dim=-1) else: log_p = F.log_softmax(self.lm_head_c(h_c).float(), dim=-1) loss = None if labels is not None: loss = F.nll_loss( log_p[..., :-1, :].reshape(-1, log_p.size(-1)), labels[..., 1:].reshape(-1), ) return CausalLMOutputWithPast(loss=loss, logits=log_p, past_key_values=past_key_values)