Publish BananaMind 2.1 Unified transformers wrapper

configuration_bananamind21unified.py

@@ -0,0 +1,86 @@
+from transformers import PretrainedConfig
+
+
+class BananaMind21UnifiedConfig(PretrainedConfig):
+    """BananaMind 2.1 Unified: three towers, a relay middle, one token out.
+
+    A and C are the outer towers and each owns an output head.  B is the relay:
+    no head, no solo loss, and the only path between A and C.
+    """
+
+    model_type = "bananamind21_unified"
+
+    def __init__(
+        self,
+        vocab_size=8192,
+        head_dim=64,
+        max_position_embeddings=4096,
+        rope_theta=100000.0,
+        rms_norm_eps=1e-6,
+        embed_width=384,
+        hidden_size_a=256,
+        num_hidden_layers_a=14,
+        num_attention_heads_a=4,
+        num_key_value_heads_a=1,
+        intermediate_size_a=704,
+        hidden_size_b=320,
+        num_hidden_layers_b=5,
+        num_attention_heads_b=5,
+        num_key_value_heads_b=1,
+        intermediate_size_b=960,
+        hidden_size_c=384,
+        num_hidden_layers_c=6,
+        num_attention_heads_c=6,
+        num_key_value_heads_c=2,
+        intermediate_size_c=1024,
+        a_read=(5, 9, 12),
+        a_land=(7, 11, 14),
+        c_read=(2, 4, 5),
+        c_land=(3, 5, 6),
+        b_land=(1, 3, 5),
+        b_read=(2, 4, 5),
+        gate_init=0.01,
+        cut_bridges=False,
+        tie_word_embeddings=False,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.head_dim = head_dim
+        self.max_position_embeddings = max_position_embeddings
+        self.rope_theta = rope_theta
+        self.rms_norm_eps = rms_norm_eps
+        self.embed_width = embed_width
+
+        self.hidden_size_a = hidden_size_a
+        self.num_hidden_layers_a = num_hidden_layers_a
+        self.num_attention_heads_a = num_attention_heads_a
+        self.num_key_value_heads_a = num_key_value_heads_a
+        self.intermediate_size_a = intermediate_size_a
+
+        self.hidden_size_b = hidden_size_b
+        self.num_hidden_layers_b = num_hidden_layers_b
+        self.num_attention_heads_b = num_attention_heads_b
+        self.num_key_value_heads_b = num_key_value_heads_b
+        self.intermediate_size_b = intermediate_size_b
+
+        self.hidden_size_c = hidden_size_c
+        self.num_hidden_layers_c = num_hidden_layers_c
+        self.num_attention_heads_c = num_attention_heads_c
+        self.num_key_value_heads_c = num_key_value_heads_c
+        self.intermediate_size_c = intermediate_size_c
+
+        self.a_read = list(a_read)
+        self.a_land = list(a_land)
+        self.c_read = list(c_read)
+        self.c_land = list(c_land)
+        self.b_land = list(b_land)
+        self.b_read = list(b_read)
+        self.gate_init = gate_init
+
+        # set true to run the three towers with every bridge removed: A and C
+        # become ordinary standalone transformers and B stops mattering
+        self.cut_bridges = cut_bridges
+
+        # `hidden_size` is what generic HF tooling looks for
+        self.hidden_size = hidden_size_c
+        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)

modeling_bananamind21unified.py

@@ -0,0 +1,290 @@
+"""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.
+
+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.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):
+        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
+        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 forward(self, x, freqs_cis, attention_mask=None):
+        bsz, seq_len, _ = x.size()
+        q = self.q_proj(x).view(bsz, seq_len, self.n_head, self.head_dim).transpose(1, 2)
+        k = self.k_proj(x).view(bsz, seq_len, self.n_kv_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x).view(bsz, seq_len, self.n_kv_heads, self.head_dim).transpose(1, 2)
+        q = self.q_norm(q)
+        k = self.k_norm(k)
+        q, k = apply_rotary_emb(q, k, freqs_cis)
+        k = k.unsqueeze(2).expand(bsz, self.n_kv_heads, self.n_rep, seq_len, self.head_dim)
+        k = k.reshape(bsz, self.n_head, seq_len, self.head_dim).contiguous()
+        v = v.unsqueeze(2).expand(bsz, self.n_kv_heads, self.n_rep, seq_len, self.head_dim)
+        v = v.reshape(bsz, self.n_head, seq_len, self.head_dim).contiguous()
+
+        attn_mask = None
+        is_causal = True
+        if attention_mask is not None:
+            key_pad = attention_mask.to(torch.bool)[:, None, None, :]
+            causal = torch.ones(seq_len, seq_len, dtype=torch.bool, device=x.device).tril()
+            attn_mask = key_pad & causal[None, None, :, :]
+            is_causal = False
+
+        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, seq_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):
+        super().__init__()
+        self.ln_1 = RMSNorm(spec.hidden_size, eps=rms_norm_eps)
+        self.attn = TowerAttention(spec, head_dim, rms_norm_eps)
+        self.ln_2 = RMSNorm(spec.hidden_size, eps=rms_norm_eps)
+        self.mlp = TowerSwiGLUMLP(spec)
+
+    def forward(self, x, freqs_cis, attention_mask=None):
+        x = x + self.attn(self.ln_1(x), freqs_cis, attention_mask=attention_mask)
+        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 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)
+
+        self.blocks_a = nn.ModuleList([TowerBlock(a, config.head_dim, config.rms_norm_eps)
+                                       for _ in range(a.num_hidden_layers)])
+        self.blocks_b = nn.ModuleList([TowerBlock(b, config.head_dim, config.rms_norm_eps)
+                                       for _ in range(b.num_hidden_layers)])
+        self.blocks_c = nn.ModuleList([TowerBlock(c, config.head_dim, config.rms_norm_eps)
+                                       for _ 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
+        self.post_init()
+
+    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 prepare_inputs_for_generation(self, input_ids, past_key_values=None, attention_mask=None, **kwargs):
+        return {"input_ids": input_ids, "attention_mask": attention_mask, "use_cache": False}
+
+    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):
+        while cursor < target:
+            arriving = pending.pop(cursor + 1, None)
+            if arriving is not None:
+                x = x + arriving
+            x = blocks[cursor](x, freqs_cis, attention_mask=attention_mask)
+            cursor += 1
+        return x, cursor
+
+    def hidden_states(self, input_ids, attention_mask=None):
+        cfg = self.config
+        _, seq_len = input_ids.size()
+        freqs_cis = self._get_freqs_cis(seq_len, input_ids.device)
+
+        embedded = self.wte(input_ids) * self._embd_scale
+        x_a = self.in_proj_a(embedded)
+        x_b = self.in_proj_b(embedded)
+        x_c = embedded
+
+        pend_a, pend_b, pend_c = {}, {}, {}
+        cur_a = cur_b = cur_c = 0
+        cut = cfg.cut_bridges
+
+        for k in range(len(cfg.a_read)):
+            x_a, cur_a = self._advance(self.blocks_a, x_a, cur_a, cfg.a_read[k], freqs_cis, pend_a, attention_mask)
+            x_c, cur_c = self._advance(self.blocks_c, x_c, cur_c, cfg.c_read[k], freqs_cis, pend_c, attention_mask)
+            if not cut:
+                pend_b[cfg.b_land[k]] = self.edges_a2b[k](x_a) + self.edges_c2b[k](x_c)
+            x_b, cur_b = self._advance(self.blocks_b, x_b, cur_b, cfg.b_read[k], freqs_cis, pend_b, attention_mask)
+            if not cut:
+                pend_a[cfg.a_land[k]] = self.edges_b2a[k](x_b)
+                pend_c[cfg.c_land[k]] = self.edges_b2c[k](x_b)
+
+        x_a, _ = self._advance(self.blocks_a, x_a, cur_a, cfg.num_hidden_layers_a, freqs_cis, pend_a, attention_mask)
+        x_c, _ = self._advance(self.blocks_c, x_c, cur_c, cfg.num_hidden_layers_c, freqs_cis, pend_c, attention_mask)
+        x_b, _ = self._advance(self.blocks_b, x_b, cur_b, cfg.num_hidden_layers_b, freqs_cis, pend_b, attention_mask)
+
+        h_a = self.ln_f_a(x_a)
+        h_c = self.ln_f_c(x_c)
+        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[object] = None,
+        use_cache=False,
+        **kwargs,
+    ):
+        h_a, h_c, mix_logit = self.hidden_states(input_ids, attention_mask=attention_mask)
+
+        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)
+
+        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=None)