Release Limen0.2B

.gitattributes

@@ -1,35 +1,3 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
 *.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+bg.png filter=lfs diff=lfs merge=lfs -text
+superword.model filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,174 @@
+---
+license: apache-2.0
+language:
+- en
+datasets:
+- HuggingFaceTB/smollm-corpus
+- openbmb/Ultra-FineWeb
+- HuggingFaceFW/finepdfs-edu
+- common-pile/peS2o_filtered
+- common-pile/stackv2_edu_filtered
+- allenai/dolma
+library_name: transformers
+tags:
+  - causal-lm
+  - decoder-only
+  - grouped-query-attention
+  - rope
+  - swiglu
+  - boundlessbpe
+  - curriculum-learning
+  - xsa
+pipeline_tag: text-generation
+---
+
+![Limen0.2B](bg.png)
+
+# Limen0.2B
+
+Limen0.2B is a 222.5M-parameter decoder-only base language model trained from scratch on 50B tokens. It supports a 1,024-token context window and uses a custom 16,384-token BoundlessBPE tokenizer. BoundlessBPE learns SuperBPE merges across whitespace, allowing frequent multi-word spans to be represented directly.
+
+This is a base completion model, not an instruction-tuned chat model.
+
+## Requirements
+
+Loading requires PyTorch, Transformers, and the Rust-backed BoundlessBPE package:
+
+```bash
+pip install torch transformers regex heapdict
+pip install "git+https://github.com/UniversalComputingResearch/fastboundlessbpe.git@perf/tokenid-training"
+```
+
+The BoundlessBPE package includes a PyO3 Rust extension. A source installation requires `rustc` and `cargo`; `pip` builds the extension automatically through Maturin when no compatible wheel is available.
+
+## Load
+
+```python
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+model_id = "UniversalComputingResearch/Limen0.2B"
+
+tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    trust_remote_code=True,
+    torch_dtype=torch.bfloat16,
+).cuda().eval()
+
+prompt = "The future of AI is"
+inputs = tokenizer(prompt, return_tensors="pt", add_special_tokens=False).to(model.device)
+with torch.inference_mode():
+    output = model.generate(**inputs, max_new_tokens=64, do_sample=True, temperature=0.8)
+print(tokenizer.decode(output[0], skip_special_tokens=True))
+```
+
+`trust_remote_code=True` is required for the custom model architecture and tokenizer implementation.
+
+## Tokenizer
+
+Superword tokenization extends conventional BPE with learned tokens that can cross pre-tokenization boundaries, allowing frequent multi-word spans, including whitespace, to be represented directly.
+
+Limen0.2B's tokenizer design and implementation build on [SuperBPE: Space Travel for Language Models](https://arxiv.org/abs/2503.13423), [Boundless Byte Pair Encoding: Breaking the Pre-tokenization Barrier](https://arxiv.org/abs/2504.00178), and [Faster Superword Tokenization](https://arxiv.org/abs/2604.05192).
+
+`superword.model` is the artifact used during pretraining. Text is encoded without automatic BOS or EOS insertion. Training documents are terminated with `<|endoftext|>` (ID 16383), which is also the generation stop token. Control-token-like text in ordinary source material is treated as text rather than as markup.
+
+## Architecture
+
+- 222,516,480 parameters; tied input/output embeddings
+- 35 transformer layers, hidden size 768
+- Grouped-query attention: 6 query heads, 2 key/value heads, head dimension 128
+- SwiGLU-style MLP with intermediate size 1,920
+- RoPE (`theta=100000`), RMSNorm, 1,024-token context
+- Linear weights and normalization/control parameters are retained in FP32;
+  matmuls use the activation dtype during BF16 inference/training.
+
+### XSA projection
+
+Each causal attention layer applies scaled-dot-product attention with grouped query attention, followed by an **XSA (exclusive self-attention) projection**. The projection removes the component of each query-head attention output that is parallel to its corresponding normalized value vector before output projection.
+
+## Training
+
+Limen0.2B was trained from scratch for 50B tokens with AdamW (`beta1=0.9`, `beta2=0.95`, `eps=1e-8`). Weight decay was `0.001` and applied only to non-embedding matrix parameters. Gradients were clipped to norm 1.0.
+
+The learning rate warmed up for 500 steps to `0.003`, remained constant through 85% of training, and then decayed linearly to zero. The released weights are an FP32 exponential moving average of the final 10% of training (`decay=0.999`, 4,768 updates), rather than the final raw optimization step.
+
+## Data curriculum
+
+The 50B-token run used a three-stage, token-weighted curriculum with a 10% transition window around stage boundaries. The source weights were:
+
+| Source | 0–33.3% | 33.3–66.7% | 66.7–100% |
+|---|---:|---:|---:|
+| FineWeb-Edu-Dedup | 30% | 30% | 40% |
+| Ultra-FineWeb | 50% | 30% | 20% |
+| FinePDFs-Edu-English | 5% | 10% | 10% |
+| Common-Pile peS2o filtered | 5% | 10% | 10% |
+| StackV2 education-filtered | 5% | 10% | 10% |
+| Dolma selected non-web | 5% | 10% | 10% |
+
+Across the full equal-duration schedule, the target composition is approximately 33.3% FineWeb-Edu-Dedup, 33.3% Ultra-FineWeb, and 8.3% from each of the four specialist sources. The final phase remains 60% web-derived data; it shifts 10 percentage points from Ultra-FineWeb to FineWeb-Edu-Dedup rather than eliminating general-web coverage.
+
+## Zero-shot final evaluation
+
+The final EMA weights were evaluated in zero-shot mode. Multiple-choice results use length-normalized accuracy (`acc_norm`); BLiMP uses pairwise accuracy. Chance-normalized scores map random guessing to 0 and perfect accuracy to 100.
+
+| Benchmark | Accuracy | Chance-normalized score |
+|---|---:|---:|
+| HellaSwag | 41.98% | 22.64 |
+| PIQA | 67.36% | 34.71 |
+| ARC-Easy | 53.37% | 37.82 |
+| ARC-Challenge | 29.69% | 6.26 |
+| CommonsenseQA | 34.15% | 17.69 |
+| BLiMP | 83.28% | 66.55 |
+
+## Reference model comparison
+
+The tables provide contextual comparisons with reported results. Multiple-choice accuracy is `acc_norm`; BLiMP uses pairwise accuracy. The second line in each row of the accuracy table lists parameter count, vocabulary size, and reported training tokens. Bold indicates the highest benchmark result in each column and the smallest value in each metadata category. Training token counts are reported figures; Qwen2.5's 18T is family-level rather than checkpoint-specific, while GPT-2 Medium's approximately 10B is an estimate.
+
+### Accuracy (`acc_norm`)
+
+| Model | HellaSwag | PIQA | ARC-Easy | ARC-Challenge | CommonsenseQA | BLiMP |
+|---|---:|---:|---:|---:|---:|---:|
+| Qwen2.5 0.5B<br>494M · 151,936 vocab · 18T† | **52.1%** | **69.4%** | 58.5% | **32.0%** | 41.0% | 84.4% |
+| Gemma 3 270M<br>270M · 262,144 vocab · 6T | 41.4% | 68.5% | 57.3% | 28.0% | **42.6%** | 82.1% |
+| SmolLM2 135M<br>135M · 49,152 vocab · 2T | 43.2% | 68.2% | **58.7%** | 29.7% | 35.5% | 81.4% |
+| **Limen0.2B**<br>223M · **16,384 vocab** · 50B | 42.0% | 67.4% | 53.4% | 29.7% | 34.2% | 83.3% |
+| GPT-X2 125M<br>**125M** · 32,768 vocab · 75B | 40.5% | 67.1% | 51.6% | 27.6% | 34.6% | 82.9% |
+| GPT-2 Medium<br>355M · 50,257 vocab · **≈10B‡** | 39.3% | 66.6% | 43.5% | 25.0% | 31.7% | **85.2%** |
+
+### Chance-normalized score
+
+| Model | HellaSwag | PIQA | ARC-Easy | ARC-Challenge | CommonsenseQA | BLiMP |
+|---|---:|---:|---:|---:|---:|---:|
+| Qwen2.5 0.5B | **36.2** | **38.8** | 44.6 | **9.3** | 26.2 | 68.7 |
+| Gemma 3 270M | 21.9 | 37.0 | 43.0 | 4.0 | **28.2** | 64.3 |
+| SmolLM2 135M | 24.3 | 36.3 | **44.9** | 6.3 | 19.4 | 62.7 |
+| **Limen0.2B** | 22.6 | 34.7 | 37.8 | 6.3 | 17.7 | 66.6 |
+| GPT-X2 125M | 20.7 | 34.2 | 35.5 | 3.4 | 18.3 | 65.8 |
+| GPT-2 Medium | 19.1 | 33.2 | 24.6 | 0.0 | 14.6 | **70.4** |
+
+- † Qwen reports 18T tokens for the Qwen2.5 pretraining corpus, without a separate total for the 0.5B checkpoint.
+- ‡ GPT-2 reports 40GB of Internet text; ≈10B tokens is a rough conversion using OpenAI's stated heuristic of about four characters per token.
+
+## Checkpoint evaluation trajectory
+
+All evaluations use zero-shot `lm-eval`, BF16 inference, and batch size 64. Task and macro scores are length-normalized accuracy (`acc_norm`). Intermediate checkpoints contain raw training weights; the final row reports the EMA model.
+
+| Checkpoint | Tokens | Val. loss | BPB | ARC-Easy | ARC-Challenge | HellaSwag | PIQA | Norm. macro |
+|---|---:|---:|---:|---:|---:|---:|---:|---:|
+| 1k | 1.049B | 3.0255 | 0.9389 | 38.68% | 22.35% | 29.33% | 59.25% | 37.40% |
+| 5k | 5.243B | 2.6211 | 0.8134 | 46.76% | 25.85% | 34.71% | 62.89% | 42.56% |
+| 10k | 10.486B | 2.5320 | 0.7857 | 47.47% | 26.71% | 37.39% | 64.85% | 44.11% |
+| 20k | 20.972B | 2.4640 | 0.7646 | 51.98% | 27.39% | 38.83% | 65.07% | 45.82% |
+| 30k | 31.457B | 2.4313 | 0.7545 | 52.36% | 30.03% | 40.11% | 66.92% | 47.36% |
+| 40k | 41.943B | 2.4118 | 0.7484 | 50.51% | 29.44% | 40.76% | 67.03% | 46.93% |
+| **Final EMA** | **50.000B** | — | — | **53.37%** | **29.69%** | **41.98%** | **67.36%** | **48.10%** |
+
+Normalized macro accuracy increased from 37.40% at 1k steps to 48.10% for the final model.
+
+## Repository contents
+
+- `model.safetensors`: EMA model weights
+- `config.json`, `config.py`, `model.py`: custom Transformers model definition
+- `superword.model`, `tokenizer_config.json`, `tokenization_superword.py`:
+  tokenizer model and Transformers integration

config.json

@@ -0,0 +1,26 @@
+{
+  "architectures": [
+    "GPTForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "config.GPTConfig",
+    "AutoModelForCausalLM": "model.GPTForCausalLM"
+  },
+  "block_size": 1024,
+  "dtype": "float32",
+  "head_dim": 128,
+  "hidden_size": 768,
+  "intermediate_size": 1920,
+  "labels_are_shifted": true,
+  "max_position_embeddings": 1024,
+  "model_type": "gpt",
+  "num_attention_heads": 6,
+  "num_hidden_layers": 35,
+  "num_key_value_heads": 2,
+  "rms_norm_eps": 1e-06,
+  "rope_theta": 100000.0,
+  "tie_word_embeddings": true,
+  "transformers_version": "5.14.1",
+  "vocab_size": 16384,
+  "xsa_projection": true
+}

config.py

@@ -0,0 +1,67 @@
+"""Standalone Transformers configuration for Limen0.2B."""
+
+from __future__ import annotations
+
+from transformers import PretrainedConfig
+
+
+DEFAULT_VOCAB_SIZE = 16_384
+DEFAULT_HIDDEN_SIZE = 768
+DEFAULT_NUM_HIDDEN_LAYERS = 35
+DEFAULT_NUM_ATTENTION_HEADS = 6
+DEFAULT_NUM_KEY_VALUE_HEADS = 2
+DEFAULT_HEAD_DIM = DEFAULT_HIDDEN_SIZE // DEFAULT_NUM_ATTENTION_HEADS
+DEFAULT_INTERMEDIATE_SIZE = DEFAULT_HIDDEN_SIZE * 5 // 2
+DEFAULT_BLOCK_SIZE = 1024
+DEFAULT_ROPE_THETA = 100_000.0
+
+
+class GPTConfig(PretrainedConfig):
+    """Configuration for the Limen0.2B decoder-only language model."""
+
+    model_type = "gpt"
+
+    def __init__(
+        self,
+        vocab_size: int = DEFAULT_VOCAB_SIZE,
+        hidden_size: int = DEFAULT_HIDDEN_SIZE,
+        num_hidden_layers: int = DEFAULT_NUM_HIDDEN_LAYERS,
+        num_attention_heads: int = DEFAULT_NUM_ATTENTION_HEADS,
+        num_key_value_heads: int | None = DEFAULT_NUM_KEY_VALUE_HEADS,
+        intermediate_size: int | None = DEFAULT_INTERMEDIATE_SIZE,
+        head_dim: int | None = None,
+        block_size: int = DEFAULT_BLOCK_SIZE,
+        rope_theta: float = DEFAULT_ROPE_THETA,
+        rms_norm_eps: float = 1e-6,
+        xsa_projection: bool = True,
+        tie_word_embeddings: bool = True,
+        labels_are_shifted: bool = False,
+        **kwargs,
+    ):
+        if num_key_value_heads is None:
+            num_key_value_heads = num_attention_heads
+        if head_dim is None:
+            if hidden_size % num_attention_heads != 0:
+                raise ValueError("hidden_size must be divisible by num_attention_heads")
+            head_dim = hidden_size // num_attention_heads
+        if intermediate_size is None:
+            intermediate_size = hidden_size * 4
+        if num_attention_heads % num_key_value_heads != 0:
+            raise ValueError("num_attention_heads must be divisible by num_key_value_heads")
+        if head_dim % 2 != 0:
+            raise ValueError("head_dim must be even for RoPE")
+
+        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
+        self.vocab_size = int(vocab_size)
+        self.hidden_size = int(hidden_size)
+        self.num_hidden_layers = int(num_hidden_layers)
+        self.num_attention_heads = int(num_attention_heads)
+        self.num_key_value_heads = int(num_key_value_heads)
+        self.intermediate_size = int(intermediate_size)
+        self.head_dim = int(head_dim)
+        self.block_size = int(block_size)
+        self.max_position_embeddings = int(block_size)
+        self.rope_theta = float(rope_theta)
+        self.rms_norm_eps = float(rms_norm_eps)
+        self.xsa_projection = bool(xsa_projection)
+        self.labels_are_shifted = bool(labels_are_shifted)

configuration_gpt.py

@@ -0,0 +1,27 @@
+"""Backward-compatible exports for the Limen0.2B configuration."""
+
+from .config import (
+    DEFAULT_BLOCK_SIZE,
+    DEFAULT_HEAD_DIM,
+    DEFAULT_HIDDEN_SIZE,
+    DEFAULT_INTERMEDIATE_SIZE,
+    DEFAULT_NUM_ATTENTION_HEADS,
+    DEFAULT_NUM_HIDDEN_LAYERS,
+    DEFAULT_NUM_KEY_VALUE_HEADS,
+    DEFAULT_ROPE_THETA,
+    DEFAULT_VOCAB_SIZE,
+    GPTConfig,
+)
+
+__all__ = [
+    "DEFAULT_BLOCK_SIZE",
+    "DEFAULT_HEAD_DIM",
+    "DEFAULT_HIDDEN_SIZE",
+    "DEFAULT_INTERMEDIATE_SIZE",
+    "DEFAULT_NUM_ATTENTION_HEADS",
+    "DEFAULT_NUM_HIDDEN_LAYERS",
+    "DEFAULT_NUM_KEY_VALUE_HEADS",
+    "DEFAULT_ROPE_THETA",
+    "DEFAULT_VOCAB_SIZE",
+    "GPTConfig",
+]

generation_config.json

@@ -0,0 +1,8 @@
+{
+  "_from_model_config": true,
+  "eos_token_id": 16383,
+  "output_attentions": false,
+  "output_hidden_states": false,
+  "pad_token_id": 16379,
+  "transformers_version": "5.14.1"
+}

model.py

@@ -0,0 +1,357 @@
+import math
+from typing import Optional
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+from torch.nn import functional as F
+from transformers import PreTrainedModel
+from transformers.cache_utils import DynamicCache
+from transformers.generation.utils import GenerationMixin
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+from .config import GPTConfig
+
+
+CONTROL_TENSOR_NAME_PATTERNS = (
+    "scale",
+    "gate",
+    "gain",
+    "norm",
+    "ln_",
+    "rms",
+)
+
+
+def _parameter_count(module: nn.Module) -> int:
+    return sum(parameter.numel() for parameter in module.parameters())
+
+
+def _format_parameter_count(count: int) -> str:
+    return f"{count:,} ({count / 1_000_000:.3f}M)"
+
+
+def parameter_breakdown(model: nn.Module) -> str:
+    """Return a tree-style parameter summary for the GPT architecture."""
+    transformer = model.transformer
+    blocks = list(transformer["h"])
+    embedding = transformer["wte"]
+    final_norm = transformer["ln_f"]
+    first = blocks[0] if blocks else None
+
+    embedding_params = _parameter_count(embedding)
+    final_norm_params = _parameter_count(final_norm)
+    block_params = _parameter_count(first) if first is not None else 0
+    attention_params = _parameter_count(first.attn) if first is not None else 0
+    mlp_params = _parameter_count(first.mlp) if first is not None else 0
+    ln1_params = _parameter_count(first.ln_1) if first is not None else 0
+    ln2_params = _parameter_count(first.ln_2) if first is not None else 0
+    lm_head_params = _parameter_count(model.lm_head)
+    tied_lm_head = bool(model.lm_head.weight is embedding.weight)
+    unique_total = _parameter_count(model)
+
+    def component_lines(prefix: str = "│  │  ") -> list[str]:
+        if first is None:
+            return []
+        attn = first.attn
+        mlp = first.mlp
+        return [
+            f"{prefix}├─ norm 1 ................ {_format_parameter_count(ln1_params)}",
+            f"{prefix}├─ attention .............. {_format_parameter_count(attention_params)}",
+            f"{prefix}│  ├─ q_proj .............. {_format_parameter_count(_parameter_count(attn.q_proj))}",
+            f"{prefix}│  ├─ k_proj .............. {_format_parameter_count(_parameter_count(attn.k_proj))}",
+            f"{prefix}│  ├─ v_proj .............. {_format_parameter_count(_parameter_count(attn.v_proj))}",
+            f"{prefix}│  └─ o_proj .............. {_format_parameter_count(_parameter_count(attn.o_proj))}",
+            f"{prefix}├─ norm 2 ................ {_format_parameter_count(ln2_params)}",
+            f"{prefix}└─ mlp ................... {_format_parameter_count(mlp_params)}",
+            f"{prefix}   ├─ w_gate ............ {_format_parameter_count(_parameter_count(mlp.w_gate))}",
+            f"{prefix}   ├─ w_up .............. {_format_parameter_count(_parameter_count(mlp.w_up))}",
+            f"{prefix}   └─ w_down ............ {_format_parameter_count(_parameter_count(mlp.w_down))}",
+        ]
+
+    lines = [
+        "parameter breakdown:",
+        f"├─ embedding (wte) ........ {_format_parameter_count(embedding_params)}",
+        f"├─ transformer blocks × {len(blocks)}",
+        f"│  ├─ one block ........... {_format_parameter_count(block_params)}",
+        "│  ├─ per-block detail",
+        *component_lines(),
+        f"│  └─ all blocks .......... {_format_parameter_count(block_params * len(blocks))}",
+        f"├─ final norm (ln_f) ....... {_format_parameter_count(final_norm_params)}",
+    ]
+    if tied_lm_head:
+        lines.append(
+            f"├─ lm_head ................ tied to embedding (+0 unique; raw {_format_parameter_count(lm_head_params)})"
+        )
+    else:
+        lines.append(f"├─ lm_head ................ {_format_parameter_count(lm_head_params)}")
+    lines.append(f"└─ total unique ............ {_format_parameter_count(unique_total)}")
+    return "\n".join(lines)
+
+
+class CastedLinear(nn.Linear):
+    """Store linear params in FP32, cast to activation dtype for matmul."""
+
+    def forward(self, x: Tensor) -> Tensor:
+        weight = self.weight.to(dtype=x.dtype)
+        bias = self.bias.to(dtype=x.dtype) if self.bias is not None else None
+        return F.linear(x, weight, bias)
+
+
+def restore_fp32_params(model: nn.Module) -> None:
+    """Keep linear weights and control params in FP32 after dtype conversion."""
+    for module in model.modules():
+        if isinstance(module, CastedLinear):
+            module.float()
+    for name, param in model.named_parameters():
+        if (
+            param.ndim < 2
+            or any(pattern in name for pattern in CONTROL_TENSOR_NAME_PATTERNS)
+        ) and param.dtype != torch.float32:
+            param.data = param.data.float()
+
+
+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):
+        rms = torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps)
+        return (x.float() * rms).to(dtype=x.dtype) * self.weight.to(dtype=x.dtype)
+
+
+def build_rope_inv_freq(head_dim, theta=2500.0):
+    return 1.0 / (theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim))
+
+
+def precompute_rope_cos_sin(head_dim, seq_len, theta=2500.0):
+    freqs = build_rope_inv_freq(head_dim, theta)
+    t = torch.arange(seq_len, dtype=torch.float32)
+    freqs = torch.outer(t, freqs)
+    return freqs.cos(), freqs.sin()
+
+
+def _apply_rope(x, cos, sin):
+    x_float = x.float()
+    x_pair = x_float.reshape(*x_float.shape[:-1], -1, 2)
+    even = x_pair[..., 0]
+    odd = x_pair[..., 1]
+    cos = cos.unsqueeze(0).unsqueeze(0)
+    sin = sin.unsqueeze(0).unsqueeze(0)
+    x_rot = torch.stack((even * cos - odd * sin, even * sin + odd * cos), dim=-1)
+    return x_rot.flatten(-2).type_as(x)
+
+
+def apply_rotary_emb(q, k, freqs_cis):
+    cos, sin = freqs_cis
+    return _apply_rope(q, cos, sin), _apply_rope(k, cos, sin)
+
+class GPTAttention(nn.Module):
+    def __init__(self, config, layer_idx):
+        super().__init__()
+        self.layer_idx = layer_idx
+        self.n_head = config.num_attention_heads
+        self.n_kv_heads = config.num_key_value_heads
+        self.head_dim = config.head_dim
+        self.n_rep = self.n_head // self.n_kv_heads
+        self.xsa_projection = config.xsa_projection
+
+        self.q_proj = CastedLinear(config.hidden_size, self.n_head * self.head_dim, bias=False)
+        self.k_proj = CastedLinear(config.hidden_size, self.n_kv_heads * self.head_dim, bias=False)
+        self.v_proj = CastedLinear(config.hidden_size, self.n_kv_heads * self.head_dim, bias=False)
+        self.o_proj = CastedLinear(self.n_head * self.head_dim, config.hidden_size, bias=False)
+
+    def _xsa_efficient(self, y: Tensor, v_current: Tensor) -> Tensor:
+        # y:         [B, H,   T, D]
+        # v_current: [B, Hkv, T, D]
+        B, H, T, D = y.shape
+        Hkv = v_current.size(1)
+        group = H // Hkv
+
+        y_g = y.reshape(B, Hkv, group, T, D)
+        v_n = F.normalize(v_current, dim=-1).unsqueeze(2)
+
+        proj = (y_g * v_n).sum(dim=-1, keepdim=True) * v_n
+        return (y_g - proj).reshape(B, H, T, D)
+
+    def forward(self, x, freqs_cis, past_key_value=None, use_cache=False, attention_mask=None):
+        B, T, _ = x.size()
+
+        q = self.q_proj(x).view(B, T, self.n_head, self.head_dim).transpose(1, 2)
+        k_current = self.k_proj(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)
+        v_current = self.v_proj(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)
+
+        q, k_current = apply_rotary_emb(q, k_current, freqs_cis)
+
+        if past_key_value is not None:
+            k, v = past_key_value.update(k_current, v_current, self.layer_idx)
+        else:
+            k, v = k_current, v_current
+
+        S = k.size(2)
+
+        is_causal = past_key_value is None or past_key_value.get_seq_length(self.layer_idx) == T
+
+        attn_mask = None
+        if attention_mask is not None:
+            key_pad = attention_mask.to(torch.bool)[:, None, None, :]
+
+            if is_causal and T > 1:
+                causal = torch.ones(T, S, dtype=torch.bool, device=x.device).tril(diagonal=S - T)
+                attn_mask = key_pad & causal[None, None, :, :]
+            else:
+                attn_mask = key_pad.expand(B, 1, T, S)
+
+            is_causal = False
+
+        y = F.scaled_dot_product_attention(
+            q,
+            k,
+            v,
+            attn_mask=attn_mask,
+            is_causal=is_causal,
+            enable_gqa=(self.n_kv_heads != self.n_head),
+        )
+
+        if self.xsa_projection:
+            y = self._xsa_efficient(y, v_current)
+
+        y = y.transpose(1, 2).contiguous().view(B, T, self.n_head * self.head_dim)
+        return self.o_proj(y)
+
+
+class GPTMLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.w_gate = CastedLinear(config.hidden_size, config.intermediate_size, bias=False)
+        self.w_up = CastedLinear(config.hidden_size, config.intermediate_size, bias=False)
+        self.w_down = CastedLinear(config.intermediate_size, config.hidden_size, bias=False)
+
+    def forward(self, x):
+        return self.w_down(F.silu(self.w_gate(x)) * self.w_up(x))
+
+
+class GPTBlock(nn.Module):
+    def __init__(self, config, layer_idx):
+        super().__init__()
+        self.ln_1 = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.attn = GPTAttention(config, layer_idx)
+        self.ln_2 = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.mlp = GPTMLP(config)
+
+    def forward(self, x, freqs_cis, past_key_value=None, use_cache=False, attention_mask=None):
+        x = x + self.attn(self.ln_1(x), freqs_cis, past_key_value, use_cache, attention_mask=attention_mask)
+        x = x + self.mlp(self.ln_2(x))
+        return x
+
+
+class GPTPreTrainedModel(PreTrainedModel):
+    config_class = GPTConfig
+    base_model_prefix = "transformer"
+    supports_gradient_checkpointing = False
+
+    def _init_weights(self, module):
+        std = self.config.hidden_size ** -0.5
+        if isinstance(module, nn.Linear):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=std)
+        elif isinstance(module, nn.Embedding):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=std)
+
+
+class GPTForCausalLM(GPTPreTrainedModel, GenerationMixin):
+    _tied_weights_keys = {"lm_head.weight": "transformer.wte.weight"}
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.config = config
+        self.transformer = nn.ModuleDict(dict(
+            wte=nn.Embedding(config.vocab_size, config.hidden_size),
+            h=nn.ModuleList([GPTBlock(config, i) for i in range(config.num_hidden_layers)]),
+            ln_f=RMSNorm(config.hidden_size, eps=config.rms_norm_eps),
+        ))
+        self.lm_head = CastedLinear(config.hidden_size, config.vocab_size, bias=False)
+        if config.tie_word_embeddings:
+            self.lm_head.weight = self.transformer["wte"].weight
+        self._freqs_cis_cache = None
+        self.post_init()
+        restore_fp32_params(self)
+
+    def _apply(self, fn):
+        module = super()._apply(fn)
+        restore_fp32_params(self)
+        return module
+
+    def get_input_embeddings(self):
+        return self.transformer["wte"]
+
+    def set_input_embeddings(self, value):
+        self.transformer["wte"] = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, attention_mask=None, **kwargs):
+        if past_key_values is not None and past_key_values.get_seq_length() > 0:
+            input_ids = input_ids[:, -1:]
+        return {
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "past_key_values": past_key_values,
+            "use_cache": True,
+        }
+
+    def _get_freqs_cis(self, seq_len, device):
+        cache = self._freqs_cis_cache
+        if cache is None or cache[0].device != device or cache[0].size(0) < seq_len:
+            cache = tuple(
+                tensor.to(device)
+                for tensor in precompute_rope_cos_sin(self.config.head_dim, seq_len, self.config.rope_theta)
+            )
+            if torch.is_inference_mode_enabled():
+                return cache[0][:seq_len], cache[1][:seq_len]
+            self._freqs_cis_cache = cache
+        return cache[0][:seq_len], cache[1][:seq_len]
+
+    def forward(
+        self,
+        input_ids,
+        attention_mask=None,
+        labels=None,
+        past_key_values: Optional[DynamicCache] = None,
+        use_cache=False,
+        **kwargs,
+    ):
+        B, T = input_ids.size()
+        if use_cache and past_key_values is None:
+            past_key_values = DynamicCache()
+
+        past_len = past_key_values.get_seq_length() if past_key_values is not None else 0
+        x = self.transformer["wte"](input_ids)
+        cos, sin = self._get_freqs_cis(past_len + T, input_ids.device)
+        freqs_cis = cos[past_len:], sin[past_len:]
+
+        for block in self.transformer["h"]:
+            x = block(x, freqs_cis, past_key_values if use_cache else None, use_cache, attention_mask=attention_mask)
+
+        x = self.transformer["ln_f"](x)
+        logits = self.lm_head(x)
+
+        loss = None
+        if labels is not None:
+            if getattr(self.config, "labels_are_shifted", False):
+                loss = F.cross_entropy(logits.float().reshape(-1, logits.size(-1)), labels.reshape(-1))
+            else:
+                shift_logits = logits[..., :-1, :].contiguous()
+                shift_labels = labels[..., 1:].contiguous()
+                loss = F.cross_entropy(shift_logits.float().view(-1, shift_logits.size(-1)), shift_labels.reshape(-1))
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=past_key_values if use_cache else None,
+        )

model.safetensors

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

requirements.txt

@@ -0,0 +1,5 @@
+torch
+transformers>=4.57
+regex
+heapdict
+boundlessbpe @ git+https://github.com/UniversalComputingResearch/fastboundlessbpe.git@53a32afc808226f81e4cb203ae3a3d7f11f7b60f

special_tokens_map.json

@@ -0,0 +1,6 @@
+{
+  "bos_token": "<|bos|>",
+  "eos_token": "<|endoftext|>",
+  "pad_token": "<|pad|>",
+  "unk_token": "<|unk|>"
+}

superword.model

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

tokenization_superword.py

@@ -0,0 +1,124 @@
+"""Hugging Face adapter for Limen0.2B's SuperBPE tokenizer.
+
+Install the Rust-backed tokenizer package before loading this tokenizer:
+
+    pip install "git+https://github.com/UniversalComputingResearch/fastboundlessbpe.git@perf/tokenid-training"
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+from transformers import PreTrainedTokenizer
+
+try:
+    from boundlessbpe import FastTokenizer, RUST_AVAILABLE
+    from boundlessbpe.vocabulary import Vocabulary
+except ImportError as exc:  # pragma: no cover - depends on the consumer environment
+    raise ImportError(
+        "Limen0.2B requires the Rust-backed `boundlessbpe` package. Install it with: "
+        'pip install "git+https://github.com/UniversalComputingResearch/fastboundlessbpe.git@perf/tokenid-training"'
+    ) from exc
+
+
+class SuperwordTokenizer(PreTrainedTokenizer):
+    """Exact inference adapter for the SuperBPE model used in pretraining."""
+
+    model_input_names = ["input_ids", "attention_mask"]
+    vocab_files_names = {"superword_model_file": "superword.model"}
+
+    def __init__(self, superword_model_file: str = "superword.model", **kwargs):
+        if not RUST_AVAILABLE or FastTokenizer is None:
+            raise RuntimeError(
+                "`boundlessbpe` is installed without its Rust extension. Reinstall the "
+                "package from https://github.com/UniversalComputingResearch/fastboundlessbpe/tree/perf/tokenid-training."
+            )
+
+        model_file = Path(superword_model_file)
+        if not model_file.is_absolute():
+            model_file = Path(kwargs.pop("name_or_path", ".")) / model_file
+        self.superword_model_file = str(model_file)
+
+        self._fast = FastTokenizer()
+        self._fast.load(str(model_file))
+        with model_file.open("r", encoding="utf-8") as model_handle:
+            header = model_handle.readline().strip()
+            if not header.startswith("BoundlessBPE v2 "):
+                raise ValueError(f"Unsupported SuperBPE model header: {header!r}")
+            self._vocabulary = Vocabulary.load(model_handle)
+
+        self._special_tokens = dict(self._vocabulary.special_tokens)
+        self._inverse_special_tokens = dict(self._vocabulary.inverse_special_tokens)
+        self._vocab = {
+            token.decode("utf-8", errors="replace"): int(token_id)
+            for token, token_id in self._vocabulary.token_to_id.items()
+        }
+        self._vocab.update(self._special_tokens)
+
+        model_max_length = int(kwargs.pop("model_max_length", 1024))
+        for key in (
+            "pad_token",
+            "bos_token",
+            "eos_token",
+            "unk_token",
+        ):
+            kwargs.pop(key, None)
+        super().__init__(
+            pad_token="<|pad|>",
+            bos_token="<|bos|>",
+            eos_token="<|endoftext|>",
+            unk_token="<|unk|>",
+            model_max_length=model_max_length,
+            **kwargs,
+        )
+
+    def get_vocab(self):
+        return dict(self._vocab)
+
+    @property
+    def vocab_size(self):
+        return int(self._fast.get_vocab_size(with_added_tokens=False))
+
+    def _id_to_token(self, token_id: int) -> str:
+        token = self._vocabulary.id_to_token.get(int(token_id))
+        if token is not None:
+            return token.decode("utf-8", errors="replace")
+        return self._inverse_special_tokens.get(int(token_id), "<|unk|>")
+
+    def _tokenize(self, text, **kwargs):
+        return [self._id_to_token(token_id) for token_id in self._fast.encode_ordinary(text)]
+
+    def _convert_token_to_id(self, token):
+        return self._vocab.get(token, self._special_tokens["<|unk|>"])
+
+    def _convert_id_to_token(self, index):
+        return self._id_to_token(int(index))
+
+    def encode(self, text, text_pair=None, add_special_tokens=False, **kwargs):
+        if text_pair is not None:
+            text = text + text_pair
+        if add_special_tokens:
+            return list(self._fast.encode(text, allowed_special="all"))
+        return list(self._fast.encode_ordinary(text))
+
+    def decode(self, token_ids, skip_special_tokens=True, **kwargs):
+        if isinstance(token_ids, int):
+            token_ids = [token_ids]
+        if skip_special_tokens:
+            token_ids = [
+                token_id
+                for token_id in token_ids
+                if int(token_id) not in self._inverse_special_tokens
+            ]
+        return self._fast.decode(list(token_ids))
+
+    def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
+        if token_ids_1 is None:
+            return list(token_ids_0)
+        return list(token_ids_0) + list(token_ids_1)
+
+    def save_vocabulary(self, save_directory, filename_prefix=None):
+        target = Path(save_directory) / (filename_prefix or "")
+        target = target.with_name(target.name + "superword.model")
+        target.write_bytes(Path(self.superword_model_file).read_bytes())
+        return (str(target),)

tokenizer_config.json

@@ -0,0 +1,12 @@
+{
+  "auto_map": [
+    "tokenization_superword.SuperwordTokenizer",
+    null
+  ],
+  "bos_token": "<|bos|>",
+  "eos_token": "<|endoftext|>",
+  "model_max_length": 1024,
+  "pad_token": "<|pad|>",
+  "superword_model_file": "superword.model",
+  "unk_token": "<|unk|>"
+}