QuixiAI/FlyGPT

1"""FlyGPT: a character-level language model whose recurrent core is a real subgraph of the2fruit-fly connectome (MaleCNS v1.0). Hugging Face `transformers` implementation; self-contained.3 4Dynamics (one scalar state per neuron, plan.md §7 of the FlyGPT spec):5 6    proposal_i = tanh( sum_j W_ij h_j / sqrt(in_degree_i) + external_input_i + bias_i )7    h_i_new    = (1 - leak_i) * h_i + leak_i * proposal_i8 9The connectome is stored in `model.safetensors` as integer tensors (`graph.*`); only the learned10per-edge values and the adapters are floating point (bf16 on disk). The sparse recurrent matmul runs11in fp32: through the fused kernels of the `connectome-kernels` package when it is installed and a12CUDA device is used (training speed), else through torch.sparse COO (rows = destination). Both give13the same logits and gradients.14"""15from __future__ import annotations16 17import math18from dataclasses import dataclass19from typing import Optional20 21import torch22import torch.nn as nn23import torch.nn.functional as F24from transformers import PreTrainedModel25from transformers.generation import GenerationMixin26from transformers.utils import ModelOutput27 28from .configuration_flygpt import FlyGPTConfig29 30 31@dataclass32class FlyGPTOutput(ModelOutput):33    loss: Optional[torch.FloatTensor] = None34    logits: Optional[torch.FloatTensor] = None35    state: Optional[torch.FloatTensor] = None   # [B, N] neuron states after the last character36 37 38class FlyGraph(nn.Module):39    """The anatomy. Integer buffers only; never trained."""40 41    def __init__(self, num_neurons: int, num_edges: int, num_input: int, num_output: int):42        super().__init__()43        self.register_buffer("edge_index", torch.zeros(2, num_edges, dtype=torch.int32))   # [source, destination]44        self.register_buffer("synapse_count", torch.zeros(num_edges, dtype=torch.int32))   # MaleCNS synaptic contacts45        self.register_buffer("node_id", torch.zeros(num_neurons, dtype=torch.int64))       # MaleCNS body ids46        self.register_buffer("input_nodes", torch.zeros(num_input, dtype=torch.int64))47        self.register_buffer("output_nodes", torch.zeros(num_output, dtype=torch.int64))48 49 50class FlyRecurrentCore(nn.Module):51    """The learned state: one value per real edge, plus per-neuron bias and leak."""52 53    def __init__(self, num_neurons: int, num_edges: int, leak_init: float, learned_leak: bool):54        super().__init__()55        self.edge_values = nn.Parameter(torch.zeros(num_edges))56        self.bias = nn.Parameter(torch.zeros(num_neurons))57        self.raw_leak = nn.Parameter(torch.full((num_neurons,), math.log(leak_init / (1 - leak_init))),58                                     requires_grad=learned_leak)59 60 61class FlyGPTPreTrainedModel(PreTrainedModel):62    config_class = FlyGPTConfig63    base_model_prefix = "flygpt"64    _is_stateful = True65    _supports_cache_class = False66    supports_gradient_checkpointing = False67 68    def _init_weights(self, module):69        if isinstance(module, FlyRecurrentCore):70            nn.init.normal_(module.edge_values, std=self.config.init_scale)71            nn.init.zeros_(module.bias)72        elif isinstance(module, nn.Linear):73            nn.init.normal_(module.weight, std=0.02)74            if module.bias is not None:75                nn.init.zeros_(module.bias)76        elif isinstance(module, nn.Embedding):77            nn.init.normal_(module.weight, std=1.0)78 79 80class FlyGPTForCausalLM(FlyGPTPreTrainedModel, GenerationMixin):81    def __init__(self, config: FlyGPTConfig):82        super().__init__(config)83        c = config84        self.graph = FlyGraph(c.num_neurons, c.num_edges, c.num_input_neurons, c.num_output_neurons)85        self.recurrent = FlyRecurrentCore(c.num_neurons, c.num_edges, c.leak_init, c.learned_leak)86        self.embed = nn.Embedding(c.vocab_size, c.embedding_dim)87        self.input_proj = nn.Linear(c.embedding_dim, c.num_input_neurons)88        self.lm_head = nn.Linear(c.num_output_neurons, c.vocab_size)89        self.post_init()90 91    # ---- sparse recurrent matrix -------------------------------------------------------------92    @property93    def num_neurons(self) -> int:94        return self.config.num_neurons95 96    def edge_scale(self) -> torch.Tensor:97        """1/sqrt(in_degree) per edge (degree normalization), or ones."""98        dst = self.graph.edge_index[1].long()99        if not self.config.degree_normalization:100            return torch.ones_like(dst, dtype=torch.float32)101        in_deg = torch.bincount(dst, minlength=self.num_neurons).clamp(min=1).float()102        return 1.0 / in_deg[dst].sqrt()103 104    def sparse_weight(self) -> torch.Tensor:105        src, dst = self.graph.edge_index[0].long(), self.graph.edge_index[1].long()106        values = self.recurrent.edge_values.float() * self.edge_scale()107        return torch.sparse_coo_tensor(torch.stack([dst, src]), values, (self.num_neurons, self.num_neurons))108 109    def dense_weight(self) -> torch.Tensor:110        """Convenience for analysis; [N, N] with rows = destination. Never used in the forward pass."""111        return self.sparse_weight().to_dense()112 113    @property114    def leak(self) -> torch.Tensor:115        return torch.sigmoid(self.recurrent.raw_leak.float())116 117    # ---- dynamics ------------------------------------------------------------------------------118    def init_state(self, batch: int, device=None) -> torch.Tensor:119        return torch.zeros(batch, self.num_neurons, device=device or self.recurrent.edge_values.device)120 121    def drive(self, x: torch.Tensor) -> torch.Tensor:122        d = torch.zeros(x.shape[0], self.num_neurons, device=x.device, dtype=torch.float32)123        d[:, self.graph.input_nodes] = self.input_proj(self.embed(x)).float()124        return d125 126    def step(self, state: torch.Tensor, x: torch.Tensor, W: torch.Tensor | None = None) -> torch.Tensor:127        W = self.sparse_weight() if W is None else W128        drive, leak, bias = self.drive(x), self.leak, self.recurrent.bias.float()129        for _ in range(self.config.microsteps):130            incoming = torch.sparse.mm(W, state.float().T).T131            proposal = torch.tanh(incoming + drive + bias)132            state = (1 - leak) * state + leak * proposal133        return state134 135    def logits_from_state(self, state: torch.Tensor) -> torch.Tensor:136        return self.lm_head(state[:, self.graph.output_nodes].to(self.lm_head.weight.dtype)).float()137 138    # ---- fused CUDA path via the connectome-kernels package (optional, used for training) ------------139    def _fused_graph(self):140        from connectome_kernels import SparseGraph141        dev = self.graph.edge_index.device142        if getattr(self, "_fg", None) is None or self._fg_device != dev:143            self._fg = SparseGraph(self.graph.edge_index[0].long(), self.graph.edge_index[1].long(), self.num_neurons,144                                   self.graph.input_nodes)145            self._fg_device = dev146        return self._fg147 148    def _fused_available(self, device) -> bool:149        if device.type != "cuda":150            return False151        if not hasattr(self, "_fused_ok"):152            try:153                from connectome_kernels import available154                self._fused_ok = available()155            except Exception:156                self._fused_ok = False157        return self._fused_ok158 159    def _forward_fused(self, input_ids, state):160        from connectome_kernels import sparse_recurrence161        drives = self.input_proj(self.embed(input_ids)).float().permute(1, 2, 0).contiguous()       # [T, n_in, B]162        vals = self.recurrent.edge_values.float() * self.edge_scale()163        out = sparse_recurrence(vals, self.leak, self.recurrent.bias.float(), drives, state,164                                self._fused_graph(), self.config.microsteps)                         # [T, B, N]165        logits = self.lm_head(out[:, :, self.graph.output_nodes].to(self.lm_head.weight.dtype)).float().permute(1, 0, 2)166        return logits, out[-1]167 168    def forward(self, input_ids: torch.LongTensor, state: Optional[torch.Tensor] = None,169                labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None,170                return_dict: Optional[bool] = None, **kwargs) -> FlyGPTOutput:171        B, T = input_ids.shape172        state = self.init_state(B, input_ids.device) if state is None else state173        if self._fused_available(input_ids.device):174            logits, state = self._forward_fused(input_ids, state)175        else:176            W = self.sparse_weight()177            outs = []178            for t in range(T):179                state = self.step(state, input_ids[:, t], W)180                outs.append(self.logits_from_state(state))181            logits = torch.stack(outs, 1)182        loss = None183        if labels is not None:184            loss = F.cross_entropy(logits[:, :-1].reshape(-1, logits.shape[-1]), labels[:, 1:].reshape(-1))185        return FlyGPTOutput(loss=loss, logits=logits, state=state)186 187    # ---- generation: carry the neuron state instead of a KV cache ------------------------------188    @classmethod189    def _supports_default_dynamic_cache(cls) -> bool:190        return False  # stateful recurrent model: no KV cache, the neuron state is carried in `state`191 192    def prepare_inputs_for_generation(self, input_ids, state=None, **kwargs):193        if state is not None:194            input_ids = input_ids[:, -1:]195        return {"input_ids": input_ids, "state": state}196 197    def _update_model_kwargs_for_generation(self, outputs, model_kwargs, is_encoder_decoder=False, **kwargs):198        model_kwargs["state"] = outputs.state199        return model_kwargs200