Upload folder using huggingface_hub

config.json

@@ -0,0 +1,20 @@
+{
+  "auto_map": {
+    "AutoConfig": "modeling.StreamMixerConfig",
+    "AutoModel": "modeling.StreamMixerModel",
+    "AutoModelForCausalLM": "modeling.StreamMixerForCausalLM"
+  },
+  "hidden_size": 384,
+  "intermediate_size": 1024,
+  "max_sequence_length": 2048,
+  "model_type": "streammixer",
+  "n_embd": 384,
+  "n_layer": 14,
+  "n_read_heads": 4,
+  "n_streams": 32,
+  "num_attention_heads": 4,
+  "num_hidden_layers": 14,
+  "stream_dim": 64,
+  "transformers_version": "5.10.2",
+  "vocab_size": 16384
+}

model.py

@@ -0,0 +1,296 @@
+"""Stream Mixer GPT — model definition only. No side effects on import.
+
+A linear-time, attention-free language model. Each layer mixes the sequence via
+M parallel content-routed memory streams updated by a stable chunked parallel
+scan. Multi-head sigmoid-gated reads with QK-norm. Strictly O(B·T·M·D) per layer.
+
+Imported by both microgpt.py (training) and infer.py (sampling).
+"""
+import os
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# ---------------------------------------------------------------------------
+# BF16 mixed precision — explicit, no torch.amp.autocast.
+# Master weights stay fp32 for optimizer precision; matmuls run in
+# COMPUTE_DTYPE (bf16 on SM80+, fp32 fallback elsewhere).
+# Override with: NANOCHAT_DTYPE=float32|bfloat16|float16
+# ---------------------------------------------------------------------------
+_DTYPE_MAP = {"bfloat16": torch.bfloat16, "float16": torch.float16, "float32": torch.float32}
+def _detect_compute_dtype():
+    env = os.environ.get("NANOCHAT_DTYPE")
+    if env is not None:
+        return _DTYPE_MAP[env]
+    if torch.cuda.is_available():
+        capability = torch.cuda.get_device_capability()
+        if capability >= (8, 0):
+            return torch.bfloat16
+        return torch.float32
+    return torch.float32
+
+COMPUTE_DTYPE = _detect_compute_dtype()
+
+
+class Linear(nn.Linear):
+    """nn.Linear that casts weights to match input dtype in forward.
+    Replaces autocast: master weights stay fp32 for optimizer precision,
+    but matmuls run in the activation dtype (typically bf16 from embeddings)."""
+    def forward(self, x):
+        b = None if self.bias is None else self.bias.to(dtype=x.dtype)
+        return F.linear(x, self.weight.to(dtype=x.dtype), b)
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+    def forward(self, x):
+        return F.rms_norm(x, (x.size(-1),))
+
+
+class StreamMixer(nn.Module):
+    """Recurrent sequence mixer over M parallel content-routed memory streams,
+    with H multi-head sigmoid-gated reads.
+
+    Per token t:
+      - write_value v[t]   = W_v(x[t])              shape (D,)
+      - read_query q[t]    = W_q(x[t])              shape (H, D)   — H read heads
+      - write_route r[t]   = softmax(W_r(x[t]))     shape (M,)
+      - per-stream log-decay log_α[t] = −LOG_A_MAX_NEG·σ(W_α(x[t]) + wa.bias),
+        where wa.bias is a per-stream learnable vector initialized so that streams
+        span a wide range of timescales at init (see microgpt.py init).
+    Recurrence (per stream i):
+      s[t, i] = α[t, i] · s[t-1, i] + r[t, i] · v[t]
+    Read (per head h):
+      score[t, h, i] = (RMSNorm(q[t, h]) · RMSNorm(s[t, i])) / √D     — QK-norm
+      out[t, h]      = Σ_i σ(score[t, h, i]) · s[t, i]
+    Output:
+      Σ_out(out[t, :].flatten()) → (n_embd,)
+
+    Solved via chunked scan: within each chunk of length C,
+      Z_local[t] = Σ log_α[j] (j from chunk_start to t)
+      s[t] = exp(Z_local[t]) · (state_in + cumsum_in_chunk(exp(−Z_local) · write))
+    Cross-chunk: state_in propagates serially through n_chunks=T/C iterations
+    (each iteration is O(B·M·D); n_chunks≈8 for T=1024, C=128).
+
+    With LOG_A_MAX_NEG=0.5 and C=128, Z within a chunk ∈ (−64, 0) — exp(±64)
+    fits in float32 range, allowing α to span ~(0.61, 1.0) per step (half-lives
+    from ~1.4 steps to ∞). The old single-cumsum scan could only handle
+    LOG_A_MAX_NEG < ~0.086 for T=1024.
+
+    Causality is built into the recurrence; position is implicit (recency + the
+    diverse per-stream timescales encode coarse position information).
+
+    The output linear is named `wo_out` (not `wo`) so the training script's
+    residual-init pattern (which matches `wo.weight` / `w_out.weight` suffix)
+    misses it — the mixer needs a louder initial residual contribution than the
+    standard 1/√(2L) scaling provides.
+    """
+    LOG_A_MAX_NEG = 0.5    # caps |log α| per step; α ∈ (exp(-0.5), 1.0) ≈ (0.61, 1.0)
+    CHUNK_SIZE    = 128    # T must be a multiple of this; controls scan numerical range
+
+    def __init__(self, n_embd, n_streams, stream_dim, n_read_heads):
+        super().__init__()
+        self.M = n_streams
+        self.D = stream_dim
+        self.H = n_read_heads
+        self.wv = Linear(n_embd, stream_dim, bias=False)
+        self.wq = Linear(n_embd, n_read_heads * stream_dim, bias=False)
+        self.wr = Linear(n_embd, n_streams, bias=False)
+        # wa.bias is per-stream; initialized in microgpt.py to spread timescales.
+        self.wa = Linear(n_embd, n_streams, bias=True)
+        self.wo_out = Linear(n_read_heads * stream_dim, n_embd, bias=False)
+
+    @staticmethod
+    def _rms_norm_last(x):
+        return F.rms_norm(x, (x.size(-1),))
+
+    def _read(self, q_flat, s):
+        """Multi-head sigmoid-gated read with QK-norm.
+        q_flat: (..., H*D); s: (..., M, D); returns (..., H*D)."""
+        lead = q_flat.shape[:-1]
+        q = q_flat.view(*lead, self.H, self.D)                      # (..., H, D)
+        q_n = self._rms_norm_last(q)                                 # (..., H, D)
+        s_n = self._rms_norm_last(s)                                 # (..., M, D)
+        # score[..., h, i] = (q_n[..., h] · s_n[..., i]) / √D
+        scores = torch.einsum('...hd,...md->...hm', q_n, s_n) * (self.D ** -0.5)
+        weights = torch.sigmoid(scores)                              # (..., H, M)
+        read = torch.einsum('...hm,...md->...hd', weights, s)        # (..., H, D)
+        return read.reshape(*lead, self.H * self.D)
+
+    def _scan(self, log_a, bv):
+        """Chunked stable scan implementing s[t] = α[t]·s[t-1] + bv[t].
+        log_a: (B, T, M); bv: (B, T, M, D). Returns s: (B, T, M, D).
+
+        Within each chunk we use the closed-form cumsum identity. Across chunks
+        the incoming state propagates serially through `n_chunks` iterations —
+        cheap (B·M·D per step, ~8 steps total at T=1024). The chunking caps Z's
+        dynamic range to ±LOG_A_MAX_NEG·CHUNK_SIZE within any single exp() call,
+        keeping everything in float range even for fast-decay streams.
+
+        T need not be a multiple of CHUNK_SIZE: we right-pad the time dim with
+        log_a=0 (α=1, no decay) and bv=0 (no write), run the scan, and slice
+        back. Matters for prompt prefill where T can be smaller than C.
+        """
+        B, T, M, D = bv.shape
+        C = self.CHUNK_SIZE
+        pad = (C - T % C) % C
+        if pad > 0:
+            log_a = F.pad(log_a, (0, 0, 0, pad))               # pad T dim only
+            bv    = F.pad(bv,    (0, 0, 0, 0, 0, pad))         # pad T dim only
+        Tp = T + pad
+        n_chunks = Tp // C
+
+        log_a_c = log_a.view(B, n_chunks, C, M)
+        bv_c    = bv.view(B, n_chunks, C, M, D)
+
+        Z = log_a_c.cumsum(dim=2)                              # (B, n_c, C, M)
+        eZ = torch.exp(Z).unsqueeze(-1)                        # (B, n_c, C, M, 1)
+        inv_eZ = torch.exp(-Z).unsqueeze(-1)                   # (B, n_c, C, M, 1)
+
+        # Intra-chunk contribution assuming zero incoming state
+        s_intra = (bv_c * inv_eZ).cumsum(dim=2) * eZ           # (B, n_c, C, M, D)
+
+        # Cross-chunk: serial propagation. chunk_decay[c] is the product of α's
+        # across the entire chunk (i.e. exp(Z_local[C-1])).
+        chunk_decay     = eZ[:, :, -1]                         # (B, n_c, M, 1)
+        chunk_end_intra = s_intra[:, :, -1]                    # (B, n_c, M, D)
+
+        incomings = []
+        state = torch.zeros(B, M, D, device=bv.device, dtype=bv.dtype)
+        for c in range(n_chunks):
+            incomings.append(state)
+            state = chunk_decay[:, c] * state + chunk_end_intra[:, c]
+        incoming = torch.stack(incomings, dim=1)               # (B, n_c, M, D)
+
+        # Add carryover: position t inside chunk c receives eZ[c, t] · incoming[c]
+        s = s_intra + eZ * incoming.unsqueeze(2)               # (B, n_c, C, M, D)
+        return s.view(B, Tp, M, D)[:, :T]                      # slice off padding
+
+    def forward(self, x):
+        out, _ = self.forward_with_state(x)
+        return out
+
+    def forward_with_state(self, x):
+        """Full parallel scan. Returns (output (B,T,C_embd), final_state (B,M,D))."""
+        v = self.wv(x)                                                   # (B, T, D)
+        q = self.wq(x)                                                   # (B, T, H*D)
+        r = F.softmax(self.wr(x), dim=-1)                                # (B, T, M)
+        log_a = -self.LOG_A_MAX_NEG * torch.sigmoid(self.wa(x))          # (B, T, M)
+        bv = r.unsqueeze(-1) * v.unsqueeze(2)                            # (B, T, M, D)
+        s = self._scan(log_a, bv)                                        # (B, T, M, D)
+        read = self._read(q, s)                                          # (B, T, H*D)
+        return self.wo_out(read), s[:, -1]                               # state @ last position
+
+    def step(self, x, state):
+        """One token advance with explicit state. x:(B,1,C_embd), state:(B,M,D).
+        Returns (output (B,1,C_embd), new_state (B,M,D)). Cost O(M·D) per token."""
+        x_t = x.squeeze(1)                                               # (B, C_embd)
+        v = self.wv(x_t)                                                 # (B, D)
+        q = self.wq(x_t)                                                 # (B, H*D)
+        r = F.softmax(self.wr(x_t), dim=-1)                              # (B, M)
+        log_a = -self.LOG_A_MAX_NEG * torch.sigmoid(self.wa(x_t))        # (B, M)
+        a = torch.exp(log_a)                                             # (B, M)
+        write = r.unsqueeze(-1) * v.unsqueeze(1)                         # (B, M, D)
+        s = a.unsqueeze(-1) * state + write                              # (B, M, D)
+        read = self._read(q, s)                                          # (B, H*D)
+        return self.wo_out(read).unsqueeze(1), s
+
+
+class MLP(nn.Module):
+    """ReLU² MLP: ReLU(W_up x)² · W_down, then W_out. Hidden ≈ 8/3 * n_embd
+    keeps param count similar to a 4x GELU MLP."""
+    def __init__(self, n_embd):
+        super().__init__()
+        hidden = (int(8 * n_embd / 3) + 15) // 16 * 16
+        self.w_up   = Linear(n_embd, hidden, bias=False)
+        self.w_down = Linear(n_embd, hidden, bias=False)
+        self.w_out  = Linear(hidden, n_embd, bias=False)
+    def forward(self, x):
+        return self.w_out(F.relu(self.w_up(x)).square() * self.w_down(x))
+
+
+class Block(nn.Module):
+    def __init__(self, n_embd, n_streams, stream_dim, n_read_heads):
+        super().__init__()
+        self.ln1 = RMSNorm(n_embd)
+        self.mix = StreamMixer(n_embd, n_streams, stream_dim, n_read_heads)
+        self.ln2 = RMSNorm(n_embd)
+        self.mlp = MLP(n_embd)
+    def forward(self, x):
+        x, _ = self.forward_with_state(x)
+        return x
+    def forward_with_state(self, x):
+        mix_out, state = self.mix.forward_with_state(self.ln1(x))
+        x = x + mix_out
+        x = x + self.mlp(self.ln2(x))
+        return x, state
+    def step(self, x, state):
+        mix_out, new_state = self.mix.step(self.ln1(x), state)
+        x = x + mix_out
+        x = x + self.mlp(self.ln2(x))
+        return x, new_state
+
+
+class GPT(nn.Module):
+    def __init__(self, vocab_size, n_embd, n_layer, n_streams, stream_dim, n_read_heads):
+        super().__init__()
+        # Pad vocab to multiple of 64 for tensor core efficiency
+        padded_vocab_size = ((vocab_size + 63) // 64) * 64
+        self.config = dict(
+            vocab_size=padded_vocab_size, n_embd=n_embd, n_layer=n_layer,
+            n_streams=n_streams, stream_dim=stream_dim, n_read_heads=n_read_heads,
+        )
+        self.wte = nn.Embedding(padded_vocab_size, n_embd, dtype=COMPUTE_DTYPE)
+        self.ln0 = RMSNorm(n_embd)
+        self.blocks = nn.ModuleList([
+            Block(n_embd, n_streams, stream_dim, n_read_heads) for _ in range(n_layer)
+        ])
+        self.lm_head = Linear(n_embd, padded_vocab_size, bias=False)
+        self.lm_head.weight = self.wte.weight  # weight tying
+
+    def forward(self, token_ids):
+        logits, _ = self.forward_with_states(token_ids)
+        return logits
+
+    def forward_with_states(self, token_ids):
+        """Full forward; returns (logits, list of per-layer (B,M,D) final states).
+        Used to seed incremental generation from a prompt in one shot."""
+        x = self.ln0(self.wte(token_ids))
+        states = []
+        for block in self.blocks:
+            x, state = block.forward_with_state(x)
+            states.append(state)
+        return self.lm_head(x), states
+
+    def step(self, token_ids, states):
+        """Advance one token using carried streams state. O(L·M·D) per token.
+        token_ids: (B, 1). states: list of (B, M, D). Returns (logits (B,1,V), new states)."""
+        x = self.ln0(self.wte(token_ids))
+        new_states = []
+        for block, state in zip(self.blocks, states):
+            x, new_state = block.step(x, state)
+            new_states.append(new_state)
+        return self.lm_head(x), new_states
+
+    def initial_states(self, batch_size, device, dtype=COMPUTE_DTYPE):
+        return [torch.zeros(batch_size, b.mix.M, b.mix.D, device=device, dtype=dtype)
+                for b in self.blocks]
+
+    def get_memory_footprint(self, return_buffers=True):
+        """Total bytes used by parameters (+ buffers). Matches HF transformers' API."""
+        mem = sum(p.nelement() * p.element_size() for p in self.parameters())
+        if return_buffers:
+            mem += sum(b.nelement() * b.element_size() for b in self.buffers())
+        return mem
+
+    @classmethod
+    def from_config(cls, config):
+        return cls(**config)
+
+    def to_compute_dtype(self, dtype=None):
+        """Cast all parameters and buffers to COMPUTE_DTYPE for inference.
+        Avoids per-call casting in Linear layers — one-time cost at load."""
+        dtype = dtype or COMPUTE_DTYPE
+        self.to(dtype)
+        return self

model_configuration.py

@@ -0,0 +1,33 @@
+from __future__ import annotations
+
+import copy
+from transformers.configuration_utils import PretrainedConfig
+
+
+class StreamMixerConfig(PretrainedConfig):
+    model_type = "streammixer"
+
+    def __init__(
+        self,
+        vocab_size=32768,
+        n_embd=768,
+        n_layer=16,
+        n_streams=48,
+        stream_dim=96,
+        n_read_heads=6,
+        max_sequence_length=2048,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.vocab_size = vocab_size
+        self.n_embd = n_embd
+        self.n_layer = n_layer
+        self.n_streams = n_streams
+        self.stream_dim = stream_dim
+        self.n_read_heads = n_read_heads
+        self.max_sequence_length = max_sequence_length
+        # Aliases expected by HF internals
+        self.hidden_size = n_embd
+        self.num_hidden_layers = n_layer
+        self.num_attention_heads = n_read_heads
+        self.intermediate_size = (int(8 * n_embd / 3) + 15) // 16 * 16

modeling.py

@@ -0,0 +1,384 @@
+from __future__ import annotations
+
+import os
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from transformers.modeling_utils import PreTrainedModel
+from transformers.generation import GenerationMixin
+from transformers.modeling_outputs import BaseModelOutput, CausalLMOutput
+from transformers.configuration_utils import PretrainedConfig
+
+from typing import Optional, Union
+
+# ---------------------------------------------------------------------------
+# Inlined config + model — fully self-contained for HF trust_remote_code.
+# ---------------------------------------------------------------------------
+
+class StreamMixerConfig(PretrainedConfig):
+    model_type = "streammixer"
+
+    def __init__(
+        self,
+        vocab_size=32768,
+        n_embd=768,
+        n_layer=16,
+        n_streams=48,
+        stream_dim=96,
+        n_read_heads=6,
+        max_sequence_length=2048,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.vocab_size = vocab_size
+        self.n_embd = n_embd
+        self.n_layer = n_layer
+        self.n_streams = n_streams
+        self.stream_dim = stream_dim
+        self.n_read_heads = n_read_heads
+        self.max_sequence_length = max_sequence_length
+        self.hidden_size = n_embd
+        self.num_hidden_layers = n_layer
+        self.num_attention_heads = n_read_heads
+        self.intermediate_size = (int(8 * n_embd / 3) + 15) // 16 * 16
+
+# ---------------------------------------------------------------------------
+# Inlined from model.py — keeps modeling.py self-contained for HF cache.
+# ---------------------------------------------------------------------------
+
+_DTYPE_MAP = {"bfloat16": torch.bfloat16, "float16": torch.float16, "float32": torch.float32}
+def _detect_compute_dtype():
+    env = os.environ.get("NANOCHAT_DTYPE")
+    if env is not None:
+        return _DTYPE_MAP[env]
+    if torch.cuda.is_available():
+        capability = torch.cuda.get_device_capability()
+        if capability >= (8, 0):
+            return torch.bfloat16
+        return torch.float32
+    return torch.float32
+
+COMPUTE_DTYPE = _detect_compute_dtype()
+
+
+class Linear(nn.Linear):
+    def forward(self, x):
+        b = None if self.bias is None else self.bias.to(dtype=x.dtype)
+        return F.linear(x, self.weight.to(dtype=x.dtype), b)
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+    def forward(self, x):
+        return F.rms_norm(x, (x.size(-1),))
+
+
+class StreamMixer(nn.Module):
+    LOG_A_MAX_NEG = 0.5
+    CHUNK_SIZE    = 128
+
+    def __init__(self, n_embd, n_streams, stream_dim, n_read_heads):
+        super().__init__()
+        self.M = n_streams
+        self.D = stream_dim
+        self.H = n_read_heads
+        self.wv = Linear(n_embd, stream_dim, bias=False)
+        self.wq = Linear(n_embd, n_read_heads * stream_dim, bias=False)
+        self.wr = Linear(n_embd, n_streams, bias=False)
+        self.wa = Linear(n_embd, n_streams, bias=True)
+        self.wo_out = Linear(n_read_heads * stream_dim, n_embd, bias=False)
+
+    @staticmethod
+    def _rms_norm_last(x):
+        return F.rms_norm(x, (x.size(-1),))
+
+    def _read(self, q_flat, s):
+        lead = q_flat.shape[:-1]
+        q = q_flat.view(*lead, self.H, self.D)
+        q_n = self._rms_norm_last(q)
+        s_n = self._rms_norm_last(s)
+        scores = torch.einsum('...hd,...md->...hm', q_n, s_n) * (self.D ** -0.5)
+        weights = torch.sigmoid(scores)
+        read = torch.einsum('...hm,...md->...hd', weights, s)
+        return read.reshape(*lead, self.H * self.D)
+
+    def _scan(self, log_a, bv):
+        B, T, M, D = bv.shape
+        C = self.CHUNK_SIZE
+        pad = (C - T % C) % C
+        if pad > 0:
+            log_a = F.pad(log_a, (0, 0, 0, pad))
+            bv    = F.pad(bv,    (0, 0, 0, 0, 0, pad))
+        Tp = T + pad
+        n_chunks = Tp // C
+
+        log_a_c = log_a.view(B, n_chunks, C, M)
+        bv_c    = bv.view(B, n_chunks, C, M, D)
+
+        Z = log_a_c.cumsum(dim=2)
+        eZ = torch.exp(Z).unsqueeze(-1)
+        inv_eZ = torch.exp(-Z).unsqueeze(-1)
+
+        s_intra = (bv_c * inv_eZ).cumsum(dim=2) * eZ
+
+        chunk_decay     = eZ[:, :, -1]
+        chunk_end_intra = s_intra[:, :, -1]
+
+        incomings = []
+        state = torch.zeros(B, M, D, device=bv.device, dtype=bv.dtype)
+        for c in range(n_chunks):
+            incomings.append(state)
+            state = chunk_decay[:, c] * state + chunk_end_intra[:, c]
+        incoming = torch.stack(incomings, dim=1)
+
+        s = s_intra + eZ * incoming.unsqueeze(2)
+        return s.view(B, Tp, M, D)[:, :T]
+
+    def forward(self, x):
+        out, _ = self.forward_with_state(x)
+        return out
+
+    def forward_with_state(self, x):
+        v = self.wv(x)
+        q = self.wq(x)
+        r = F.softmax(self.wr(x), dim=-1)
+        log_a = -self.LOG_A_MAX_NEG * torch.sigmoid(self.wa(x))
+        bv = r.unsqueeze(-1) * v.unsqueeze(2)
+        s = self._scan(log_a, bv)
+        read = self._read(q, s)
+        return self.wo_out(read), s[:, -1]
+
+    def step(self, x, state):
+        x_t = x.squeeze(1)
+        v = self.wv(x_t)
+        q = self.wq(x_t)
+        r = F.softmax(self.wr(x_t), dim=-1)
+        log_a = -self.LOG_A_MAX_NEG * torch.sigmoid(self.wa(x_t))
+        a = torch.exp(log_a)
+        write = r.unsqueeze(-1) * v.unsqueeze(1)
+        s = a.unsqueeze(-1) * state + write
+        read = self._read(q, s)
+        return self.wo_out(read).unsqueeze(1), s
+
+
+class MLP(nn.Module):
+    def __init__(self, n_embd):
+        super().__init__()
+        hidden = (int(8 * n_embd / 3) + 15) // 16 * 16
+        self.w_up   = Linear(n_embd, hidden, bias=False)
+        self.w_down = Linear(n_embd, hidden, bias=False)
+        self.w_out  = Linear(hidden, n_embd, bias=False)
+    def forward(self, x):
+        return self.w_out(F.relu(self.w_up(x)).square() * self.w_down(x))
+
+
+class Block(nn.Module):
+    def __init__(self, n_embd, n_streams, stream_dim, n_read_heads):
+        super().__init__()
+        self.ln1 = RMSNorm(n_embd)
+        self.mix = StreamMixer(n_embd, n_streams, stream_dim, n_read_heads)
+        self.ln2 = RMSNorm(n_embd)
+        self.mlp = MLP(n_embd)
+    def forward(self, x):
+        x, _ = self.forward_with_state(x)
+        return x
+    def forward_with_state(self, x):
+        mix_out, state = self.mix.forward_with_state(self.ln1(x))
+        x = x + mix_out
+        x = x + self.mlp(self.ln2(x))
+        return x, state
+    def step(self, x, state):
+        mix_out, new_state = self.mix.step(self.ln1(x), state)
+        x = x + mix_out
+        x = x + self.mlp(self.ln2(x))
+        return x, new_state
+
+
+class GPT(nn.Module):
+    def __init__(self, vocab_size, n_embd, n_layer, n_streams, stream_dim, n_read_heads):
+        super().__init__()
+        padded_vocab_size = ((vocab_size + 63) // 64) * 64
+        self.config = dict(
+            vocab_size=padded_vocab_size, n_embd=n_embd, n_layer=n_layer,
+            n_streams=n_streams, stream_dim=stream_dim, n_read_heads=n_read_heads,
+        )
+        self.wte = nn.Embedding(padded_vocab_size, n_embd, dtype=COMPUTE_DTYPE)
+        self.ln0 = RMSNorm(n_embd)
+        self.blocks = nn.ModuleList([
+            Block(n_embd, n_streams, stream_dim, n_read_heads) for _ in range(n_layer)
+        ])
+        self.lm_head = Linear(n_embd, padded_vocab_size, bias=False)
+        self.lm_head.weight = self.wte.weight
+
+    def forward(self, token_ids):
+        logits, _ = self.forward_with_states(token_ids)
+        return logits
+
+    def forward_with_states(self, token_ids):
+        x = self.ln0(self.wte(token_ids))
+        states = []
+        for block in self.blocks:
+            x, state = block.forward_with_state(x)
+            states.append(state)
+        return self.lm_head(x), states
+
+    def step(self, token_ids, states):
+        x = self.ln0(self.wte(token_ids))
+        new_states = []
+        for block, state in zip(self.blocks, states):
+            x, new_state = block.step(x, state)
+            new_states.append(new_state)
+        return self.lm_head(x), new_states
+
+    def initial_states(self, batch_size, device, dtype=COMPUTE_DTYPE):
+        return [torch.zeros(batch_size, b.mix.M, b.mix.D, device=device, dtype=dtype)
+                for b in self.blocks]
+
+    def get_memory_footprint(self, return_buffers=True):
+        mem = sum(p.nelement() * p.element_size() for p in self.parameters())
+        if return_buffers:
+            mem += sum(b.nelement() * b.element_size() for b in self.buffers())
+        return mem
+
+    @classmethod
+    def from_config(cls, config):
+        return cls(**config)
+
+    def to_compute_dtype(self, dtype=None):
+        dtype = dtype or COMPUTE_DTYPE
+        self.to(dtype)
+        return self
+
+
+# ---------------------------------------------------------------------------
+# HF Model wrappers
+# ---------------------------------------------------------------------------
+
+class StreamMixerPreTrainedModel(PreTrainedModel):
+    config_class = StreamMixerConfig
+    supports_gradient_checkpointing = False
+    base_model_prefix = "model"
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        raise NotImplementedError("Gradient checkpointing is not supported")
+
+    def _init_weights(self, module):
+        pass
+
+
+class StreamMixerModel(StreamMixerPreTrainedModel):
+    def __init__(self, config: StreamMixerConfig, **kwargs):
+        super().__init__(config, **kwargs)
+        self.config = config
+        inner_cfg = {
+            'vocab_size': config.vocab_size,
+            'n_embd': config.n_embd,
+            'n_layer': config.n_layer,
+            'n_streams': config.n_streams,
+            'stream_dim': config.stream_dim,
+            'n_read_heads': config.n_read_heads,
+        }
+        self.inner = GPT.from_config(inner_cfg)
+        self.inner.to(COMPUTE_DTYPE)
+
+    def get_input_embeddings(self):
+        return self.inner.wte
+
+    def set_input_embeddings(self, value):
+        self.inner.wte = value
+
+    def get_output_embeddings(self):
+        return self.inner.lm_head
+
+    def set_output_embeddings(self, value):
+        self.inner.lm_head = value
+
+    def _get_hidden_states(self, input_ids):
+        x = self.inner.ln0(self.inner.wte(input_ids))
+        for block in self.inner.blocks:
+            x = block(x)
+        return x
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        token_type_ids: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        output_hidden_states: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs,
+    ) -> Union[tuple, BaseModelOutput]:
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        hidden_states = self._get_hidden_states(input_ids)
+        if not return_dict:
+            return (hidden_states,)
+        return BaseModelOutput(last_hidden_state=hidden_states)
+
+
+class StreamMixerForCausalLM(GenerationMixin, StreamMixerPreTrainedModel):
+    _keys_to_ignore_on_load_unexpected = set()
+
+    def __init__(self, config: StreamMixerConfig, **kwargs):
+        super().__init__(config, **kwargs)
+        self.model = StreamMixerModel(config, **kwargs)
+        self.vocab_size = config.vocab_size
+        self.config = config
+
+    def __getattr__(self, name):
+        if name == "all_tied_weights_keys":
+            return {}
+        return super().__getattr__(name)
+
+    def get_input_embeddings(self):
+        return self.model.inner.wte
+
+    def set_input_embeddings(self, value):
+        self.model.inner.wte = value
+
+    def get_output_embeddings(self):
+        return self.model.inner.lm_head
+
+    def set_output_embeddings(self, value):
+        self.model.inner.lm_head = value
+
+    def can_generate(self):
+        return True
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        token_type_ids: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        output_hidden_states: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        labels: Optional[torch.LongTensor] = None,
+        **kwargs,
+    ) -> Union[tuple, CausalLMOutput]:
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        hidden_states = self.model._get_hidden_states(input_ids)
+        logits = self.model.inner.lm_head(hidden_states)
+
+        loss = None
+        if labels is not None:
+            shift_logits = logits[:, :-1, :].contiguous()
+            shift_labels = labels[:, 1:].contiguous()
+            loss = F.cross_entropy(
+                shift_logits.view(-1, shift_logits.size(-1)),
+                shift_labels.view(-1),
+            )
+
+        if not return_dict:
+            output = (logits,)
+            return ((loss,) + output) if loss is not None else output
+
+        return CausalLMOutput(
+            loss=loss,
+            logits=logits,
+        )

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ea53044d5dfe2a6cf2005e1e84fea3d18844b5620b973aac965e664824c690a4
+size 92448927

special_tokens_map.json

@@ -0,0 +1,8 @@
+{
+  "bos_token": "<|bos|>",
+  "eos_token": "<|bos|>",
+  "pad_token": "<|bos|>",
+  "additional_special_tokens": [
+    "<|bos|>"
+  ]
+}

tokenizer_config.json

@@ -0,0 +1,16 @@
+{
+  "tokenizer_class": "PreTrainedTokenizerFast",
+  "model_max_length": 1024,
+  "bos_token": "<|bos|>",
+  "eos_token": "<|bos|>",
+  "pad_token": "<|bos|>",
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<|bos|>",
+      "lstrip": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  }
+}