Upload model files

README.md

@@ -0,0 +1,320 @@
+---
+library_name: transformers
+pipeline_tag: text-generation
+language:
+  - en
+tags:
+  - causal-lm
+  - base-model
+  - custom-code
+  - safetensors
+  - research
+  - fixed-token-codes
+  - frozen-input-representations
+---
+
+# AB-EXT Binary16 — 1.711B parameters, 100B-token target
+
+A **base pretrained decoder-only causal language model** released for
+research on trainable input embedding tables and fixed token identities.
+
+It is not an instruction-tuned or preference-optimized assistant.
+
+## Research question
+
+Can a shared contextual network learn useful language-modeling behavior
+without independently trainable token-specific input vectors?
+
+The controlled family contains a learned-input model, a canonical
+16-bit-code model, and an invertibly recoded GF2 model. They share the
+contextual backbone and output-head architecture, but not the same total
+trainable parameter count.
+
+**The results support viability, not performance equivalence.**
+The fixed-code models retain substantial capability while the learned
+model performs better on several informative evaluations.
+
+## Model specification
+
+| Property | Value |
+|---|---|
+| Input mode | `binary16` |
+| Trainable parameters | 1,711,376,384 |
+| Trainable input parameters | 0 |
+| Trainable body parameters, excluding input/output | 1,610,713,088 |
+| Trainable untied output-head parameters | 100,663,296 |
+| Persistent input-buffer values | 786,432 |
+| Hidden width | 2048 |
+| Decoder blocks | 24 |
+| Attention heads | 32 |
+| FFN intermediate width | 8192 |
+| Training context | 2048 tokens |
+| Position encoding | RoPE |
+| Normalization / activation | RMSNorm / SwiGLU |
+| Tokenizer source | `HuggingFaceTB/SmolLM2-1.7B` |
+| Exported tokenizer revision | effd688a12921b4cc83e3312b6feb579f70f9c71 |
+| Evaluated training runs for this interface | One |
+
+The stored tensor-value count includes buffers and must not be reported
+as the trainable parameter count.
+
+## Input representation
+
+
+Each token ID is represented by its canonical little-endian binary code:
+
+$$
+c(t)_j =
+\left\lfloor \frac{t}{2^j} \right\rfloor \bmod 2,
+\qquad j=0,\ldots,15.
+$$
+
+Since the vocabulary contains 49,152 entries, an injective fixed-length
+binary code requires 16 bits.
+
+The code is repeated 128 times to width 2048:
+
+$$
+x(t)=
+\underbrace{c(t)\Vert\cdots\Vert c(t)}_{128\text{ copies}}.
+$$
+
+There are **zero trainable input-interface parameters** and no additional
+trainable input projection before the standard backbone.
+The backbone and the untied output vocabulary projection remain trainable.
+
+The evaluated implementation stores the 49,152-by-16 codebook as a
+persistent non-trainable buffer. It is not literally lookup-free.
+“Minimal” describes fixed-length binary identity width, not the storage
+of the complete model or an entropy-optimal token code.
+
+
+
+## Training
+
+- **Training tokens:** Approximately 100B according to the run report; target budget 100,000,000,000 prediction targets. 
+- **Training precision:** FP32 parameters with BF16 autocast in the supplied trainer.
+- **Reported recipe:** AdamW; peak learning rate 0.00015; minimum
+  scheduled learning rate 0.00001; 2000 warmup steps; cosine decay;
+  weight decay 0.01; betas 0.9 and 0.95; gradient clipping 1.0.
+- **Reported launch geometry:** two GPUs per run, microbatch eight per
+  GPU, eight accumulation steps, sequence length 2048.
+
+The launch geometry corresponds to 262,144 prediction targets per
+optimizer step. Exact final counts must come from the checkpoint,
+not from the requested budget.
+
+The supplied sampler selects within-document windows from eligible
+documents of at least 2049 tokens. Sampling can repeat or overlap
+windows; 100B processed targets does not imply 100B unique corpus tokens.
+The original trainer does not fully restore per-rank sampling state
+on resume. A shared recipe alone does not establish identical realized
+sample order across interrupted runs.
+
+The model weights were NOT initialized from SmolLM2.
+SmolLM2 supplies tokenizer artifacts, not pretrained model weights.
+
+## Evaluation results
+
+These scores are transcribed from the supplied completed evaluation
+summary; the generator does not rerun benchmarks. Raw, unrounded
+harness outputs remain authoritative.
+
+Accuracy entries are percentages. Their reported `±` values are
+evaluation standard errors, **not variation across training seeds**.
+Perplexities and bits per byte are not percentages.
+
+| Metric | Shots | Result |
+|---|---:|---:|
+| HellaSwag acc (%) | 0 | 40.70 ± 0.49 |
+| HellaSwag acc_norm (%) | 0 | 52.40 ± 0.50 |
+| ARC-Easy acc (%) | 0 | 67.59 ± 0.96 |
+| ARC-Easy acc_norm (%) | 0 | 61.53 ± 1.00 |
+| ARC-Challenge acc (%) | 0 | 32.34 ± 1.37 |
+| ARC-Challenge acc_norm (%) | 0 | 34.04 ± 1.38 |
+| PIQA acc (%) | 0 | 70.51 ± 1.06 |
+| PIQA acc_norm (%) | 0 | 71.11 ± 1.06 |
+| WinoGrande acc (%) | 0 | 55.33 ± 1.40 |
+| OpenBookQA acc (%) | 0 | 28.20 ± 2.01 |
+| OpenBookQA acc_norm (%) | 0 | 38.00 ± 2.17 |
+| CommonsenseQA acc (%) | 0 | 20.56 ± 1.16 |
+| MMLU acc (%) | 0 | 25.88 ± 0.37 |
+| MMLU acc (%; some prompts truncated) | 5 | 25.57 ± 0.37 |
+| LAMBADA accuracy (%) | 0 | 42.75 ± 0.69 |
+| LAMBADA perplexity ↓ | 0 | 17.91 ± 0.62 |
+| WikiText word perplexity ↓ | — | 18.58 |
+| WikiText byte perplexity ↓ | — | 1.73 |
+| WikiText bits/byte ↓ | — | 0.79 |
+
+### Audit and coverage limitations
+
+The supplied audit reports identical sample/prompt multisets across
+all six models in each completed task group.
+
+For MMLU 5-shot, **1,508 / 56,168 candidate log-likelihood requests**
+were marked as truncated for each model, approximately 2.68%.
+These are candidate requests, not necessarily distinct questions.
+The displayed MMLU 5-shot score therefore includes truncated prompts.
+
+No truncations were reported for the other groups by that audit.
+For WikiText rolling likelihood, this does not mean that whole documents
+fit into one model context: rolling windowing is part of scoring.
+
+<details>
+<summary>Full six-model comparison</summary>
+
+| Metric | AB-EXT Learned | AB-EXT Binary16 | AB-EXT GF2 | SmolLM2-135M | SmolLM2-360M | SmolLM2-1.7B |
+|---|---:|---:|---:|---:|---:|---:|
+| HellaSwag acc (%); shots=0 | 44.21 ± 0.50 | 40.70 ± 0.49 | 40.24 ± 0.49 | 35.36 ± 0.48 | 43.05 ± 0.49 | 53.38 ± 0.50 |
+| HellaSwag acc_norm (%); shots=0 | 57.79 ± 0.49 | 52.40 ± 0.50 | 51.44 ± 0.50 | 43.02 ± 0.49 | 56.28 ± 0.50 | 71.43 ± 0.45 |
+| ARC-Easy acc (%); shots=0 | 71.63 ± 0.92 | 67.59 ± 0.96 | 66.84 ± 0.97 | 64.44 ± 0.98 | 70.24 ± 0.94 | 77.86 ± 0.85 |
+| ARC-Easy acc_norm (%); shots=0 | 66.04 ± 0.97 | 61.53 ± 1.00 | 60.73 ± 1.00 | 58.75 ± 1.01 | 68.18 ± 0.96 | 73.36 ± 0.91 |
+| ARC-Challenge acc (%); shots=0 | 35.92 ± 1.40 | 32.34 ± 1.37 | 30.55 ± 1.35 | 28.07 ± 1.31 | 36.26 ± 1.40 | 44.37 ± 1.45 |
+| ARC-Challenge acc_norm (%); shots=0 | 37.63 ± 1.42 | 34.04 ± 1.38 | 34.22 ± 1.39 | 29.61 ± 1.33 | 38.05 ± 1.42 | 47.27 ± 1.46 |
+| PIQA acc (%); shots=0 | 72.69 ± 1.04 | 70.51 ± 1.06 | 71.16 ± 1.06 | 68.44 ± 1.08 | 71.38 ± 1.05 | 76.99 ± 0.98 |
+| PIQA acc_norm (%); shots=0 | 72.14 ± 1.05 | 71.11 ± 1.06 | 72.14 ± 1.05 | 68.39 ± 1.08 | 71.82 ± 1.05 | 77.20 ± 0.98 |
+| WinoGrande acc (%); shots=0 | 58.56 ± 1.38 | 55.33 ± 1.40 | 55.01 ± 1.40 | 52.57 ± 1.40 | 59.35 ± 1.38 | 65.98 ± 1.33 |
+| OpenBookQA acc (%); shots=0 | 27.60 ± 2.00 | 28.20 ± 2.01 | 25.20 ± 1.94 | 22.00 ± 1.85 | 24.80 ± 1.93 | 32.20 ± 2.09 |
+| OpenBookQA acc_norm (%); shots=0 | 37.80 ± 2.17 | 38.00 ± 2.17 | 36.80 ± 2.16 | 32.60 ± 2.10 | 37.80 ± 2.17 | 44.40 ± 2.22 |
+| CommonsenseQA acc (%); shots=0 | 19.82 ± 1.14 | 20.56 ± 1.16 | 19.74 ± 1.14 | 19.90 ± 1.14 | 21.05 ± 1.17 | 41.69 ± 1.41 |
+| MMLU acc (%); shots=0 | 25.32 ± 0.37 | 25.88 ± 0.37 | 26.11 ± 0.37 | 24.25 ± 0.36 | 25.47 ± 0.37 | 48.40 ± 0.41 |
+| MMLU acc (%; some prompts truncated); shots=5 | 25.48 ± 0.37 | 25.57 ± 0.37 | 24.66 ± 0.36 | 25.15 ± 0.36 | 25.03 ± 0.37 | 50.06 ± 0.41 |
+| LAMBADA accuracy (%); shots=0 | 47.72 ± 0.70 | 42.75 ± 0.69 | 42.29 ± 0.69 | 42.97 ± 0.69 | 53.31 ± 0.70 | 67.51 ± 0.65 |
+| LAMBADA perplexity ↓; shots=0 | 12.88 ± 0.42 | 17.91 ± 0.62 | 18.47 ± 0.63 | 19.06 ± 0.63 | 9.38 ± 0.27 | 4.44 ± 0.10 |
+| WikiText word perplexity ↓; shots=— | 16.50 | 18.58 | 19.03 | 23.14 | 17.12 | 11.62 |
+| WikiText byte perplexity ↓; shots=— | 1.69 | 1.73 | 1.73 | 1.80 | 1.70 | 1.58 |
+| WikiText bits/byte ↓; shots=— | 0.76 | 0.79 | 0.79 | 0.85 | 0.77 | 0.66 |
+
+</details>
+
+### How to interpret SmolLM2 comparisons
+
+All scores above are from the supplied local evaluation summary, not
+copied leaderboard scores.
+
+The SmolLM2 technical report gives approximate training budgets of:
+
+| External reference | Published budget | Relative to 100B |
+|---|---:|---:|
+| SmolLM2-135M | 2T tokens | 20× |
+| SmolLM2-360M | 4T tokens | 40× |
+| SmolLM2-1.7B | 11T tokens | 110× |
+
+Source: https://arxiv.org/abs/2502.02737
+
+These models differ in architecture, size, data, training schedule,
+and compute. They are quality references, **not matched controls** and
+not proof of a sample-efficiency advantage.
+
+## Usage
+
+Review the custom Python files before enabling `trust_remote_code=True`.
+Use a tested Transformers version and pin the Hub revision for
+reproducible deployment.
+
+```python
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+model_id = 'E6E831728/ab_ext_binary16'
+# For published Hub use, pin revision to a reviewed commit.
+revision = None
+
+tokenizer = AutoTokenizer.from_pretrained(
+    model_id,
+    revision=revision,
+)
+
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    revision=revision,
+    trust_remote_code=True,
+    dtype=torch.bfloat16,
+).to("cuda").eval()
+
+inputs = tokenizer(
+    "Gravity is",
+    return_tensors="pt",
+    add_special_tokens=False,
+    return_attention_mask=True,
+).to("cuda")
+
+pad_id = tokenizer.pad_token_id
+if pad_id is None:
+    pad_id = tokenizer.eos_token_id
+
+with torch.inference_mode():
+    output = model.generate(
+        input_ids=inputs["input_ids"],
+        attention_mask=inputs["attention_mask"],
+        max_new_tokens=32,
+        do_sample=False,
+        use_cache=False,
+        eos_token_id=tokenizer.eos_token_id,
+        pad_token_id=pad_id,
+    )
+
+print(tokenizer.decode(output[0], skip_special_tokens=True))
+```
+
+The implementation does not provide a KV cache.
+The trained context is 2048 tokens; the supplied generation adapter
+uses a sliding window when the context grows beyond its limit.
+This is not evidence of trained long-context capability.
+
+### Loss API
+
+The original training model consumes already-shifted targets.
+The HF runtime is intended to expose the usual causal-LM convention
+with an internal label shift. Do not pass already-shifted labels to
+such a runtime.
+
+Before fine-tuning, verify the actual runtime's loss implementation.
+Forward-logit equivalence does not by itself test label conventions.
+
+## Verification and integrity
+
+The supplied verification logs report:
+
+- successful BF16 loading and generation for all three releases;
+- exactly matching original/exported forward logits on four short
+  prompts for each model, in the tested verification configuration.
+
+These are smoke and implementation-parity checks, not an exhaustive
+test across padding, context lengths, dtypes, or generation modes.
+
+- Weight file: `model.safetensors`
+- Weight SHA-256: `c3e78f24f03dff36b6174985dc5189e8ce37a52685d022ade36fa68c3239801b`
+- Stored tensor values: 1,712,162,816
+- Stored values by dtype: `{"F32": 1712162816}`
+- Trainable parameter count: 1,711,376,384
+- Persistent input-buffer values: 786,432
+
+This card update does not modify the weights, tokenizer, model code,
+or configuration.
+
+## Limitations and intended use
+
+- Research use and text completion; not a validated high-stakes assistant.
+- One evaluated training run per input interface at this scale.
+- Fixed-code and learned-input models are backbone-matched, not
+  total-parameter-matched.
+- No measured runtime or energy advantage is established by parameter
+  counts alone.
+- One GF2 recoding does not establish invariance to arbitrary codes.
+- The output vocabulary matrix remains trainable and token-specific.
+- Input-code structure is not fitted to the pretraining objective, but
+  the tokenizer and its ID assignment can contain corpus-derived structure.
+- Benchmark contamination has not been independently certified absent.
+- Generated text can be false, biased, or harmful.
+
+- Zero/one coding maps token ID zero to a zero input vector; the
+  zero-offset GF2 transform preserves it. In the supplied bias-free
+  architecture, a context made entirely of zero-code tokens gives
+  uniform logits. This does not apply to arbitrary contexts ending
+  in that token.
+
+
+## Attribution and licensing
+
+Tokenizer artifacts are sourced from `HuggingFaceTB/SmolLM2-1.7B`.

config.json

@@ -0,0 +1,44 @@
+{
+  "architectures": [
+    "AttnExtForCausalLM"
+  ],
+  "attention_bias": false,
+  "auto_map": {
+    "AutoConfig": "configuration_attn_ext.AttnExtConfig",
+    "AutoModelForCausalLM": "modeling_attn_ext.AttnExtForCausalLM"
+  },
+  "binary_dim": 16,
+  "binary_encoding": "zero_one",
+  "binary_repeat": 128,
+  "binary_scale": 1.0,
+  "block_size": 2048,
+  "bos_token_id": 0,
+  "code_seed": 12345,
+  "d_model": 2048,
+  "dropout": 0.0,
+  "eos_token_id": 0,
+  "ffn_multiplier": 4.0,
+  "head_dim": 64,
+  "hidden_size": 2048,
+  "initializer_range": 0.02,
+  "input_mode": "binary16",
+  "is_decoder": true,
+  "is_encoder_decoder": false,
+  "max_position_embeddings": 2048,
+  "min_col_weight": 4,
+  "min_row_weight": 4,
+  "mlp_bias": false,
+  "model_type": "attn_ext",
+  "multiple_of": 256,
+  "n_head": 32,
+  "n_layer": 24,
+  "num_attention_heads": 32,
+  "num_hidden_layers": 24,
+  "pad_token_id": null,
+  "rms_norm_eps": 1e-05,
+  "rope_theta": 10000.0,
+  "tie_word_embeddings": false,
+  "transformers_version": null,
+  "use_cache": false,
+  "vocab_size": 49152
+}

configuration_attn_ext.py

@@ -0,0 +1,113 @@
+from transformers import PretrainedConfig
+
+
+class AttnExtConfig(PretrainedConfig):
+    model_type = "attn_ext"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        vocab_size=49152,
+        d_model=2048,
+        n_layer=24,
+        n_head=32,
+        ffn_multiplier=4.0,
+        multiple_of=256,
+        block_size=2048,
+        rope_theta=10000.0,
+        dropout=0.0,
+        rms_norm_eps=1e-5,
+        initializer_range=0.02,
+        attention_bias=False,
+        mlp_bias=False,
+        input_mode="learned",
+        binary_dim=16,
+        binary_encoding="zero_one",
+        binary_scale=1.0,
+        code_seed=12345,
+        min_row_weight=4,
+        min_col_weight=4,
+        pad_token_id=None,
+        bos_token_id=None,
+        eos_token_id=None,
+        tie_word_embeddings=False,
+        use_cache=False,
+        **kwargs,
+    ):
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+
+        if d_model % n_head != 0:
+            raise ValueError("d_model must be divisible by n_head")
+
+        head_dim = d_model // n_head
+        if head_dim % 2 != 0:
+            raise ValueError("RoPE requires an even head dimension")
+
+        if input_mode not in {"learned", "binary16", "gf2"}:
+            raise ValueError(
+                "input_mode must be learned, binary16, or gf2"
+            )
+
+        if input_mode != "learned":
+            if binary_dim != 16:
+                raise ValueError("Frozen-code models require binary_dim=16")
+            if vocab_size > 2**binary_dim:
+                raise ValueError("Vocabulary does not fit in 16 bits")
+            if d_model % binary_dim != 0:
+                raise ValueError(
+                    "d_model must be divisible by binary_dim"
+                )
+            if tie_word_embeddings:
+                raise ValueError(
+                    "Frozen input codes cannot be tied to lm_head"
+                )
+
+        if binary_encoding not in {"zero_one", "bipolar"}:
+            raise ValueError(
+                "binary_encoding must be zero_one or bipolar"
+            )
+
+        self.vocab_size = vocab_size
+
+        self.d_model = d_model
+        self.hidden_size = d_model
+
+        self.n_layer = n_layer
+        self.num_hidden_layers = n_layer
+
+        self.n_head = n_head
+        self.num_attention_heads = n_head
+        self.head_dim = head_dim
+
+        self.ffn_multiplier = ffn_multiplier
+        self.multiple_of = multiple_of
+
+        self.block_size = block_size
+        self.max_position_embeddings = block_size
+        self.rope_theta = rope_theta
+
+        self.dropout = dropout
+        self.rms_norm_eps = rms_norm_eps
+        self.initializer_range = initializer_range
+        self.attention_bias = attention_bias
+        self.mlp_bias = mlp_bias
+
+        self.input_mode = input_mode
+        self.binary_dim = binary_dim
+        self.binary_encoding = binary_encoding
+        self.binary_scale = binary_scale
+        self.binary_repeat = d_model // binary_dim
+
+        self.code_seed = code_seed
+        self.min_row_weight = min_row_weight
+        self.min_col_weight = min_col_weight
+
+        self.use_cache = use_cache
+        self.is_decoder = True
+        self.is_encoder_decoder = False

generation_config.json

@@ -0,0 +1,8 @@
+{
+  "bos_token_id": 0,
+  "do_sample": false,
+  "eos_token_id": 0,
+  "pad_token_id": null,
+  "transformers_version": null,
+  "use_cache": false
+}

model.safetensors

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

modeling_attn_ext.py

@@ -0,0 +1,733 @@
+import math
+from typing import Optional
+
+import torch
+import torch.utils.checkpoint
+import torch.nn as nn
+import torch.nn.functional as F
+
+from transformers import PreTrainedModel
+from transformers.generation import GenerationMixin
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+from .configuration_attn_ext import AttnExtConfig
+
+
+def round_up(value: int, multiple: int) -> int:
+    return multiple * math.ceil(value / multiple)
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.eps = eps
+
+    def forward(self, x):
+        dtype = x.dtype
+        xf = x.float()
+        xf = xf * torch.rsqrt(
+            xf.pow(2).mean(dim=-1, keepdim=True) + self.eps
+        )
+        return (xf * self.weight.float()).to(dtype)
+
+
+def rotate_half(x):
+    x1 = x[..., ::2]
+    x2 = x[..., 1::2]
+    return torch.stack((-x2, x1), dim=-1).flatten(-2)
+
+
+class RotaryEmbedding(nn.Module):
+    def __init__(self, dim, max_position, theta):
+        super().__init__()
+
+        inv_freq = 1.0 / (
+            theta
+            ** (
+                torch.arange(0, dim, 2, dtype=torch.float32)
+                / dim
+            )
+        )
+
+        positions = torch.arange(
+            max_position,
+            dtype=torch.float32,
+        )
+
+        frequencies = torch.outer(positions, inv_freq)
+        embedding = torch.repeat_interleave(
+            frequencies,
+            repeats=2,
+            dim=-1,
+        )
+
+        self.register_buffer(
+            "cos_cached",
+            embedding.cos(),
+            persistent=False,
+        )
+        self.register_buffer(
+            "sin_cached",
+            embedding.sin(),
+            persistent=False,
+        )
+
+    def forward(self, q, k, position_ids=None):
+        sequence_length = q.shape[-2]
+
+        if position_ids is None:
+            cos = self.cos_cached[:sequence_length][
+                None, None, :, :
+            ]
+            sin = self.sin_cached[:sequence_length][
+                None, None, :, :
+            ]
+        else:
+            cos = self.cos_cached[position_ids][:, None, :, :]
+            sin = self.sin_cached[position_ids][:, None, :, :]
+
+        cos = cos.to(device=q.device, dtype=q.dtype)
+        sin = sin.to(device=q.device, dtype=q.dtype)
+
+        q = q * cos + rotate_half(q) * sin
+        k = k * cos + rotate_half(k) * sin
+        return q, k
+
+
+class CausalSelfAttention(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+
+        self.d_model = config.d_model
+        self.n_head = config.n_head
+        self.head_dim = config.head_dim
+        self.dropout_p = config.dropout
+
+        self.q_proj = nn.Linear(
+            config.d_model,
+            config.d_model,
+            bias=config.attention_bias,
+        )
+        self.k_proj = nn.Linear(
+            config.d_model,
+            config.d_model,
+            bias=config.attention_bias,
+        )
+        self.v_proj = nn.Linear(
+            config.d_model,
+            config.d_model,
+            bias=config.attention_bias,
+        )
+        self.o_proj = nn.Linear(
+            config.d_model,
+            config.d_model,
+            bias=config.attention_bias,
+        )
+
+        self.rope = RotaryEmbedding(
+            config.head_dim,
+            config.block_size,
+            config.rope_theta,
+        )
+
+    def forward(
+        self,
+        x,
+        attention_mask=None,
+        position_ids=None,
+    ):
+        batch_size, sequence_length, channels = x.shape
+
+        q = self.q_proj(x).view(
+            batch_size,
+            sequence_length,
+            self.n_head,
+            self.head_dim,
+        ).transpose(1, 2)
+
+        k = self.k_proj(x).view(
+            batch_size,
+            sequence_length,
+            self.n_head,
+            self.head_dim,
+        ).transpose(1, 2)
+
+        v = self.v_proj(x).view(
+            batch_size,
+            sequence_length,
+            self.n_head,
+            self.head_dim,
+        ).transpose(1, 2)
+
+        q, k = self.rope(
+            q,
+            k,
+            position_ids=position_ids,
+        )
+
+        dropout_p = self.dropout_p if self.training else 0.0
+
+        if attention_mask is None or bool(attention_mask.all()):
+            output = F.scaled_dot_product_attention(
+                q,
+                k,
+                v,
+                attn_mask=None,
+                dropout_p=dropout_p,
+                is_causal=True,
+            )
+        else:
+            if attention_mask.shape != (
+                batch_size,
+                sequence_length,
+            ):
+                raise ValueError(
+                    "attention_mask must have shape "
+                    f"{(batch_size, sequence_length)}"
+                )
+
+            causal = torch.ones(
+                sequence_length,
+                sequence_length,
+                device=x.device,
+                dtype=torch.bool,
+            ).tril()
+
+            allowed = (
+                causal[None, None, :, :]
+                & attention_mask[:, None, None, :].bool()
+            )
+
+            output = F.scaled_dot_product_attention(
+                q,
+                k,
+                v,
+                attn_mask=allowed,
+                dropout_p=dropout_p,
+                is_causal=False,
+            )
+
+        output = output.transpose(1, 2).contiguous().view(
+            batch_size,
+            sequence_length,
+            channels,
+        )
+
+        return self.o_proj(output)
+
+
+class SwiGLU(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+
+        hidden_dim = round_up(
+            int(config.ffn_multiplier * config.d_model),
+            config.multiple_of,
+        )
+
+        self.gate_proj = nn.Linear(
+            config.d_model,
+            hidden_dim,
+            bias=config.mlp_bias,
+        )
+        self.up_proj = nn.Linear(
+            config.d_model,
+            hidden_dim,
+            bias=config.mlp_bias,
+        )
+        self.down_proj = nn.Linear(
+            hidden_dim,
+            config.d_model,
+            bias=config.mlp_bias,
+        )
+        self.dropout = nn.Dropout(config.dropout)
+
+    def forward(self, x):
+        x = F.silu(self.gate_proj(x)) * self.up_proj(x)
+        return self.dropout(self.down_proj(x))
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+
+        self.input_norm = RMSNorm(
+            config.d_model,
+            config.rms_norm_eps,
+        )
+        self.post_attention_norm = RMSNorm(
+            config.d_model,
+            config.rms_norm_eps,
+        )
+
+        self.attention = CausalSelfAttention(config)
+        self.mlp = SwiGLU(config)
+
+    def forward(
+        self,
+        x,
+        attention_mask=None,
+        position_ids=None,
+    ):
+        x = x + self.attention(
+            self.input_norm(x),
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+        )
+
+        x = x + self.mlp(
+            self.post_attention_norm(x)
+        )
+
+        return x
+
+
+def canonical_binary_codebook(
+    vocab_size,
+    bits,
+    encoding,
+):
+    token_ids = torch.arange(
+        vocab_size,
+        dtype=torch.int64,
+    )
+    shifts = torch.arange(
+        bits,
+        dtype=torch.int64,
+    )
+
+    codebook = (
+        (token_ids[:, None] >> shifts[None, :]) & 1
+    ).to(torch.float32)
+
+    if encoding == "bipolar":
+        codebook = codebook.mul(2.0).sub(1.0)
+
+    return codebook.contiguous()
+
+
+def gf2_rank(matrix):
+    matrix = matrix.detach().cpu().to(
+        torch.uint8
+    ).clone()
+    matrix &= 1
+
+    rows, columns = matrix.shape
+    rank = 0
+
+    for column in range(columns):
+        pivot = None
+
+        for row in range(rank, rows):
+            if int(matrix[row, column]) == 1:
+                pivot = row
+                break
+
+        if pivot is None:
+            continue
+
+        if pivot != rank:
+            temporary = matrix[rank].clone()
+            matrix[rank] = matrix[pivot]
+            matrix[pivot] = temporary
+
+        for row in range(rows):
+            if row != rank and int(
+                matrix[row, column]
+            ) == 1:
+                matrix[row] ^= matrix[rank]
+
+        rank += 1
+
+        if rank == rows:
+            break
+
+    return rank
+
+
+def make_invertible_gf2_matrix(
+    bits,
+    seed,
+    min_row_weight,
+    min_col_weight,
+):
+    generator = torch.Generator(device="cpu")
+    generator.manual_seed(seed)
+
+    for _ in range(1_000_000):
+        matrix = torch.randint(
+            0,
+            2,
+            (bits, bits),
+            generator=generator,
+            dtype=torch.uint8,
+        )
+
+        if bool(
+            torch.any(
+                matrix.sum(dim=1) < min_row_weight
+            )
+        ):
+            continue
+
+        if bool(
+            torch.any(
+                matrix.sum(dim=0) < min_col_weight
+            )
+        ):
+            continue
+
+        if gf2_rank(matrix) == bits:
+            return matrix.contiguous()
+
+    raise RuntimeError(
+        "Could not construct an invertible GF(2) matrix"
+    )
+
+
+def gf2_binary_codebook(config):
+    source = canonical_binary_codebook(
+        config.vocab_size,
+        config.binary_dim,
+        "zero_one",
+    ).to(torch.uint8)
+
+    matrix = make_invertible_gf2_matrix(
+        bits=config.binary_dim,
+        seed=config.code_seed,
+        min_row_weight=config.min_row_weight,
+        min_col_weight=config.min_col_weight,
+    )
+
+    shift = torch.zeros(
+        config.binary_dim,
+        dtype=torch.uint8,
+    )
+
+    codebook = (
+        source.to(torch.int16)
+        @ matrix.to(torch.int16).T
+    ).remainder(2).to(torch.uint8)
+
+    codebook = codebook ^ shift
+
+    if config.binary_encoding == "bipolar":
+        codebook = (
+            codebook.float().mul(2.0).sub(1.0)
+        )
+    else:
+        codebook = codebook.float()
+
+    return (
+        codebook.contiguous(),
+        matrix.contiguous(),
+        shift.contiguous(),
+    )
+
+
+class FixedBinaryEmbedding(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+
+        if config.input_mode == "binary16":
+            codebook = canonical_binary_codebook(
+                config.vocab_size,
+                config.binary_dim,
+                config.binary_encoding,
+            )
+            matrix = None
+            shift = None
+
+        elif config.input_mode == "gf2":
+            codebook, matrix, shift = (
+                gf2_binary_codebook(config)
+            )
+
+        else:
+            raise ValueError(
+                "FixedBinaryEmbedding requires a "
+                "frozen-code input mode"
+            )
+
+        self.register_buffer(
+            "codebook",
+            codebook,
+            persistent=True,
+        )
+
+        if matrix is not None:
+            self.register_buffer(
+                "A_gf2",
+                matrix,
+                persistent=True,
+            )
+            self.register_buffer(
+                "b_gf2",
+                shift,
+                persistent=True,
+            )
+
+        self.repeat = config.binary_repeat
+        self.binary_scale = config.binary_scale
+
+    @property
+    def weight(self):
+        return self.codebook
+
+    def forward(self, input_ids):
+        code = self.codebook[input_ids.long()]
+
+        output = code.repeat(
+            *([1] * (code.ndim - 1)),
+            self.repeat,
+        )
+
+        if self.binary_scale != 1.0:
+            output = output * self.binary_scale
+
+        return output
+
+
+class AttnExtPreTrainedModel(PreTrainedModel):
+    config_class = AttnExtConfig
+    base_model_prefix = "attn_ext"
+    supports_gradient_checkpointing = True
+    _supports_sdpa = True
+    _no_split_modules = ["TransformerBlock"]
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            nn.init.normal_(
+                module.weight,
+                mean=0.0,
+                std=self.config.initializer_range,
+            )
+            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=self.config.initializer_range,
+            )
+
+
+class AttnExtForCausalLM(
+    AttnExtPreTrainedModel,
+    GenerationMixin,
+):
+    main_input_name = "input_ids"
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        if config.input_mode == "learned":
+            self.token_embeddings = nn.Embedding(
+                config.vocab_size,
+                config.d_model,
+            )
+        else:
+            self.token_embeddings = FixedBinaryEmbedding(
+                config
+            )
+
+        self.layers = nn.ModuleList(
+            [
+                TransformerBlock(config)
+                for _ in range(config.n_layer)
+            ]
+        )
+
+        self.final_norm = RMSNorm(
+            config.d_model,
+            config.rms_norm_eps,
+        )
+
+        self.lm_head = nn.Linear(
+            config.d_model,
+            config.vocab_size,
+            bias=False,
+        )
+
+        self.gradient_checkpointing = False
+        self.post_init()
+
+        residual_std = (
+            config.initializer_range
+            / math.sqrt(2 * config.n_layer)
+        )
+
+        for layer in self.layers:
+            nn.init.normal_(
+                layer.attention.o_proj.weight,
+                mean=0.0,
+                std=residual_std,
+            )
+            nn.init.normal_(
+                layer.mlp.down_proj.weight,
+                mean=0.0,
+                std=residual_std,
+            )
+
+    def get_input_embeddings(self):
+        return self.token_embeddings
+
+    def set_input_embeddings(self, value):
+        if self.config.input_mode != "learned":
+            raise RuntimeError(
+                "Frozen input codes cannot be replaced "
+                "through set_input_embeddings"
+            )
+        self.token_embeddings = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, value):
+        self.lm_head = value
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids,
+        attention_mask=None,
+        **kwargs,
+    ):
+        if input_ids.shape[1] > self.config.block_size:
+            input_ids = input_ids[
+                :, -self.config.block_size:
+            ]
+
+            if attention_mask is not None:
+                attention_mask = attention_mask[
+                    :, -self.config.block_size:
+                ]
+
+        position_ids = None
+
+        if attention_mask is not None:
+            position_ids = (
+                attention_mask.long().cumsum(-1) - 1
+            )
+            position_ids.masked_fill_(
+                attention_mask == 0,
+                0,
+            )
+
+        return {
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "position_ids": position_ids,
+            "use_cache": False,
+        }
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        labels=None,
+        position_ids=None,
+        inputs_embeds=None,
+        use_cache=None,
+        return_dict=None,
+        **kwargs,
+    ):
+        if input_ids is None and inputs_embeds is None:
+            raise ValueError(
+                "input_ids or inputs_embeds is required"
+            )
+
+        if inputs_embeds is not None:
+            x = inputs_embeds
+            batch_size, sequence_length, _ = x.shape
+        else:
+            batch_size, sequence_length = input_ids.shape
+            x = self.token_embeddings(input_ids)
+
+        if sequence_length > self.config.block_size:
+            raise ValueError(
+                f"Sequence length {sequence_length} exceeds "
+                f"block_size={self.config.block_size}"
+            )
+
+        if attention_mask is not None:
+            expected = (batch_size, sequence_length)
+            if attention_mask.shape != expected:
+                raise ValueError(
+                    f"attention_mask must have shape {expected}"
+                )
+
+        # HF_EXPORT_INPUT_DTYPE_FIX
+        # Frozen floating-point buffers may remain FP32 after loading.
+        # Match the residual stream to the backbone parameter dtype.
+        x = x.to(dtype=self.layers[0].attention.q_proj.weight.dtype)
+
+        for layer in self.layers:
+            if self.gradient_checkpointing and self.training:
+                def custom_forward(hidden_states, current_layer=layer):
+                    return current_layer(
+                        hidden_states,
+                        attention_mask=attention_mask,
+                        position_ids=position_ids,
+                    )
+
+                x = torch.utils.checkpoint.checkpoint(
+                    custom_forward,
+                    x,
+                    use_reentrant=False,
+                )
+            else:
+                x = layer(
+                    x,
+                    attention_mask=attention_mask,
+                    position_ids=position_ids,
+                )
+
+        x = self.final_norm(x)
+        logits = self.lm_head(x)
+
+        loss = None
+
+        if labels is not None:
+            if labels.shape != (
+                batch_size,
+                sequence_length,
+            ):
+                raise ValueError(
+                    "labels must have the same shape as input_ids"
+                )
+
+            shift_logits = logits[:, :-1, :].contiguous()
+            shift_labels = labels[:, 1:].contiguous().clone()
+
+            if attention_mask is not None:
+                shift_labels.masked_fill_(
+                    attention_mask[:, 1:].eq(0),
+                    -100,
+                )
+
+            loss = F.cross_entropy(
+                shift_logits.float().view(
+                    -1,
+                    self.config.vocab_size,
+                ),
+                shift_labels.view(-1),
+                ignore_index=-100,
+            )
+
+        return_dict = (
+            self.config.use_return_dict
+            if return_dict is None
+            else return_dict
+        )
+
+        if not return_dict:
+            output = (logits,)
+            return ((loss,) + output) if loss is not None else output
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=None,
+        )

special_tokens_map.json

@@ -0,0 +1,42 @@
+{
+  "additional_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>"
+  ],
+  "bos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer_config.json

@@ -0,0 +1,168 @@
+{
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "<|im_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "3": {
+      "content": "<repo_name>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "4": {
+      "content": "<reponame>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "5": {
+      "content": "<file_sep>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "6": {
+      "content": "<filename>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "7": {
+      "content": "<gh_stars>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "8": {
+      "content": "<issue_start>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "9": {
+      "content": "<issue_comment>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "10": {
+      "content": "<issue_closed>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "11": {
+      "content": "<jupyter_start>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "12": {
+      "content": "<jupyter_text>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "13": {
+      "content": "<jupyter_code>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "14": {
+      "content": "<jupyter_output>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "15": {
+      "content": "<jupyter_script>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "16": {
+      "content": "<empty_output>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "additional_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>"
+  ],
+  "bos_token": "<|endoftext|>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|endoftext|>",
+  "extra_special_tokens": {},
+  "model_max_length": 8192,
+  "tokenizer_class": "GPT2Tokenizer",
+  "unk_token": "<|endoftext|>",
+  "vocab_size": 49152
+}