Upload lilm 16K recovery pilot

README.md

@@ -0,0 +1,9 @@
+---
+license: mit
+pipeline_tag: text-generation
+---
+# lilm1-200m 16K recovery pilot
+
+Controlled 299.63M-token continuation from the 2K/27B base checkpoint.
+The mixture is 34% document-contiguous 16K Longmino data and 66% 2K replay.
+This is an experimental base checkpoint and requires the custom `model.py`.

config.json

@@ -0,0 +1,16 @@
+{
+  "vocab_size": 49152,
+  "n_layer": 16,
+  "n_embd": 1024,
+  "n_head": 16,
+  "kv_heads": 4,
+  "intermediate_size": 2752,
+  "kernel_size": 4,
+  "attn_dropout_p": 0.1,
+  "resid_dropout_p": 0.05,
+  "n_conv_layers": 10,
+  "max_seq_len": 16384,
+  "rope_theta": 10000.0,
+  "use_native_gqa": true,
+  "activation_checkpointing": true
+}

model.py

@@ -0,0 +1,434 @@
+"""
+lilm1-200m model definition — lilm hybrid local-convolution/attention model
+═════════════════════════════════════════════════════════════════════════════
+
+Architecture:
+  - 10 local causal-convolution layers + 6 attention layers (16 total)
+  - GQA with n_head=16, kv_heads=4  (4:1 ratio)
+  - SwiGLU FFN (intermediate_size=2752)
+  - RMSNorm (learnable, eps=1e-6)
+  - RoPE positional encoding
+  - Attention via PyTorch SDPA (FlashAttention/math backend auto-dispatch)
+  - Residual dropout on all residual paths
+  - Untied embedding / lm_head weights
+  - Causal depthwise 1D convolution for local token mixing
+
+Paste this entire cell into the "PASTE YOUR MODEL DEFINITION BELOW" cell.
+═════════════════════════════════════════════════════════════════════════════
+"""
+
+from dataclasses import dataclass
+from typing import Optional
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+# ── Config ────────────────────────────────────────────────────────────────────
+
+@dataclass
+class LilmConfig:
+    vocab_size:        int   = 49152
+    n_layer:           int   = 16          # 10 Conv + 6 Attention
+    n_embd:            int   = 1024
+    n_head:            int   = 16
+    kv_heads:          int   = 4           # GQA 4:1
+    intermediate_size: int   = 2752        # SwiGLU hidden dim
+    kernel_size:       int   = 4           # Conv window
+    attn_dropout_p:    float = 0.1         # Passed to PyTorch SDPA
+    resid_dropout_p:   float = 0.05        # nn.Dropout on residual paths
+    n_conv_layers:     int   = 10          # First 10 layers are Conv
+    max_seq_len:       int   = 2048        # For RoPE precomputation
+    rope_theta:        float = 10_000.0    # RoPE base frequency
+    use_native_gqa:     bool  = True        # Use SDPA's native GQA when available
+    activation_checkpointing: bool = False  # Recompute layer activations to save memory
+
+    @property
+    def head_dim(self) -> int:
+        return self.n_embd // self.n_head
+
+    @property
+    def n_attn_layers(self) -> int:
+        return self.n_layer - self.n_conv_layers
+
+
+# ── RMSNorm ───────────────────────────────────────────────────────────────────
+
+class RMSNorm(nn.Module):
+    """Root Mean Square Layer Normalization (Zhang & Sennrich, 2019)."""
+
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        norm = x.float().pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt()
+        return (x.float() * norm).type_as(x) * self.weight
+
+
+# ── RoPE ──────────────────────────────────────────────────────────────────────
+
+def precompute_rope_freqs(
+    head_dim: int,
+    max_seq_len: int,
+    theta: float = 10_000.0,
+    device: Optional[torch.device] = None,
+) -> torch.Tensor:
+    """Precompute complex exponentials for RoPE: shape (max_seq_len, head_dim//2)."""
+    freqs = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim))
+    t = torch.arange(max_seq_len, device=device).float()
+    freqs = torch.outer(t, freqs)                       # (T, head_dim//2)
+    return torch.polar(torch.ones_like(freqs), freqs)    # complex64
+
+
+def apply_rope(
+    x: torch.Tensor,                   # (B, T, n_head, head_dim)
+    freqs_cis: torch.Tensor,           # (T, head_dim//2)
+) -> torch.Tensor:
+    """Apply rotary positional embeddings to q or k."""
+    B, T, H, D = x.shape
+    # View as pairs of (real, imag)
+    x_complex = torch.view_as_complex(x.float().reshape(B, T, H, D // 2, 2))
+    freqs = freqs_cis[:T].unsqueeze(0).unsqueeze(2)     # (1, T, 1, head_dim//2)
+    x_rotated = x_complex * freqs
+    return torch.view_as_real(x_rotated).reshape(B, T, H, D).type_as(x)
+
+
+# ── SwiGLU FFN ────────────────────────────────────────────────────────────────
+
+class SwiGLUFFN(nn.Module):
+    """
+    SwiGLU Feed-Forward Network (Shazeer 2020, Touvron et al. 2023).
+    gate = Swish(x W_gate)
+    out  = (gate ⊙ (x W_up)) W_down
+    """
+
+    def __init__(self, config: LilmConfig):
+        super().__init__()
+        self.w_gate = nn.Linear(config.n_embd, config.intermediate_size, bias=False)
+        self.w_up   = nn.Linear(config.n_embd, config.intermediate_size, bias=False)
+        self.w_down = nn.Linear(config.intermediate_size, config.n_embd, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.w_down(F.silu(self.w_gate(x)) * self.w_up(x))
+
+
+# ── GQA Attention ────────────────────────────────────────────────────────────
+
+class LilmAttention(nn.Module):
+    """
+    Grouped Query Attention using PyTorch scaled_dot_product_attention.
+
+    n_head=16 query heads, kv_heads=4 key/value heads → 4:1 GQA ratio.
+    KV heads are expanded explicitly for broad PyTorch compatibility; SDPA
+    still dispatches to the best available CUDA backend for the resulting
+    dense-head causal attention.
+    """
+
+    def __init__(self, config: LilmConfig):
+        super().__init__()
+        self.n_head   = config.n_head
+        self.kv_heads = config.kv_heads
+        self.head_dim = config.head_dim
+        self.n_rep    = config.n_head // config.kv_heads   # 4 — repeat factor for GQA
+        self.attn_dropout_p = config.attn_dropout_p
+        self.use_native_gqa = (
+            config.use_native_gqa
+            and "enable_gqa" in (F.scaled_dot_product_attention.__doc__ or "")
+        )
+
+        # Projections
+        self.q_proj = nn.Linear(config.n_embd, config.n_head   * config.head_dim, bias=False)
+        self.k_proj = nn.Linear(config.n_embd, config.kv_heads * config.head_dim, bias=False)
+        self.v_proj = nn.Linear(config.n_embd, config.kv_heads * config.head_dim, bias=False)
+        self.o_proj = nn.Linear(config.n_head * config.head_dim, config.n_embd,   bias=False)
+
+    def _expand_kv(self, x: torch.Tensor) -> torch.Tensor:
+        """Repeat KV heads to match query head count: (B,T,kv_heads,D) → (B,T,n_head,D)."""
+        if self.n_rep == 1:
+            return x
+        B, T, H, D = x.shape
+        return x[:, :, :, None, :].expand(B, T, H, self.n_rep, D).reshape(B, T, H * self.n_rep, D)
+
+    def forward(
+        self,
+        x: torch.Tensor,                       # (B, T, n_embd)
+        freqs_cis: torch.Tensor,                # (max_seq_len, head_dim//2)
+    ) -> torch.Tensor:
+        B, T, _ = x.shape
+
+        # Project → (B, T, n_heads, head_dim)
+        q = self.q_proj(x).reshape(B, T, self.n_head,   self.head_dim)
+        k = self.k_proj(x).reshape(B, T, self.kv_heads, self.head_dim)
+        v = self.v_proj(x).reshape(B, T, self.kv_heads, self.head_dim)
+
+        # Apply RoPE to queries and keys
+        q = apply_rope(q, freqs_cis)
+        k = apply_rope(k, freqs_cis)
+
+        dropout_p = self.attn_dropout_p if self.training else 0.0
+
+        # Use PyTorch SDPA, which auto-dispatches to the best available
+        # attention backend on supported hardware.
+        # SDPA expects (B, H, T, D) layout
+        q = q.transpose(1, 2)
+        if self.use_native_gqa:
+            k = k.transpose(1, 2)
+            v = v.transpose(1, 2)
+            out = F.scaled_dot_product_attention(
+                q, k, v,
+                dropout_p=dropout_p,
+                is_causal=True,
+                enable_gqa=True,
+            )
+        else:
+            # Compatibility fallback for PyTorch builds without native GQA.
+            k = self._expand_kv(k).transpose(1, 2)
+            v = self._expand_kv(v).transpose(1, 2)
+            out = F.scaled_dot_product_attention(
+                q, k, v,
+                dropout_p=dropout_p,
+                is_causal=True,
+            )
+        out = out.transpose(1, 2)    # back to (B, T, H, D)
+
+        # Merge heads → project out
+        out = out.reshape(B, T, self.n_head * self.head_dim)
+        return self.o_proj(out)
+
+
+# ── Causal Depthwise Conv ────────────────────────────────────────────────────
+
+class CausalDepthwiseConv1d(nn.Module):
+    """
+    Causal depthwise 1D convolution (Mamba/Hyena style).
+
+    Each channel is convolved independently with its own kernel.
+    Left-padded to preserve causality — output at position t depends
+    only on inputs at positions [t - kernel_size + 1, ..., t].
+    """
+
+    def __init__(self, channels: int, kernel_size: int):
+        super().__init__()
+        self.kernel_size = kernel_size
+        # groups=channels → depthwise (each channel has its own filter)
+        self.conv = nn.Conv1d(
+            in_channels  = channels,
+            out_channels = channels,
+            kernel_size  = kernel_size,
+            padding      = 0,          # We handle causal padding manually
+            groups       = channels,
+            bias         = True,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """x: (B, T, C) → (B, T, C)"""
+        # Transpose to (B, C, T) for Conv1d
+        x = x.transpose(1, 2)
+        # Causal left-padding: pad (kernel_size - 1) zeros on the left
+        x = F.pad(x, (self.kernel_size - 1, 0))
+        x = self.conv(x)
+        return x.transpose(1, 2)       # Back to (B, T, C)
+
+
+# ── Conv Mixer Block (for Conv layers) ───────────────────────────────────────
+
+class ConvMixerBlock(nn.Module):
+    """
+    Conv-based block (used in layers 0–9):
+      x → RMSNorm → CausalDepthwiseConv → SiLU → Linear(project) → residual dropout → +x
+      x → RMSNorm → SwiGLU FFN → residual dropout → +x
+
+    The conv branch replaces attention with a local causal mixing mechanism.
+    SiLU activation after conv + a linear projection keeps dimensionality matched.
+    """
+
+    def __init__(self, config: LilmConfig):
+        super().__init__()
+        self.norm1 = RMSNorm(config.n_embd)
+        self.norm2 = RMSNorm(config.n_embd)
+
+        # Conv branch: depthwise conv → activation → pointwise projection
+        self.conv1d    = CausalDepthwiseConv1d(config.n_embd, config.kernel_size)
+        self.conv_proj = nn.Linear(config.n_embd, config.n_embd, bias=False)
+
+        # FFN
+        self.mlp = SwiGLUFFN(config)
+
+        # Residual dropout
+        self.resid_drop = nn.Dropout(p=config.resid_dropout_p)
+
+    def forward(self, x: torch.Tensor, **kwargs) -> torch.Tensor:
+        # Conv mixing path
+        h = self.norm1(x)
+        h = self.conv1d(h)
+        h = F.silu(h)
+        h = self.conv_proj(h)
+        x = x + self.resid_drop(h)
+
+        # FFN path
+        x = x + self.resid_drop(self.mlp(self.norm2(x)))
+        return x
+
+
+# ── Attention Block (for Attention layers) ───────────────────────────────────
+
+class AttentionBlock(nn.Module):
+    """
+    Standard transformer block (used in layers 10–15):
+      x → RMSNorm → GQA Attention (SDPA) → residual dropout → +x
+      x → RMSNorm → SwiGLU FFN → residual dropout → +x
+    """
+
+    def __init__(self, config: LilmConfig):
+        super().__init__()
+        self.norm1 = RMSNorm(config.n_embd)
+        self.norm2 = RMSNorm(config.n_embd)
+
+        self.attn = LilmAttention(config)
+        self.mlp  = SwiGLUFFN(config)
+
+        self.resid_drop = nn.Dropout(p=config.resid_dropout_p)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        freqs_cis: torch.Tensor,
+    ) -> torch.Tensor:
+        x = x + self.resid_drop(self.attn(self.norm1(x), freqs_cis))
+        x = x + self.resid_drop(self.mlp(self.norm2(x)))
+        return x
+
+
+# ── Full Model ────────────────────────────────────────────────────────────────
+
+class LilmModel(nn.Module):
+    """
+    lilm hybrid local-convolution/attention model.
+
+    - Untied embedding + lm_head
+    - RoPE applied only in attention layers
+    - Pre-norm (RMSNorm) architecture
+    - Final RMSNorm before lm_head
+
+    Forward returns logits: (B, T, vocab_size)
+    """
+
+    def __init__(self, config: LilmConfig):
+        super().__init__()
+        self.config = config
+
+        # Token embedding (untied — separate from lm_head)
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.n_embd)
+
+        # Hybrid layers: first n_conv_layers are Conv, rest are Attention
+        self.layers = nn.ModuleList()
+        for i in range(config.n_layer):
+            if i < config.n_conv_layers:
+                self.layers.append(ConvMixerBlock(config))
+            else:
+                self.layers.append(AttentionBlock(config))
+
+        # Final norm + output head (untied)
+        self.final_norm = RMSNorm(config.n_embd)
+        self.lm_head    = nn.Linear(config.n_embd, config.vocab_size, bias=False)
+
+        # Precompute RoPE frequencies (registered as buffer — not a parameter)
+        freqs_cis = precompute_rope_freqs(
+            config.head_dim,
+            config.max_seq_len,
+            config.rope_theta,
+        )
+        self.register_buffer("freqs_cis", freqs_cis, persistent=False)
+
+        # Initialise weights
+        self._init_weights()
+
+        # Report parameter count
+        n_params = sum(p.numel() for p in self.parameters())
+        print(f"LilmModel — {n_params / 1e6:.2f}M parameters")
+        print(f"  Conv layers : {config.n_conv_layers} (layers 0–{config.n_conv_layers - 1})")
+        print(f"  Attn layers : {config.n_attn_layers} (layers {config.n_conv_layers}–{config.n_layer - 1})")
+        print(f"  GQA ratio   : {config.n_head}:{config.kv_heads}")
+        print(f"  Head dim    : {config.head_dim}")
+
+    def _init_weights(self):
+        """
+        Weight initialisation following GPT-NeoX / LLaMA conventions:
+        - Embeddings:  N(0, 0.02)
+        - Linear:      N(0, 0.02) with residual scaling on output projections
+        - Conv1d:      N(0, 0.02)
+        - RMSNorm:     ones (already default)
+        - Biases:      zeros
+        """
+        residual_scale = 1.0 / math.sqrt(2.0 * self.config.n_layer)
+
+        for name, module in self.named_modules():
+            if isinstance(module, nn.Linear):
+                nn.init.normal_(module.weight, mean=0.0, std=0.02)
+                # Scale down residual-path output projections to stabilise deep nets
+                if any(tag in name for tag in ("o_proj", "w_down", "conv_proj")):
+                    module.weight.data.mul_(residual_scale)
+                if module.bias is not None:
+                    nn.init.zeros_(module.bias)
+            elif isinstance(module, nn.Embedding):
+                nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            elif isinstance(module, nn.Conv1d):
+                nn.init.normal_(module.weight, mean=0.0, std=0.02)
+                if module.bias is not None:
+                    nn.init.zeros_(module.bias)
+
+    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            input_ids: (B, T) token indices
+        Returns:
+            logits: (B, T, vocab_size)
+        """
+        B, T = input_ids.shape
+        assert T <= self.config.max_seq_len, (
+            f"Sequence length {T} exceeds max_seq_len {self.config.max_seq_len}"
+        )
+
+        # Embed tokens
+        x = self.embed_tokens(input_ids)        # (B, T, n_embd)
+
+        # Pass through all layers
+        if self.config.activation_checkpointing and self.training:
+            from torch.utils.checkpoint import checkpoint
+
+            for layer in self.layers:
+                if isinstance(layer, AttentionBlock):
+                    x = checkpoint(
+                        lambda hidden, layer=layer: layer(hidden, freqs_cis=self.freqs_cis),
+                        x,
+                        use_reentrant=False,
+                    )
+                else:
+                    x = checkpoint(layer, x, use_reentrant=False)
+        else:
+            for layer in self.layers:
+                if isinstance(layer, AttentionBlock):
+                    x = layer(x, freqs_cis=self.freqs_cis)
+                else:
+                    x = layer(x)
+
+        # Final norm + logits
+        x = self.final_norm(x)
+        logits = self.lm_head(x)                # (B, T, vocab_size)
+
+        return logits
+
+    def gradient_checkpointing_enable(self):
+        """Enable gradient checkpointing for memory-efficient training."""
+        self.config.activation_checkpointing = True
+
+
+# Deprecated compatibility aliases for older notebooks/checkpoints that imported
+# previous class names. New code should use LilmConfig and LilmModel.
+LFMUpgradedConfig = LilmConfig
+LiquidAttention = LilmAttention
+LiquidUpgradedModel = LilmModel

model.safetensors

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

tokenizer_config.json

@@ -0,0 +1,34 @@
+{
+  "add_prefix_space": false,
+  "backend": "tokenizers",
+  "bos_token": "<|endoftext|>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|endoftext|>",
+  "errors": "replace",
+  "extra_special_tokens": [
+    "<|endoftext|>",
+    "<|im_start|>",
+    "<|im_end|>",
+    "<repo_name>",
+    "<reponame>",
+    "<file_sep>",
+    "<filename>",
+    "<gh_stars>",
+    "<issue_start>",
+    "<issue_comment>",
+    "<issue_closed>",
+    "<jupyter_start>",
+    "<jupyter_text>",
+    "<jupyter_code>",
+    "<jupyter_output>",
+    "<jupyter_script>",
+    "<empty_output>"
+  ],
+  "is_local": false,
+  "local_files_only": false,
+  "model_max_length": 8192,
+  "pad_token": null,
+  "tokenizer_class": "GPT2Tokenizer",
+  "unk_token": "<|endoftext|>",
+  "vocab_size": 49152
+}