CMA-1M Mini release

.gitattributes

@@ -0,0 +1,35 @@
+*.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

README.md

@@ -0,0 +1,266 @@
+---
+library_name: transformers
+pipeline_tag: text-generation
+language:
+- en
+tags:
+- custom_code
+- causal-lm
+- cma
+- small-language-model
+- base-model
+- byte-level
+- safetensors
+datasets:
+- HuggingFaceFW/fineweb_edu_100BT-shuffled
+- mlfoundations/dclm-baseline-1.0-parquet
+- HuggingFaceTB/dclm-edu
+- HuggingFaceTB/smollm-corpus
+- HuggingFaceTB/finemath
+widget:
+- text: "The process of photosynthesis"
+  example_title: Science
+- text: "Once upon a time, in a quiet village"
+  example_title: Story
+- text: "A computer program is"
+  example_title: Technology
+---
+
+# CMA-1M Mini
+
+> [!WARNING]
+> **Experimental research model ? generations are fully unreliable.** CMA-1M Mini
+> exists only as a test bed for understanding how Channel-Mixing Attention works,
+> where it helps, and where its limits and failure modes appear. Do not treat its
+> output as factual, safe, coherent, or suitable for production or real-world
+> decisions.
+
+CMA-1M Mini is a **958,692-parameter base causal language model** built to test
+Channel-Mixing Attention (CMA) at very small scale. It combines causal grouped-query
+token attention with content-dependent mixing across each token's hidden channels.
+The model uses a lossless byte-level tokenizer, tied embeddings, native BF16 weights,
+and a 2,048-token context window.
+
+| | |
+|---|---|
+| Parameters | 958,692 |
+| Architecture | Decoder-only CMA causal LM |
+| Context | 2,048 byte tokens |
+| Vocabulary | 260 tokens: 256 bytes + PAD/BOS/EOS/UNK |
+| Weight format | BF16 Safetensors |
+| Intended interface | Plain-text completion |
+
+> This is a pretrained base model, not a chat or instruction model. Give it ordinary
+> text to continue; no chat template or role markers are required.
+
+## Quick start
+
+The architecture is provided as custom Transformers code, so
+`trust_remote_code=True` is required. PyTorch 2.5 or newer is recommended.
+
+```python
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+repo_id = "User01110/CMA-1M-Mini"
+device = "cuda" if torch.cuda.is_available() else "cpu"
+dtype = (
+    torch.bfloat16
+    if device == "cuda" and torch.cuda.is_bf16_supported()
+    else torch.float32
+)
+
+tokenizer = AutoTokenizer.from_pretrained(
+    repo_id,
+    trust_remote_code=True,
+)
+model = AutoModelForCausalLM.from_pretrained(
+    repo_id,
+    trust_remote_code=True,
+    dtype=dtype,
+).to(device).eval()
+
+prompt = "The process of photosynthesis"
+inputs = tokenizer(prompt, return_tensors="pt")
+inputs = {name: tensor.to(device) for name, tensor in inputs.items()}
+
+with torch.inference_mode():
+    output = model.generate(
+        **inputs,
+        max_new_tokens=96,
+        do_sample=False,
+    )
+
+print(tokenizer.decode(output[0], skip_special_tokens=True))
+```
+
+The tokenizer automatically prepends `<bos>` during normal encoding. It does not
+append `<eos>` to a prompt; generation ends when the model emits EOS or reaches the
+requested length. If you deliberately use `add_special_tokens=False`, prepend BOS
+yourself.
+
+## Generation options
+
+The included generation defaults are deterministic decoding with a repetition
+penalty of 1.2. Override them per request as needed.
+
+| Goal | Recommended settings |
+|---|---|
+| Reproducible completion | `do_sample=False` |
+| Balanced sampling | `do_sample=True, temperature=0.8, top_p=0.9, top_k=50` |
+| More varied text | `do_sample=True, temperature=1.0, top_p=0.95` |
+| Reduce loops | `repetition_penalty=1.1` to `1.2` |
+| Beam search | `do_sample=False, num_beams=4` |
+| Output length | Set `max_new_tokens`; keep prompt + output within 2,048 tokens |
+
+Example with sampling:
+
+```python
+with torch.inference_mode():
+    output = model.generate(
+        **inputs,
+        max_new_tokens=128,
+        do_sample=True,
+        temperature=0.8,
+        top_p=0.9,
+        top_k=50,
+        repetition_penalty=1.15,
+    )
+```
+
+For the high-level pipeline API:
+
+```python
+import torch
+from transformers import pipeline
+
+generator = pipeline(
+    "text-generation",
+    model="User01110/CMA-1M-Mini",
+    tokenizer="User01110/CMA-1M-Mini",
+    trust_remote_code=True,
+    dtype="auto",
+    device=0 if torch.cuda.is_available() else -1,
+)
+
+result = generator(
+    "In a distant future,",
+    max_new_tokens=80,
+    do_sample=True,
+    temperature=0.8,
+    top_p=0.9,
+)
+print(result[0]["generated_text"])
+```
+
+To score text rather than generate it:
+
+```python
+encoded = tokenizer("CMA models text one byte at a time.", return_tensors="pt")
+encoded = {name: tensor.to(device) for name, tensor in encoded.items()}
+
+with torch.inference_mode():
+    result = model(**encoded, labels=encoded["input_ids"])
+
+print(float(result.loss))
+```
+
+## Tokenizer and context
+
+- Every UTF-8 byte has a token, so ordinary text cannot become out-of-vocabulary.
+- The four control tokens are `<pad>` (0), `<bos>` (1), `<eos>` (2), and `<unk>` (3).
+- The context limit is 2,048 **byte tokens**, not 2,048 words or subword tokens.
+- Non-ASCII text usually consumes multiple byte tokens per character.
+- For long inputs, explicitly keep the most recent 2,048 tokens rather than relying
+  on implicit truncation.
+- The tokenizer has no arithmetic-specific splitting, chat template, or hidden prompt
+  transformation.
+
+## Architecture
+
+| Component | Configuration |
+|---|---|
+| Hidden width / layers | 128 / 6 |
+| Token attention | 4 query heads, 2 key-value heads |
+| Position encoding | Contiguous-half rotary embeddings, no scaling |
+| CMA layout | 8 channel slots x 16 channels |
+| CMA routing | 2 heads, expansion 2, content-dependent softmax mixing |
+| CMA initialization | 90% diagonal routing prior with a dense base path |
+| Feed-forward gate | SiLU-gated routed values |
+| Normalization | RMSNorm |
+| Embeddings | Input and output weights tied |
+| Attention runtime | Native PyTorch scaled-dot-product attention |
+
+For each token, CMA projects dense values, reshapes them into channel slots, and
+learns a softmax mixing matrix between those slots. A bounded learned coefficient
+controls the routed difference from the dense base value, so routing enriches rather
+than replaces the fallback path.
+
+The exported generation implementation does not maintain a KV cache. This keeps the
+custom model compact and straightforward, but long autoregressive generations will
+recompute the active context and are slower than cached generation.
+
+## Training data
+
+The model was pretrained as a general causal language model on the following mixture.
+Percentages describe the trained-token mixture.
+
+| Source | Share |
+|---|---:|
+| [FineWeb-Edu 100BT](https://huggingface.co/datasets/HuggingFaceFW/fineweb_edu_100BT-shuffled) | 45% |
+| [DCLM-Baseline 1.0](https://huggingface.co/datasets/mlfoundations/dclm-baseline-1.0-parquet) | 25% |
+| [DCLM-Edu](https://huggingface.co/datasets/HuggingFaceTB/dclm-edu) | 10% |
+| [Cosmopedia v2](https://huggingface.co/datasets/HuggingFaceTB/smollm-corpus) | 10% |
+| [FineMath 4+](https://huggingface.co/datasets/HuggingFaceTB/finemath) | 10% |
+
+No benchmark-specific prompts, task detectors, arithmetic vocabulary, or
+inference-time answer rules are built into the model.
+
+## Evaluation
+
+Evaluation is zero-shot. The four lm-eval tasks use normalized accuracy when
+provided by lm-eval 0.4.12. ArithMark-2 uses raw continuation log-likelihood sums.
+Weights are evaluated in BF16 with FP32 likelihood softmax and an automatic BOS
+prefix.
+
+| Benchmark | Accuracy |
+|---|---:|
+| HellaSwag | 29.35% |
+| ARC-Easy | 29.29% |
+| ARC-Challenge | 21.76% |
+| PIQA | 54.62% |
+| ArithMark-2 | 27.44% |
+| **Open SLM Leaderboard-style average** | **34.23%** |
+
+The combined score is
+`(HellaSwag + mean(ARC-Easy, ARC-Challenge) + PIQA + ArithMark-2) / 4`.
+This is a report-only reproduction of the leaderboard formula, not an official
+leaderboard submission. WikiText-103 normalized validation BPB is **1.6974**.
+Exact machine-readable results are available in
+[`benchmark_results.json`](./benchmark_results.json).
+
+## Intended use and limitations
+
+CMA-1M Mini is intended for architecture research, educational experiments,
+lightweight language-model tooling, and controlled comparisons at tiny scale.
+
+- At fewer than one million parameters, generations are short-range and frequently
+  incoherent; the model should not be treated as a knowledge source.
+- It is not instruction-tuned, conversationally aligned, tool-using, or safety-tuned.
+- Training data is predominantly English even though byte tokenization can represent
+  any UTF-8 text.
+- Outputs may reproduce biases, inaccuracies, or undesirable patterns from public
+  training corpora.
+- Do not use it for consequential medical, legal, financial, or safety decisions.
+- Loading custom code executes files from the repository. Review the code or pin a
+  trusted revision in security-sensitive environments.
+
+## Repository contents
+
+- `model.safetensors` — BF16 model weights
+- `modeling_cma.py` — Transformers-compatible CMA implementation
+- `config.json` and `generation_config.json` — architecture and decoding defaults
+- `tokenizer.json` and `tokenizer_config.json` — deterministic byte tokenizer
+- `benchmark_results.json` — exact evaluation metrics and protocol metadata
+- `training_config.json` — reproducibility configuration
+

benchmark_results.json

@@ -0,0 +1,52 @@
+{
+  "step": 12000,
+  "parameters": 958692,
+  "lm_eval_version": "0.4.12",
+  "evaluation_dtype": "bfloat16",
+  "softmax_dtype": "float32",
+  "evaluation_autocast": false,
+  "num_fewshot": 0,
+  "bos_prefix": true,
+  "leaderboard_accuracy": {
+    "arc_easy": 0.29292929292929293,
+    "arc_challenge": 0.2175767918088737,
+    "hellaswag": 0.29346743676558457,
+    "piqa": 0.5462459194776932
+  },
+  "lm_eval_metric_policy": "acc_norm when available, otherwise acc",
+  "selection_metric": "validation.normalized_bpb (lower is better)",
+  "arithmark_2.0": {
+    "acc": 0.2744,
+    "correct": 686,
+    "total": 2500,
+    "by_operator_count": {
+      "1": {
+        "acc": 0.2552,
+        "correct": 319,
+        "total": 1250
+      },
+      "2": {
+        "acc": 0.30666666666666664,
+        "correct": 230,
+        "total": 750
+      },
+      "3": {
+        "acc": 0.274,
+        "correct": 137,
+        "total": 500
+      }
+    }
+  },
+  "arc_average": 0.2552530423690833,
+  "open_slm_leaderboard_average": 0.34234159965309024,
+  "average_formula": "(hellaswag + mean(arc_easy, arc_challenge) + piqa + arithmark_2.0) / 4",
+  "validation": {
+    "loss": 1.1765458586014317,
+    "perplexity": 3.2431525260156,
+    "normalized_bpb": 1.6973968755827331,
+    "tokens": 1144831,
+    "normalized_utf8_bytes": 1144831,
+    "window": 2048,
+    "stride": 1024
+  }
+}

config.json

@@ -0,0 +1,31 @@
+{
+  "architectures": [
+    "CMAForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "modeling_cma.CMAConfig",
+    "AutoModelForCausalLM": "modeling_cma.CMAForCausalLM"
+  },
+  "model_type": "cma",
+  "vocab_size": 260,
+  "seq_len": 2048,
+  "max_position_embeddings": 2048,
+  "n_positions": 2048,
+  "n_ctx": 2048,
+  "d_model": 128,
+  "n_layers": 6,
+  "n_heads": 4,
+  "n_kv_heads": 2,
+  "chunk": 16,
+  "cma_heads": 2,
+  "expand": 2,
+  "cma_identity_prob": 0.9,
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "pad_token_id": 0,
+  "tie_word_embeddings": true,
+  "safe_serialization": true,
+  "minimum_torch_version": "2.5",
+  "torch_dtype": "bfloat16",
+  "transformers_version": "5.13.0"
+}

generation_config.json

@@ -0,0 +1,7 @@
+{
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "pad_token_id": 0,
+  "do_sample": false,
+  "repetition_penalty": 1.2
+}

model.safetensors

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

modeling_cma.py

@@ -0,0 +1,375 @@
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+# CMA is text-only. Disable Transformers' optional vision backend before it
+# imports PreTrainedModel: some cloud images expose a system torchvision built
+# against a different PyTorch, which otherwise crashes on torchvision::nms.
+import transformers as _transformers
+import transformers.utils as _transformers_utils
+from transformers.utils import import_utils as _transformers_import_utils
+
+
+def _cma_torchvision_unavailable():
+    return False
+
+
+_transformers_import_utils._torchvision_available = False
+_transformers_import_utils.is_torchvision_available = _cma_torchvision_unavailable
+_transformers_import_utils.is_torchvision_v2_available = _cma_torchvision_unavailable
+_transformers_utils.is_torchvision_available = _cma_torchvision_unavailable
+_transformers_utils.is_torchvision_v2_available = _cma_torchvision_unavailable
+_transformers.is_torchvision_available = _cma_torchvision_unavailable
+_transformers.is_torchvision_v2_available = _cma_torchvision_unavailable
+try:
+    from transformers import GenerationMixin
+except ImportError:
+    from transformers.generation import GenerationMixin
+from transformers import PretrainedConfig, PreTrainedModel
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+
+class CMAConfig(PretrainedConfig):
+    model_type = "cma"
+
+    def __init__(
+        self,
+        vocab_size=260,
+        seq_len=2048,
+        d_model=128,
+        n_layers=6,
+        n_heads=4,
+        n_kv_heads=2,
+        chunk=16,
+        cma_heads=2,
+        expand=2,
+        cma_identity_prob=0.90,
+        max_position_embeddings=None,
+        n_positions=None,
+        n_ctx=None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.vocab_size = vocab_size
+        self.seq_len = seq_len
+        self.max_position_embeddings = max_position_embeddings or seq_len
+        self.n_positions = n_positions or self.max_position_embeddings
+        self.n_ctx = n_ctx or self.max_position_embeddings
+        self.d_model = d_model
+        self.n_layers = n_layers
+        self.n_heads = n_heads
+        self.n_kv_heads = n_kv_heads
+        self.chunk = chunk
+        self.cma_heads = cma_heads
+        self.expand = expand
+        self.cma_identity_prob = cma_identity_prob
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, d):
+        super().__init__()
+        self.w = nn.Parameter(torch.ones(d))
+
+    def forward(self, x):
+        return F.rms_norm(
+            x, (x.shape[-1],), self.w.to(dtype=x.dtype), eps=1e-6
+        )
+
+
+def rope_cache(seq_len, head_dim, device, base=10000.0):
+    inv = 1.0 / (
+        base ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim)
+    )
+    t = torch.arange(seq_len, device=device).float()
+    freqs = torch.outer(t, inv)
+    return torch.cos(freqs), torch.sin(freqs)
+
+
+def apply_rope(x, cos, sin):
+    x1, x2 = x.chunk(2, dim=-1)
+    return torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1)
+
+
+class TokenAttention(nn.Module):
+    def __init__(self, d, n_heads, n_kv_heads):
+        super().__init__()
+        if d % n_heads or n_heads % n_kv_heads:
+            raise ValueError("d_model % n_heads and n_heads % n_kv_heads must be zero.")
+        self.h, self.kv_h, self.hd = n_heads, n_kv_heads, d // n_heads
+        self.q = nn.Linear(d, d, bias=False)
+        self.k = nn.Linear(d, n_kv_heads * self.hd, bias=False)
+        self.v = nn.Linear(d, n_kv_heads * self.hd, bias=False)
+        self.o = nn.Linear(d, d, bias=False)
+        self.qn, self.kn = RMSNorm(self.hd), RMSNorm(self.hd)
+
+    def forward(self, x, cos, sin):
+        B, T, d = x.shape
+        q = self.q(x).view(B, T, self.h, self.hd).transpose(1, 2)
+        k = self.k(x).view(B, T, self.kv_h, self.hd).transpose(1, 2)
+        v = self.v(x).view(B, T, self.kv_h, self.hd).transpose(1, 2)
+        q, k = self.qn(q), self.kn(k)
+        q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin)
+        # Native GQA keeps compact K/V heads and lets SDPA select a fused CUDA
+        # implementation without materializing repeated K/V activations.
+        y = F.scaled_dot_product_attention(
+            q,
+            k,
+            v,
+            is_causal=True,
+            enable_gqa=self.h != self.kv_h,
+        )
+        return self.o(y.transpose(1, 2).reshape(B, T, d))
+
+
+class CMA(nn.Module):
+    def __init__(self, d, chunk=16, heads=2, expand=2, identity_prob=0.90):
+        super().__init__()
+        if min(d, chunk, heads, expand) <= 0:
+            raise ValueError("CMA dimensions and expansion must be positive.")
+        if d % chunk or chunk % heads:
+            raise ValueError("CMA requires d % chunk == 0 and chunk % heads == 0.")
+        if d // chunk < 2:
+            raise ValueError("CMA requires at least two channel slots.")
+        if not 0.0 < identity_prob < 1.0:
+            raise ValueError("cma_identity_prob must be between zero and one.")
+
+        self.d, self.n, self.c, self.h = d, d // chunk, chunk, heads
+        self.hd, self.expand = chunk // heads, expand
+        self.chunk_emb = nn.Parameter(torch.randn(self.n, chunk) * 0.02)
+        self.wqk = nn.Parameter(torch.randn(self.n, chunk, 2 * chunk) * 0.02)
+        self.global_proj = nn.Linear(d, chunk, bias=False)
+        self.wv = nn.Linear(d, d * expand, bias=False)
+        self.bias = nn.Parameter(torch.zeros(heads, self.n, self.n))
+        self.qn, self.kn = RMSNorm(self.hd), RMSNorm(self.hd)
+        self.logit_scale = nn.Parameter(torch.zeros(heads))
+        self.layer_gain = nn.Parameter(torch.zeros(heads))
+        self.route_gate_weight = nn.Parameter(torch.randn(heads, self.hd) * 0.02)
+        self.route_gate_bias = nn.Parameter(torch.zeros(heads))
+        self.gate = nn.Linear(d, d * expand, bias=False)
+        self.o = nn.Linear(d * expand, d, bias=False)
+        nn.init.zeros_(self.o.weight)
+
+        diagonal_bias = math.log(
+            (self.n - 1) * identity_prob / (1.0 - identity_prob)
+        )
+        with torch.no_grad():
+            self.bias.add_(torch.eye(self.n) * diagonal_bias)
+
+    def _routing(self, x):
+        batch_tokens = x.numel() // self.d
+        xc = x.reshape(batch_tokens, self.n, self.c)
+        global_state = self.global_proj(x).reshape(batch_tokens, 1, self.c)
+        q_input = xc + self.chunk_emb + global_state
+        value_slots = self.wv(x).reshape(
+            batch_tokens, self.n, self.h, self.hd, self.expand
+        )
+        key_input = value_slots.mean(dim=-1).reshape(batch_tokens, self.n, self.c)
+        key_input = key_input + self.chunk_emb
+
+        q = torch.einsum("bnc,nco->bno", q_input, self.wqk[..., : self.c])
+        k = torch.einsum("bnc,nco->bno", key_input, self.wqk[..., self.c :])
+        v = value_slots.reshape(
+            batch_tokens, self.n, self.h, self.hd * self.expand
+        ).transpose(1, 2)
+        q = self.qn(q.reshape(batch_tokens, self.n, self.h, self.hd)).transpose(1, 2)
+        k = self.kn(k.reshape(batch_tokens, self.n, self.h, self.hd)).transpose(1, 2)
+        q = F.normalize(q.float(), dim=-1).to(v.dtype)
+        k = F.normalize(k.float(), dim=-1).to(v.dtype)
+
+        scale = self.logit_scale.clamp(max=math.log(100.0)).exp()
+        logits = (q * scale.view(1, self.h, 1, 1).to(q.dtype)) @ k.transpose(-2, -1)
+        logits = logits + self.bias.unsqueeze(0).to(logits.dtype)
+        attn = F.softmax(logits, dim=-1, dtype=torch.float32).to(v.dtype)
+        routed = attn @ v
+        route_signal = (
+            (q.float() * self.route_gate_weight.float().view(1, self.h, 1, self.hd)).sum(-1)
+            + self.route_gate_bias.float().view(1, self.h, 1)
+            + self.layer_gain.float().view(1, self.h, 1)
+        )
+        route_coeff = torch.tanh(route_signal).to(v.dtype)
+        contribution = route_coeff.unsqueeze(-1) * (routed - v)
+        return v, attn, route_coeff, contribution, logits
+
+    def forward(self, x):
+        B, T, _ = x.shape
+        v, _, _, contribution, _ = self._routing(x)
+        y = (v + contribution).transpose(1, 2).reshape(B, T, self.d * self.expand)
+        return self.o(y * F.silu(self.gate(x)))
+
+    @torch.no_grad()
+    def diagnostic_stats(self, x):
+        probe_count = x.shape[0] if x.ndim > 2 else 1
+        v, attn, route_coeff, contribution, logits = self._routing(x)
+        B, T, _ = x.shape
+        base = v.transpose(1, 2).reshape(B, T, self.d * self.expand)
+        routed = (v + contribution).transpose(1, 2).reshape(
+            B, T, self.d * self.expand
+        )
+        gate = F.silu(self.gate(x))
+        base_output = self.o(base * gate)
+        routed_output = self.o(routed * gate)
+        output_delta = (routed_output - base_output).reshape(-1, self.d)
+        token_effect = (
+            output_delta.float().norm(dim=-1)
+            / routed_output.reshape(-1, self.d).float().norm(dim=-1).clamp_min(1e-12)
+        )
+        probe_effect = token_effect.reshape(probe_count, -1).mean(dim=-1)
+        probs = attn.float().clamp_min(1e-9)
+        entropy = -(probs * probs.log()).sum(dim=-1) / math.log(self.n)
+        diagonal_mass = probs.diagonal(dim1=-2, dim2=-1).mean()
+        if self.h > 1:
+            flattened = probs.transpose(0, 1).reshape(self.h, -1)
+            normalized = F.normalize(flattened, dim=-1)
+            similarity = normalized @ normalized.mT
+            head_similarity = (
+                similarity.sum() - similarity.diagonal().sum()
+            ) / (self.h * (self.h - 1))
+        else:
+            head_similarity = probs.new_zeros(())
+        return {
+            "token_effect": token_effect,
+            "probe_effect": probe_effect,
+            "entropy": entropy.mean().item(),
+            "diagonal_mass": diagonal_mass.item(),
+            "head_similarity": head_similarity.item(),
+            "gate_abs": route_coeff.float().abs().mean().item(),
+            "gate_saturation": (route_coeff.float().abs() > 0.95).float().mean().item(),
+            "logit_max": logits.float().abs().max().item(),
+        }
+
+
+class Block(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        d = config.d_model
+        self.n1, self.n2 = RMSNorm(d), RMSNorm(d)
+        self.attn = TokenAttention(d, config.n_heads, config.n_kv_heads)
+        self.mix = CMA(
+            d,
+            config.chunk,
+            config.cma_heads,
+            config.expand,
+            config.cma_identity_prob,
+        )
+
+    def forward(self, x, cos, sin):
+        x = x + self.attn(self.n1(x), cos, sin)
+        x = x + self.mix(self.n2(x))
+        return x
+
+
+class CMAModel(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        d = config.d_model
+        self.config = config
+        self.emb = nn.Embedding(config.vocab_size, d)
+        self.blocks = nn.ModuleList(Block(config) for _ in range(config.n_layers))
+        self.norm = RMSNorm(d)
+        hd = d // config.n_heads
+        cos, sin = rope_cache(config.seq_len, hd, "cpu")
+        self.register_buffer("cos", cos)
+        self.register_buffer("sin", sin)
+
+    def forward_hidden(self, idx):
+        if idx.size(1) > self.config.seq_len:
+            idx = idx[:, -self.config.seq_len :]
+        x = self.emb(idx)
+        device_type = x.device.type
+        rope_dtype = (
+            torch.get_autocast_dtype(device_type)
+            if torch.is_autocast_enabled(device_type)
+            else x.dtype
+        )
+        cos = self.cos[: idx.size(1)].to(device=idx.device, dtype=rope_dtype)
+        sin = self.sin[: idx.size(1)].to(device=idx.device, dtype=rope_dtype)
+        for b in self.blocks:
+            x = b(x, cos, sin)
+        return self.norm(x)
+
+
+class CMAForCausalLM(PreTrainedModel, GenerationMixin):
+    config_class = CMAConfig
+    base_model_prefix = "model"
+    _tied_weights_keys = {"head.weight": "model.emb.weight"}
+    all_tied_weights_keys = {"head.weight": "model.emb.weight"}
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = CMAModel(config)
+        self.head = nn.Linear(config.d_model, config.vocab_size, bias=False)
+        # Modern Transformers creates loader metadata and performs configured
+        # tying in post_init(); omitting it leaves all_tied_weights_keys absent.
+        self.post_init()
+        self.head.weight = self.model.emb.weight
+
+    def get_input_embeddings(self):
+        return self.model.emb
+
+    def set_input_embeddings(self, value):
+        self.model.emb = value
+
+    def get_output_embeddings(self):
+        return self.head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.head = new_embeddings
+
+    def raw_logits(self, idx):
+        return self.head(self.model.forward_hidden(idx))
+
+    def logits(self, idx):
+        return self.raw_logits(idx)
+
+    def _masked_logits(self, input_ids, attention_mask):
+        if attention_mask is None or bool(attention_mask.all()):
+            return self.logits(input_ids)
+
+        B, T = input_ids.shape
+        out = None
+        for i in range(B):
+            keep = attention_mask[i].bool().nonzero(as_tuple=False).flatten()
+            if keep.numel() == 0:
+                keep = torch.tensor([T - 1], device=input_ids.device)
+            trimmed = input_ids[i, keep].unsqueeze(0)
+            logits_i = self.logits(trimmed)
+            if out is None:
+                out = logits_i.new_zeros(B, T, logits_i.size(-1))
+            out[i, keep, :] = logits_i[0, -keep.numel() :, :]
+        return out
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        labels=None,
+        use_cache=False,
+        past_key_values=None,
+        **kwargs,
+    ):
+        if input_ids.size(1) > self.config.seq_len:
+            input_ids = input_ids[:, -self.config.seq_len :]
+            if attention_mask is not None:
+                attention_mask = attention_mask[:, -self.config.seq_len :]
+            if labels is not None:
+                labels = labels[:, -self.config.seq_len :]
+        logits = self._masked_logits(input_ids, attention_mask)
+        loss = None
+        if labels is not None:
+            shift_logits = logits[:, :-1, :].contiguous()
+            shift_labels = labels[:, 1:].contiguous()
+            loss = F.cross_entropy(
+                shift_logits.view(-1, shift_logits.size(-1)).float(),
+                shift_labels.view(-1),
+                ignore_index=-100,
+            )
+        return CausalLMOutputWithPast(loss=loss, logits=logits, past_key_values=None)
+
+    def prepare_inputs_for_generation(self, input_ids, **kwargs):
+        attention_mask = kwargs.get("attention_mask")
+        result = {"input_ids": input_ids[:, -self.config.seq_len :]}
+        if attention_mask is not None:
+            result["attention_mask"] = attention_mask[:, -self.config.seq_len :]
+        return result

tokenizer.json

@@ -0,0 +1,377 @@
+{
+  "version": "1.0",
+  "truncation": null,
+  "padding": null,
+  "added_tokens": [
+    {
+      "id": 0,
+      "content": "<pad>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    {
+      "id": 1,
+      "content": "<bos>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    {
+      "id": 2,
+      "content": "<eos>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    {
+      "id": 3,
+      "content": "<unk>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    }
+  ],
+  "normalizer": null,
+  "pre_tokenizer": {
+    "type": "ByteLevel",
+    "add_prefix_space": false,
+    "trim_offsets": true,
+    "use_regex": true
+  },
+  "post_processor": {
+    "type": "TemplateProcessing",
+    "single": [
+      {
+        "SpecialToken": {
+          "id": "<bos>",
+          "type_id": 0
+        }
+      },
+      {
+        "Sequence": {
+          "id": "A",
+          "type_id": 0
+        }
+      }
+    ],
+    "pair": [
+      {
+        "SpecialToken": {
+          "id": "<bos>",
+          "type_id": 0
+        }
+      },
+      {
+        "Sequence": {
+          "id": "A",
+          "type_id": 0
+        }
+      },
+      {
+        "Sequence": {
+          "id": "B",
+          "type_id": 0
+        }
+      }
+    ],
+    "special_tokens": {
+      "<bos>": {
+        "id": "<bos>",
+        "ids": [
+          1
+        ],
+        "tokens": [
+          "<bos>"
+        ]
+      }
+    }
+  },
+  "decoder": {
+    "type": "ByteLevel",
+    "add_prefix_space": true,
+    "trim_offsets": true,
+    "use_regex": true
+  },
+  "model": {
+    "type": "BPE",
+    "dropout": null,
+    "unk_token": "<unk>",
+    "continuing_subword_prefix": null,
+    "end_of_word_suffix": null,
+    "fuse_unk": false,
+    "byte_fallback": false,
+    "ignore_merges": false,
+    "vocab": {
+      "<pad>": 0,
+      "<bos>": 1,
+      "<eos>": 2,
+      "<unk>": 3,
+      "!": 4,
+      "\"": 5,
+      "#": 6,
+      "$": 7,
+      "%": 8,
+      "&": 9,
+      "'": 10,
+      "(": 11,
+      ")": 12,
+      "*": 13,
+      "+": 14,
+      ",": 15,
+      "-": 16,
+      ".": 17,
+      "/": 18,
+      "0": 19,
+      "1": 20,
+      "2": 21,
+      "3": 22,
+      "4": 23,
+      "5": 24,
+      "6": 25,
+      "7": 26,
+      "8": 27,
+      "9": 28,
+      ":": 29,
+      ";": 30,
+      "<": 31,
+      "=": 32,
+      ">": 33,
+      "?": 34,
+      "@": 35,
+      "A": 36,
+      "B": 37,
+      "C": 38,
+      "D": 39,
+      "E": 40,
+      "F": 41,
+      "G": 42,
+      "H": 43,
+      "I": 44,
+      "J": 45,
+      "K": 46,
+      "L": 47,
+      "M": 48,
+      "N": 49,
+      "O": 50,
+      "P": 51,
+      "Q": 52,
+      "R": 53,
+      "S": 54,
+      "T": 55,
+      "U": 56,
+      "V": 57,
+      "W": 58,
+      "X": 59,
+      "Y": 60,
+      "Z": 61,
+      "[": 62,
+      "\\": 63,
+      "]": 64,
+      "^": 65,
+      "_": 66,
+      "`": 67,
+      "a": 68,
+      "b": 69,
+      "c": 70,
+      "d": 71,
+      "e": 72,
+      "f": 73,
+      "g": 74,
+      "h": 75,
+      "i": 76,
+      "j": 77,
+      "k": 78,
+      "l": 79,
+      "m": 80,
+      "n": 81,
+      "o": 82,
+      "p": 83,
+      "q": 84,
+      "r": 85,
+      "s": 86,
+      "t": 87,
+      "u": 88,
+      "v": 89,
+      "w": 90,
+      "x": 91,
+      "y": 92,
+      "z": 93,
+      "{": 94,
+      "|": 95,
+      "}": 96,
+      "~": 97,
+      "¡": 98,
+      "¢": 99,
+      "£": 100,
+      "¤": 101,
+      "¥": 102,
+      "¦": 103,
+      "§": 104,
+      "¨": 105,
+      "©": 106,
+      "ª": 107,
+      "«": 108,
+      "¬": 109,
+      "®": 110,
+      "¯": 111,
+      "°": 112,
+      "±": 113,
+      "²": 114,
+      "³": 115,
+      "´": 116,
+      "µ": 117,
+      "¶": 118,
+      "·": 119,
+      "¸": 120,
+      "¹": 121,
+      "º": 122,
+      "»": 123,
+      "¼": 124,
+      "½": 125,
+      "¾": 126,
+      "¿": 127,
+      "À": 128,
+      "Á": 129,
+      "Â": 130,
+      "Ã": 131,
+      "Ä": 132,
+      "Å": 133,
+      "Æ": 134,
+      "Ç": 135,
+      "È": 136,
+      "É": 137,
+      "Ê": 138,
+      "Ë": 139,
+      "Ì": 140,
+      "Í": 141,
+      "Î": 142,
+      "Ï": 143,
+      "Ð": 144,
+      "Ñ": 145,
+      "Ò": 146,
+      "Ó": 147,
+      "Ô": 148,
+      "Õ": 149,
+      "Ö": 150,
+      "×": 151,
+      "Ø": 152,
+      "Ù": 153,
+      "Ú": 154,
+      "Û": 155,
+      "Ü": 156,
+      "Ý": 157,
+      "Þ": 158,
+      "ß": 159,
+      "à": 160,
+      "á": 161,
+      "â": 162,
+      "ã": 163,
+      "ä": 164,
+      "å": 165,
+      "æ": 166,
+      "ç": 167,
+      "è": 168,
+      "é": 169,
+      "ê": 170,
+      "ë": 171,
+      "ì": 172,
+      "í": 173,
+      "î": 174,
+      "ï": 175,
+      "ð": 176,
+      "ñ": 177,
+      "ò": 178,
+      "ó": 179,
+      "ô": 180,
+      "õ": 181,
+      "ö": 182,
+      "÷": 183,
+      "ø": 184,
+      "ù": 185,
+      "ú": 186,
+      "û": 187,
+      "ü": 188,
+      "ý": 189,
+      "þ": 190,
+      "ÿ": 191,
+      "Ā": 192,
+      "ā": 193,
+      "Ă": 194,
+      "ă": 195,
+      "Ą": 196,
+      "ą": 197,
+      "Ć": 198,
+      "ć": 199,
+      "Ĉ": 200,
+      "ĉ": 201,
+      "Ċ": 202,
+      "ċ": 203,
+      "Č": 204,
+      "č": 205,
+      "Ď": 206,
+      "ď": 207,
+      "Đ": 208,
+      "đ": 209,
+      "Ē": 210,
+      "ē": 211,
+      "Ĕ": 212,
+      "ĕ": 213,
+      "Ė": 214,
+      "ė": 215,
+      "Ę": 216,
+      "ę": 217,
+      "Ě": 218,
+      "ě": 219,
+      "Ĝ": 220,
+      "ĝ": 221,
+      "Ğ": 222,
+      "ğ": 223,
+      "Ġ": 224,
+      "ġ": 225,
+      "Ģ": 226,
+      "ģ": 227,
+      "Ĥ": 228,
+      "ĥ": 229,
+      "Ħ": 230,
+      "ħ": 231,
+      "Ĩ": 232,
+      "ĩ": 233,
+      "Ī": 234,
+      "ī": 235,
+      "Ĭ": 236,
+      "ĭ": 237,
+      "Į": 238,
+      "į": 239,
+      "İ": 240,
+      "ı": 241,
+      "IJ": 242,
+      "ij": 243,
+      "Ĵ": 244,
+      "ĵ": 245,
+      "Ķ": 246,
+      "ķ": 247,
+      "ĸ": 248,
+      "Ĺ": 249,
+      "ĺ": 250,
+      "Ļ": 251,
+      "ļ": 252,
+      "Ľ": 253,
+      "ľ": 254,
+      "Ŀ": 255,
+      "ŀ": 256,
+      "Ł": 257,
+      "ł": 258,
+      "Ń": 259
+    },
+    "merges": []
+  }
+}

tokenizer_config.json

@@ -0,0 +1,10 @@
+{
+  "backend": "tokenizers",
+  "bos_token": "<bos>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<eos>",
+  "model_max_length": 2048,
+  "pad_token": "<pad>",
+  "tokenizer_class": "TokenizersBackend",
+  "unk_token": "<unk>"
+}

training_config.json

@@ -0,0 +1,54 @@
+{
+  "steps": 20000,
+  "batch_size": 512,
+  "grad_accum": 4,
+  "seq_len": 2048,
+  "eval_seq_len": 2048,
+  "d_model": 128,
+  "n_layers": 6,
+  "n_heads": 4,
+  "n_kv_heads": 2,
+  "chunk": 16,
+  "cma_heads": 2,
+  "expand": 2,
+  "cma_identity_prob": 0.9,
+  "lr": 0.001,
+  "muon_momentum": 0.95,
+  "muon_ns_steps": 5,
+  "muon_adjust_lr_fn": "match_rms_adamw",
+  "warmup": 500,
+  "wd": 0.01,
+  "grad_clip": 5.0,
+  "log_every": 5,
+  "early_eval_step": 1000,
+  "diag_every": 1000,
+  "eval_every": 2000,
+  "val_batch_size": 32,
+  "val_batches": 0,
+  "val_stride": 1024,
+  "infer_tokens": 512,
+  "infer_repeat_penalty": 1.2,
+  "infer_prompt": "The process of photosynthesis",
+  "lm_eval_tasks": "arc_easy,arc_challenge,hellaswag,piqa",
+  "lm_eval_batch_size": "auto",
+  "lm_eval_device": "cuda",
+  "lm_eval_dtype": "bfloat16",
+  "lm_eval_softmax_dtype": "float32",
+  "lm_eval_expected_version": "0.4.12",
+  "lm_eval_retries": 3,
+  "lm_eval_export_dir": "CMA_1M_Mini_lm_eval_hf",
+  "lm_eval_output_dir": "lm_eval_results_CMA_1M_Mini",
+  "arithmark_batch_size": 16,
+  "arithmark_data_path": "benchmark_cache/arithmark_2.0.jsonl",
+  "arithmark_force_download": false,
+  "recipe_version": "CMA_1M_mini_byte_c16_h2_e2_muon_w500_cosine20k_b512_ga4_ctx2k_async_v7",
+  "hf_repo_id": "User01110/CMA-1M-Mini",
+  "hf_repo_private": false,
+  "hub_upload_retries": 3,
+  "tokenizer_name": "local-byte-level",
+  "tokenizer_revision": "byte-v1",
+  "data_seed": 1337,
+  "shuffle_buffer": 50000,
+  "tokenize_batch_size": 64,
+  "compile": true
+}