Upload 13 files

README.md

@@ -1,3 +1,226 @@
 ---
 license: apache-2.0
+language:
+- en
+library_name: transformers
+pipeline_tag: text-generation
+datasets:
+- HuggingFaceFW/fineweb-edu
+tags:
+- causal-lm
+- language-model
+- base-model
+- small-language-model
+- bananamind
+- bananamind2
+- ternary
+- int8-embeddings
+- digit-tokenizer
+- pytorch
+- safetensors
+- custom-code
+- trust-remote-code
+- custom-architecture
 ---
+
+# TernaryBananaMind-10M
+
+**TernaryBananaMind-10M** is a compact decoder-only base language model from BananaMind. Its inference checkpoint stores projection and output weights as packed ternary values and input embeddings as signed 8-bit values. The intended model ID is `BananaMind/TernaryBananaMind-10M`.
+
+The model has **8,428,032 logical weights**, a **4,096-token context window**, and a custom **2,048-token digit-aware byte-level BPE tokenizer**. Its `model.safetensors` file is **2,175,676 bytes** (about **2.07 bits per logical weight**, including scales, float normalization weights, and file overhead). The model name is a size class; the exact logical weight count is 8.43M.
+
+## Model Details
+
+| Field | Value |
+|---|---:|
+| Architecture | BananaMind 2 style decoder-only Transformer |
+| Logical weights | 8,428,032 |
+| Layers | 10 |
+| Hidden size | 256 |
+| Intermediate size | 704 |
+| Attention heads | 4 |
+| KV heads | 2 |
+| Head dimension | 64 |
+| Attention | Grouped-query attention with QK norm |
+| MLP | SwiGLU |
+| Position embeddings | RoPE, theta 100,000 |
+| Normalization | RMSNorm, epsilon 1e-6 |
+| Vocabulary size | 2,048 |
+| Context length | 4,096 |
+| Input embeddings | Separate from output head; signed 8-bit with one scale per row |
+| Projection and output weights | Ternary (-1, 0, +1), five values per byte |
+| Activations in linear layers | Quantized to 8 bits during inference |
+| Checkpoint | `model.safetensors`, 2,175,676 bytes |
+| HF architecture | `BananaAllForCausalLM` |
+| HF model type | `bananaall` |
+
+The packed weights are decoded to temporary tensors for matrix multiplication. **2.07 bits per weight describes checkpoint storage, not arithmetic precision or peak inference memory.** Packed weights are registered as buffers, so `sum(p.numel() for p in model.parameters())` reports only the 6,656 trainable normalization parameters. Use the logical weight count above for model-size comparisons.
+
+## Tokenizer
+
+The custom 2k byte-level BPE tokenizer isolates digits during pre-tokenization.
+
+| Special token | ID |
+|---|---:|
+| `<\|pad\|>` | 0 |
+| `<\|bos\|>` | 1 |
+| `<\|eos\|>` | 2 |
+| `<\|unk\|>` | 3 |
+
+## Training
+
+The model was trained on **2 billion tokens of FineWeb-Edu** with the **BananaAll training framework**. The included `training_args.bin` records the settings below. Its maximum-step value is a configured limit, not a separately verified final step count.
+
+| Setting | Recorded value |
+|---|---:|
+| Maximum optimizer steps | 61,036 |
+| Per-device micro batch | 8 sequences |
+| Gradient accumulation | 2 |
+| Optimizer | PyTorch fused AdamW |
+| Betas | 0.9, 0.999 |
+| Peak learning rate | 0.0018 |
+| Learning-rate schedule | Cosine |
+| Warmup steps | 1,831 |
+| Weight decay | 0.01 |
+| Gradient clipping | 1.0 |
+| BF16 training | Enabled |
+| PyTorch compile | Enabled |
+| Seed | 37 |
+
+## Evaluation
+
+The `lm_eval` and BananaMind Base Bench 1.1 scores below were supplied for the **current packed checkpoint**. Harness version, dtype, and runtime settings can affect results.
+
+### Standard benchmarks (`lm_eval`)
+
+| Benchmark | Acc | Acc norm | Samples |
+|---|---:|---:|---:|
+| PIQA | 54.46% | 53.54% | 1,838 |
+| ARC Easy | 30.98% | 31.57% | 2,376 |
+| ARC Challenge | 18.09% | 21.08% | 1,172 |
+| HellaSwag | 26.94% | 27.76% | 10,042 |
+| **Four-task mean** | — | **33.49%** | — |
+
+### BananaMind Base Bench 1.1
+
+The reported official complete run scored **37.14% accuracy (130/350)**, **35.95% weighted accuracy**, and **904 overall Elo**.
+
+| Category | Elo | Accuracy | Weighted accuracy |
+|---|---:|---:|---:|
+| Language completion | 963 | 58.00% | 58.28% |
+| Commonsense | 828 | 32.00% | 33.12% |
+| World knowledge | 869 | 36.00% | 38.24% |
+| Context tracking | 832 | 32.00% | 28.34% |
+| Quantitative | 785 | 20.00% | 18.30% |
+| Logical reasoning | 1,043 | 46.00% | 42.32% |
+| Code completion | 1,000 | 36.00% | 37.26% |
+
+### ArithMark 3.0
+
+The [official benchmark script](https://huggingface.co/datasets/AxiomicLabs/Arithmark-3.0/raw/main/bencharithmark-3.py) was run against this local packed checkpoint on all **1,000 examples** with its default batch size and context limit, using CUDA and BF16. The model loaded without missing or unexpected checkpoint keys.
+
+| Metric | Score |
+|---|---:|
+| Raw continuation accuracy | 31.10% (311/1,000) |
+| Length-normalized accuracy | 31.20% (312/1,000) |
+
+Dataset SHA-256: `bf8ab1a5193d52cdf0e05ff0b3ca226bdfcf416cb6e75562dcbe72e7e4559435`.
+
+### Intelligence Index
+
+The **Intelligence Index is 4.90**. Each component is adjusted for its chance floor with `N(score, chance) = 100 × (score - chance) / (100 - chance)`. The calculation uses the length-normalized scores above. Combined ARC is the mean of ARC Easy and ARC Challenge before chance adjustment.
+
+| Component | Score used | Chance floor | Adjusted score | Weight |
+|---|---:|---:|---:|---:|
+| HellaSwag | 27.76% | 25% | 3.68 | 1.00 |
+| Combined ARC | 26.325% | 25% | 1.77 | 1.00 |
+| PIQA | 53.54% | 50% | 7.08 | 1.00 |
+| ArithMark 3.0 | 31.20% | 25% | 8.27 | 0.65 |
+
+`(3.68 + 1.7667 + 7.08 + 0.65 × 8.2667) / 3.65 = 4.9041`, rounded to **4.90**.
+
+## Comparison with BananaMind-2-Nano
+
+| Property | TernaryBananaMind-10M | BananaMind-2-Nano |
+|---|---:|---:|
+| Logical weights | 8,428,032 | 9,968,128 |
+| Vocabulary size | 2,048 | 8,192 |
+| Intermediate size | 704 | 768 |
+| Input/output embeddings | Separate; input is 8-bit | Tied |
+| Training tokens seen | 2B FineWeb-Edu | 30B across four datasets |
+
+| Shared `lm_eval` benchmark | TernaryBananaMind-10M acc norm | BananaMind-2-Nano acc norm | Difference |
+|---|---:|---:|---:|
+| PIQA | 53.54% | 55.98% | -2.44 points |
+| ARC Easy | 31.57% | 36.20% | -4.63 points |
+| ARC Challenge | 21.08% | 23.38% | -2.30 points |
+| HellaSwag | 27.76% | 27.50% | +0.26 points |
+| **Four-task mean** | **33.49%** | **35.77%** | **-2.28 points** |
+
+The two models differ in vocabulary, MLP width, training token count, and weight format, so these scores are a model-level comparison rather than an isolated measure of quantization.
+
+| Additional metric | TernaryBananaMind-10M | BananaMind-2-Nano | Difference |
+|---|---:|---:|---:|
+| Intelligence Index | **4.90** | **8.01** | -3.11 |
+| ArithMark 3.0 acc norm | 31.20% | 33.70% | -2.50 points |
+| BananaMind Base Bench 1.1 Elo | 904 | 917 | -13 |
+| BananaMind Base Bench 1.1 accuracy | 37.14% (130/350) | 40.86% (143/350) | -3.72 points |
+| BananaMind Base Bench 1.1 weighted accuracy | 35.95% | 37.52% | -1.57 points |
+
+### BananaMind Base Bench 1.1 category comparison
+
+| Category | Ternary Elo | Nano Elo | Ternary accuracy | Nano accuracy | Ternary weighted | Nano weighted |
+|---|---:|---:|---:|---:|---:|---:|
+| Language completion | 963 | 1,147 | 58.00% | 80.00% | 58.28% | 80.25% |
+| Commonsense | 828 | 865 | 32.00% | 38.00% | 33.12% | 37.85% |
+| World knowledge | 869 | 895 | 36.00% | 46.00% | 38.24% | 41.69% |
+| Context tracking | 832 | 804 | 32.00% | 26.00% | 28.34% | 25.16% |
+| Quantitative | 785 | 881 | 20.00% | 32.00% | 18.30% | 28.08% |
+| Logical reasoning | 1,043 | 1,043 | 46.00% | 46.00% | 42.32% | 42.32% |
+| Code completion | 1,000 | 827 | 36.00% | 18.00% | 37.26% | 18.15% |
+
+## Usage
+
+This model uses custom architecture code. Load it with `trust_remote_code=True`.
+
+```bash
+pip install -U transformers safetensors torch
+```
+
+```python
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+model_id = "BananaMind/TernaryBananaMind-10M"
+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(model_id, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    trust_remote_code=True,
+    dtype=dtype,
+).to(device).eval()
+
+inputs = tokenizer("The color of the sky is", return_tensors="pt").to(device)
+with torch.inference_mode():
+    output = model.generate(
+        **inputs,
+        max_new_tokens=96,
+        do_sample=True,
+        temperature=0.7,
+        top_p=0.9,
+        pad_token_id=tokenizer.eos_token_id,
+        eos_token_id=tokenizer.eos_token_id,
+    )
+
+print(tokenizer.decode(output[0], skip_special_tokens=True))
+```
+
+## License
+
+Apache 2.0.

config.json

@@ -0,0 +1,32 @@
+{
+  "architecture_style": "bananamind2",
+  "architectures": [
+    "BananaAllForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_bananaall.BananaAllConfig",
+    "AutoModelForCausalLM": "modeling_bananaall_int8_v2.BananaAllForCausalLM"
+  },
+  "bos_token_id": 1,
+  "dtype": "float32",
+  "eos_token_id": 2,
+  "head_dim": 64,
+  "hidden_size": 256,
+  "intermediate_size": 704,
+  "lft": false,
+  "max_position_embeddings": 4096,
+  "model_type": "bananaall",
+  "num_attention_heads": 4,
+  "num_hidden_layers": 10,
+  "num_key_value_heads": 2,
+  "pad_token_id": 0,
+  "rms_norm_eps": 1e-06,
+  "rope_theta": 100000.0,
+  "ternary": true,
+  "tie_word_embeddings": false,
+  "transformers_version": "5.14.1",
+  "use_cache": false,
+  "vocab_size": 2048,
+  "packed_ternary": true,
+  "embedding_bits": 8
+}

configuration_bananaall.py

@@ -0,0 +1,29 @@
+from transformers import PretrainedConfig
+
+
+class BananaAllConfig(PretrainedConfig):
+    model_type = "bananaall"
+
+    def __init__(self, vocab_size=8192, hidden_size=384, num_hidden_layers=14,
+                 num_attention_heads=6, num_key_value_heads=2, head_dim=64,
+                 intermediate_size=1024, max_position_embeddings=4096,
+                 rope_theta=100000.0, rms_norm_eps=1e-6, architecture_style="bananamind2",
+                 lft=False, ternary=False, packed_ternary=False, embedding_bits=None, **kwargs):
+        kwargs.setdefault("tie_word_embeddings", not packed_ternary)
+        super().__init__(**kwargs)
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.head_dim = head_dim
+        self.intermediate_size = intermediate_size
+        self.max_position_embeddings = max_position_embeddings
+        self.rope_theta = rope_theta
+        self.rms_norm_eps = rms_norm_eps
+        self.architecture_style = architecture_style
+        self.lft = lft
+        self.ternary = ternary
+        self.packed_ternary = packed_ternary
+        self.embedding_bits = embedding_bits
+        self.use_cache = False

dataset_tokens.json

@@ -0,0 +1,6 @@
+[
+  {
+    "dataset": "HuggingFaceFW/fineweb-edu",
+    "tokens": 2000000000
+  }
+]

generation_config.json

@@ -0,0 +1,12 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 1,
+  "eos_token_id": [
+    2
+  ],
+  "output_attentions": false,
+  "output_hidden_states": false,
+  "pad_token_id": 0,
+  "transformers_version": "5.14.1",
+  "use_cache": false
+}

model.safetensors

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

model_backup_full.safetensors

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

modeling_bananaall.py

@@ -0,0 +1,270 @@
+"""BananaMind 2 style decoder with optional LFT or ternary quantization.
+
+The model uses pre-RMSNorm, RoPE, grouped-query attention and SwiGLU.
+BananaMind 2 mode also applies QK norm. LFT changes only routing.
+Packed mode keeps five ternary projection weights per byte and can store
+input embeddings as int8 with a separate scale for each vocabulary row.
+"""
+import math
+from pathlib import Path
+import torch
+from torch import nn
+from torch.nn import functional as F
+from transformers import PreTrainedModel
+from transformers.generation import GenerationMixin
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+try:
+    from .configuration_bananaall import BananaAllConfig
+except ImportError:
+    from configuration_bananaall import BananaAllConfig
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, size, eps=1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(size))
+        self.eps = eps
+
+    def forward(self, x):
+        y = x.float()
+        return (y * torch.rsqrt(y.square().mean(-1, keepdim=True) + self.eps) * self.weight.float()).to(x.dtype)
+
+
+class TernaryLinear(nn.Linear):
+    """W1.58A8 fake quantization with straight-through gradients.
+
+    The master weights remain floating point for optimization and checkpoints.
+    This layer does not pack ternary weights or use a low-bit inference kernel.
+    """
+
+    def forward(self, x):
+        weights = self.weight.float()
+        weight_scale = weights.detach().abs().mean().clamp_min(1e-6)
+        quantized_weights = (weights / weight_scale).round().clamp(-1, 1) * weight_scale
+        fake_weights = self.weight + (quantized_weights.to(self.weight.dtype) - self.weight).detach()
+
+        activations = x.float()
+        activation_scale = activations.detach().abs().amax(dim=-1, keepdim=True).clamp_min(1e-6) / 127
+        quantized_activations = (activations / activation_scale).round().clamp(-127, 127) * activation_scale
+        fake_activations = x + (quantized_activations.to(x.dtype) - x).detach()
+        return F.linear(fake_activations, fake_weights, self.bias)
+
+
+def unpack_ternary(packed, width, scale, dtype):
+    """Decode five base-3 weights per byte, with values -1, 0, and 1."""
+    powers = torch.tensor((1, 3, 9, 27, 81), device=packed.device, dtype=torch.int32)
+    values = (packed.to(torch.int32).unsqueeze(-1) // powers) % 3 - 1
+    values = values.flatten(-2)[..., :width]
+    return (values.to(torch.float32) * scale).to(dtype)
+
+
+class PackedTernaryLinear(nn.Module):
+    """Packed resident weights; decoding creates a temporary compute tensor."""
+
+    def __init__(self, in_features, out_features, bias=False):
+        super().__init__()
+        if bias:
+            raise ValueError("PackedTernaryLinear does not support bias")
+        self.in_features = in_features
+        self.out_features = out_features
+        self.register_buffer("packed_weight", torch.zeros(out_features, (in_features + 4) // 5, dtype=torch.uint8))
+        self.register_buffer("weight_scale", torch.ones((), dtype=torch.float32))
+
+    def forward(self, x):
+        weight = unpack_ternary(self.packed_weight, self.in_features, self.weight_scale, x.dtype)
+        activations = x.float()
+        activation_scale = activations.abs().amax(dim=-1, keepdim=True).clamp_min(1e-6) / 127
+        activations = (activations / activation_scale).round().clamp(-127, 127) * activation_scale
+        return F.linear(activations.to(x.dtype), weight)
+
+
+class PackedTernaryEmbedding(nn.Module):
+    def __init__(self, num_embeddings, embedding_dim):
+        super().__init__()
+        self.num_embeddings = num_embeddings
+        self.embedding_dim = embedding_dim
+        self.register_buffer("packed_weight", torch.zeros(num_embeddings, (embedding_dim + 4) // 5, dtype=torch.uint8))
+        self.register_buffer("weight_scale", torch.ones((), dtype=torch.float32))
+
+    def forward(self, input_ids):
+        rows = self.packed_weight[input_ids]
+        return unpack_ternary(rows, self.embedding_dim, self.weight_scale, self.weight_scale.dtype)
+
+
+class Int8Embedding(nn.Module):
+    """Signed 8-bit input embeddings with one float scale per vocabulary row."""
+
+    def __init__(self, num_embeddings, embedding_dim):
+        super().__init__()
+        self.num_embeddings = num_embeddings
+        self.embedding_dim = embedding_dim
+        self.register_buffer("quantized_weight", torch.zeros(num_embeddings, embedding_dim, dtype=torch.int8))
+        self.register_buffer("weight_scale", torch.ones(num_embeddings, 1, dtype=torch.float32))
+
+    def forward(self, input_ids):
+        return self.quantized_weight[input_ids].to(self.weight_scale.dtype) * self.weight_scale[input_ids]
+
+
+def apply_rope(x, theta, position_ids):
+    dim = x.shape[-1]
+    inv = 1.0 / (theta ** (torch.arange(0, dim, 2, device=x.device, dtype=torch.float32) / dim))
+    angles = position_ids.float().unsqueeze(-1) * inv
+    cos = angles.cos().unsqueeze(1).to(x.dtype)
+    sin = angles.sin().unsqueeze(1).to(x.dtype)
+    even, odd = x[..., ::2], x[..., 1::2]
+    return torch.stack((even * cos - odd * sin, even * sin + odd * cos), dim=-1).flatten(-2)
+
+
+class Attention(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        h, d, kv = config.num_attention_heads, config.head_dim, config.num_key_value_heads
+        self.h, self.d, self.kv, self.theta = h, d, kv, config.rope_theta
+        linear = (PackedTernaryLinear if config.packed_ternary else TernaryLinear) if config.ternary else nn.Linear
+        self.q_proj = linear(config.hidden_size, h * d, bias=False)
+        self.k_proj = linear(config.hidden_size, kv * d, bias=False)
+        self.v_proj = linear(config.hidden_size, kv * d, bias=False)
+        self.o_proj = linear(h * d, config.hidden_size, bias=False)
+        self.q_norm = RMSNorm(d, config.rms_norm_eps) if config.architecture_style == "bananamind2" else nn.Identity()
+        self.k_norm = RMSNorm(d, config.rms_norm_eps) if config.architecture_style == "bananamind2" else nn.Identity()
+
+    def forward(self, x, attention_mask=None):
+        b, t, _ = x.shape
+        q = self.q_norm(self.q_proj(x).view(b, t, self.h, self.d).transpose(1, 2))
+        k = self.k_norm(self.k_proj(x).view(b, t, self.kv, self.d).transpose(1, 2))
+        v = self.v_proj(x).view(b, t, self.kv, self.d).transpose(1, 2)
+        positions = torch.arange(t, device=x.device).unsqueeze(0)
+        q, k = apply_rope(q, self.theta, positions), apply_rope(k, self.theta, positions)
+        k = k.repeat_interleave(self.h // self.kv, dim=1)
+        v = v.repeat_interleave(self.h // self.kv, dim=1)
+        if attention_mask is None:
+            y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+        else:
+            causal = torch.ones(t, t, device=x.device, dtype=torch.bool).tril()
+            mask = causal[None, None] & attention_mask[:, None, None, :].bool()
+            y = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
+        return self.o_proj(y.transpose(1, 2).contiguous().view(b, t, self.h * self.d))
+
+
+class Block(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.norm1 = RMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.attn = Attention(config)
+        self.norm2 = RMSNorm(config.hidden_size, config.rms_norm_eps)
+        linear = (PackedTernaryLinear if config.packed_ternary else TernaryLinear) if config.ternary else nn.Linear
+        self.gate_proj = linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.up_proj = linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.down_proj = linear(config.intermediate_size, config.hidden_size, bias=False)
+
+    def forward(self, x, attention_mask=None):
+        x = x + self.attn(self.norm1(x), attention_mask)
+        z = self.norm2(x)
+        return x + self.down_proj(F.silu(self.gate_proj(z)) * self.up_proj(z))
+
+
+class BananaAllForCausalLM(PreTrainedModel, GenerationMixin):
+    config_class = BananaAllConfig
+    base_model_prefix = "model"
+    _supports_sdpa = True
+    # forward() returns a mean loss per microbatch and ignores **kwargs.
+    # Tell Trainer to divide it by the gradient-accumulation count.
+    accepts_loss_kwargs = False
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
+        """Use the weight file's layout even when a caller supplies a default config."""
+        packed_checkpoint = False
+        if isinstance(pretrained_model_name_or_path, (str, Path)):
+            checkpoint = Path(pretrained_model_name_or_path) / "model.safetensors"
+            if checkpoint.is_file():
+                from safetensors import safe_open
+
+                with safe_open(str(checkpoint), framework="pt", device="cpu") as weights:
+                    keys = set(weights.keys())
+                if "layers.0.attn.q_proj.packed_weight" in keys:
+                    packed_checkpoint = True
+                    config = kwargs.get("config")
+                    if not isinstance(config, cls.config_class):
+                        config = cls.config_class.from_pretrained(pretrained_model_name_or_path)
+                    config.ternary = True
+                    config.packed_ternary = True
+                    config.embedding_bits = 8 if "embed_tokens.quantized_weight" in keys else None
+                    config.tie_word_embeddings = False
+                    kwargs["config"] = config
+        if not packed_checkpoint:
+            return super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)
+
+        return_loading_info = kwargs.get("output_loading_info", False)
+        kwargs["output_loading_info"] = True
+        model, loading_info = super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)
+        if loading_info["missing_keys"] or loading_info["unexpected_keys"]:
+            raise RuntimeError(
+                "Packed BananaAll checkpoint did not load cleanly: "
+                f"{len(loading_info['missing_keys'])} missing and "
+                f"{len(loading_info['unexpected_keys'])} unexpected keys"
+            )
+        return (model, loading_info) if return_loading_info else model
+
+    def __init__(self, config):
+        super().__init__(config)
+        if config.packed_ternary and config.embedding_bits == 8:
+            self.embed_tokens = Int8Embedding(config.vocab_size, config.hidden_size)
+        elif config.packed_ternary:
+            self.embed_tokens = PackedTernaryEmbedding(config.vocab_size, config.hidden_size)
+        else:
+            self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.layers = nn.ModuleList([Block(config) for _ in range(config.num_hidden_layers)])
+        self.norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
+        linear = (PackedTernaryLinear if config.packed_ternary else TernaryLinear) if config.ternary else nn.Linear
+        self.lm_head = linear(config.hidden_size, config.vocab_size, bias=False)
+        self.post_init()
+        self.tie_weights()
+
+    def get_input_embeddings(self):
+        return self.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.embed_tokens = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, value):
+        self.lm_head = value
+
+    def forward(self, input_ids=None, attention_mask=None, labels=None, **kwargs):
+        x = self.embed_tokens(input_ids)
+        if self.config.lft and len(self.layers) > 1:
+            x = self.layers[0](x, attention_mask)
+            for i in range(1, len(self.layers)):
+                x = self.layers[i](x, attention_mask)
+                if i < len(self.layers) - 1:
+                    x = self.layers[i - 1](x, attention_mask)
+                    x = self.layers[i](x, attention_mask)
+        else:
+            for layer in self.layers:
+                x = layer(x, attention_mask)
+        logits = self.lm_head(self.norm(x))
+        loss = None
+        if labels is not None:
+            shifted_logits = logits[:, :-1, :].contiguous().float()
+            shifted_labels = labels[:, 1:].contiguous()
+            loss = F.cross_entropy(shifted_logits.view(-1, shifted_logits.size(-1)), shifted_labels.view(-1), ignore_index=-100)
+        return CausalLMOutputWithPast(loss=loss, logits=logits, past_key_values=None)
+
+    def prepare_inputs_for_generation(self, input_ids, attention_mask=None, **kwargs):
+        return {"input_ids": input_ids, "attention_mask": attention_mask}
+
+
+def register():
+    from transformers import AutoConfig, AutoModelForCausalLM
+    try:
+        AutoConfig.register("bananaall", BananaAllConfig)
+    except ValueError:
+        pass
+    try:
+        AutoModelForCausalLM.register(BananaAllConfig, BananaAllForCausalLM)
+    except ValueError:
+        pass

modeling_bananaall_int8_v2.py

@@ -0,0 +1,270 @@
+"""BananaMind 2 style decoder with optional LFT or ternary quantization.
+
+The model uses pre-RMSNorm, RoPE, grouped-query attention and SwiGLU.
+BananaMind 2 mode also applies QK norm. LFT changes only routing.
+Packed mode keeps five ternary projection weights per byte and can store
+input embeddings as int8 with a separate scale for each vocabulary row.
+"""
+import math
+from pathlib import Path
+import torch
+from torch import nn
+from torch.nn import functional as F
+from transformers import PreTrainedModel
+from transformers.generation import GenerationMixin
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+try:
+    from .configuration_bananaall import BananaAllConfig
+except ImportError:
+    from configuration_bananaall import BananaAllConfig
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, size, eps=1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(size))
+        self.eps = eps
+
+    def forward(self, x):
+        y = x.float()
+        return (y * torch.rsqrt(y.square().mean(-1, keepdim=True) + self.eps) * self.weight.float()).to(x.dtype)
+
+
+class TernaryLinear(nn.Linear):
+    """W1.58A8 fake quantization with straight-through gradients.
+
+    The master weights remain floating point for optimization and checkpoints.
+    This layer does not pack ternary weights or use a low-bit inference kernel.
+    """
+
+    def forward(self, x):
+        weights = self.weight.float()
+        weight_scale = weights.detach().abs().mean().clamp_min(1e-6)
+        quantized_weights = (weights / weight_scale).round().clamp(-1, 1) * weight_scale
+        fake_weights = self.weight + (quantized_weights.to(self.weight.dtype) - self.weight).detach()
+
+        activations = x.float()
+        activation_scale = activations.detach().abs().amax(dim=-1, keepdim=True).clamp_min(1e-6) / 127
+        quantized_activations = (activations / activation_scale).round().clamp(-127, 127) * activation_scale
+        fake_activations = x + (quantized_activations.to(x.dtype) - x).detach()
+        return F.linear(fake_activations, fake_weights, self.bias)
+
+
+def unpack_ternary(packed, width, scale, dtype):
+    """Decode five base-3 weights per byte, with values -1, 0, and 1."""
+    powers = torch.tensor((1, 3, 9, 27, 81), device=packed.device, dtype=torch.int32)
+    values = (packed.to(torch.int32).unsqueeze(-1) // powers) % 3 - 1
+    values = values.flatten(-2)[..., :width]
+    return (values.to(torch.float32) * scale).to(dtype)
+
+
+class PackedTernaryLinear(nn.Module):
+    """Packed resident weights; decoding creates a temporary compute tensor."""
+
+    def __init__(self, in_features, out_features, bias=False):
+        super().__init__()
+        if bias:
+            raise ValueError("PackedTernaryLinear does not support bias")
+        self.in_features = in_features
+        self.out_features = out_features
+        self.register_buffer("packed_weight", torch.zeros(out_features, (in_features + 4) // 5, dtype=torch.uint8))
+        self.register_buffer("weight_scale", torch.ones((), dtype=torch.float32))
+
+    def forward(self, x):
+        weight = unpack_ternary(self.packed_weight, self.in_features, self.weight_scale, x.dtype)
+        activations = x.float()
+        activation_scale = activations.abs().amax(dim=-1, keepdim=True).clamp_min(1e-6) / 127
+        activations = (activations / activation_scale).round().clamp(-127, 127) * activation_scale
+        return F.linear(activations.to(x.dtype), weight)
+
+
+class PackedTernaryEmbedding(nn.Module):
+    def __init__(self, num_embeddings, embedding_dim):
+        super().__init__()
+        self.num_embeddings = num_embeddings
+        self.embedding_dim = embedding_dim
+        self.register_buffer("packed_weight", torch.zeros(num_embeddings, (embedding_dim + 4) // 5, dtype=torch.uint8))
+        self.register_buffer("weight_scale", torch.ones((), dtype=torch.float32))
+
+    def forward(self, input_ids):
+        rows = self.packed_weight[input_ids]
+        return unpack_ternary(rows, self.embedding_dim, self.weight_scale, self.weight_scale.dtype)
+
+
+class Int8Embedding(nn.Module):
+    """Signed 8-bit input embeddings with one float scale per vocabulary row."""
+
+    def __init__(self, num_embeddings, embedding_dim):
+        super().__init__()
+        self.num_embeddings = num_embeddings
+        self.embedding_dim = embedding_dim
+        self.register_buffer("quantized_weight", torch.zeros(num_embeddings, embedding_dim, dtype=torch.int8))
+        self.register_buffer("weight_scale", torch.ones(num_embeddings, 1, dtype=torch.float32))
+
+    def forward(self, input_ids):
+        return self.quantized_weight[input_ids].to(self.weight_scale.dtype) * self.weight_scale[input_ids]
+
+
+def apply_rope(x, theta, position_ids):
+    dim = x.shape[-1]
+    inv = 1.0 / (theta ** (torch.arange(0, dim, 2, device=x.device, dtype=torch.float32) / dim))
+    angles = position_ids.float().unsqueeze(-1) * inv
+    cos = angles.cos().unsqueeze(1).to(x.dtype)
+    sin = angles.sin().unsqueeze(1).to(x.dtype)
+    even, odd = x[..., ::2], x[..., 1::2]
+    return torch.stack((even * cos - odd * sin, even * sin + odd * cos), dim=-1).flatten(-2)
+
+
+class Attention(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        h, d, kv = config.num_attention_heads, config.head_dim, config.num_key_value_heads
+        self.h, self.d, self.kv, self.theta = h, d, kv, config.rope_theta
+        linear = (PackedTernaryLinear if config.packed_ternary else TernaryLinear) if config.ternary else nn.Linear
+        self.q_proj = linear(config.hidden_size, h * d, bias=False)
+        self.k_proj = linear(config.hidden_size, kv * d, bias=False)
+        self.v_proj = linear(config.hidden_size, kv * d, bias=False)
+        self.o_proj = linear(h * d, config.hidden_size, bias=False)
+        self.q_norm = RMSNorm(d, config.rms_norm_eps) if config.architecture_style == "bananamind2" else nn.Identity()
+        self.k_norm = RMSNorm(d, config.rms_norm_eps) if config.architecture_style == "bananamind2" else nn.Identity()
+
+    def forward(self, x, attention_mask=None):
+        b, t, _ = x.shape
+        q = self.q_norm(self.q_proj(x).view(b, t, self.h, self.d).transpose(1, 2))
+        k = self.k_norm(self.k_proj(x).view(b, t, self.kv, self.d).transpose(1, 2))
+        v = self.v_proj(x).view(b, t, self.kv, self.d).transpose(1, 2)
+        positions = torch.arange(t, device=x.device).unsqueeze(0)
+        q, k = apply_rope(q, self.theta, positions), apply_rope(k, self.theta, positions)
+        k = k.repeat_interleave(self.h // self.kv, dim=1)
+        v = v.repeat_interleave(self.h // self.kv, dim=1)
+        if attention_mask is None:
+            y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+        else:
+            causal = torch.ones(t, t, device=x.device, dtype=torch.bool).tril()
+            mask = causal[None, None] & attention_mask[:, None, None, :].bool()
+            y = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
+        return self.o_proj(y.transpose(1, 2).contiguous().view(b, t, self.h * self.d))
+
+
+class Block(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.norm1 = RMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.attn = Attention(config)
+        self.norm2 = RMSNorm(config.hidden_size, config.rms_norm_eps)
+        linear = (PackedTernaryLinear if config.packed_ternary else TernaryLinear) if config.ternary else nn.Linear
+        self.gate_proj = linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.up_proj = linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.down_proj = linear(config.intermediate_size, config.hidden_size, bias=False)
+
+    def forward(self, x, attention_mask=None):
+        x = x + self.attn(self.norm1(x), attention_mask)
+        z = self.norm2(x)
+        return x + self.down_proj(F.silu(self.gate_proj(z)) * self.up_proj(z))
+
+
+class BananaAllForCausalLM(PreTrainedModel, GenerationMixin):
+    config_class = BananaAllConfig
+    base_model_prefix = "model"
+    _supports_sdpa = True
+    # forward() returns a mean loss per microbatch and ignores **kwargs.
+    # Tell Trainer to divide it by the gradient-accumulation count.
+    accepts_loss_kwargs = False
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
+        """Use the weight file's layout even when a caller supplies a default config."""
+        packed_checkpoint = False
+        if isinstance(pretrained_model_name_or_path, (str, Path)):
+            checkpoint = Path(pretrained_model_name_or_path) / "model.safetensors"
+            if checkpoint.is_file():
+                from safetensors import safe_open
+
+                with safe_open(str(checkpoint), framework="pt", device="cpu") as weights:
+                    keys = set(weights.keys())
+                if "layers.0.attn.q_proj.packed_weight" in keys:
+                    packed_checkpoint = True
+                    config = kwargs.get("config")
+                    if not isinstance(config, cls.config_class):
+                        config = cls.config_class.from_pretrained(pretrained_model_name_or_path)
+                    config.ternary = True
+                    config.packed_ternary = True
+                    config.embedding_bits = 8 if "embed_tokens.quantized_weight" in keys else None
+                    config.tie_word_embeddings = False
+                    kwargs["config"] = config
+        if not packed_checkpoint:
+            return super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)
+
+        return_loading_info = kwargs.get("output_loading_info", False)
+        kwargs["output_loading_info"] = True
+        model, loading_info = super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)
+        if loading_info["missing_keys"] or loading_info["unexpected_keys"]:
+            raise RuntimeError(
+                "Packed BananaAll checkpoint did not load cleanly: "
+                f"{len(loading_info['missing_keys'])} missing and "
+                f"{len(loading_info['unexpected_keys'])} unexpected keys"
+            )
+        return (model, loading_info) if return_loading_info else model
+
+    def __init__(self, config):
+        super().__init__(config)
+        if config.packed_ternary and config.embedding_bits == 8:
+            self.embed_tokens = Int8Embedding(config.vocab_size, config.hidden_size)
+        elif config.packed_ternary:
+            self.embed_tokens = PackedTernaryEmbedding(config.vocab_size, config.hidden_size)
+        else:
+            self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.layers = nn.ModuleList([Block(config) for _ in range(config.num_hidden_layers)])
+        self.norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
+        linear = (PackedTernaryLinear if config.packed_ternary else TernaryLinear) if config.ternary else nn.Linear
+        self.lm_head = linear(config.hidden_size, config.vocab_size, bias=False)
+        self.post_init()
+        self.tie_weights()
+
+    def get_input_embeddings(self):
+        return self.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.embed_tokens = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, value):
+        self.lm_head = value
+
+    def forward(self, input_ids=None, attention_mask=None, labels=None, **kwargs):
+        x = self.embed_tokens(input_ids)
+        if self.config.lft and len(self.layers) > 1:
+            x = self.layers[0](x, attention_mask)
+            for i in range(1, len(self.layers)):
+                x = self.layers[i](x, attention_mask)
+                if i < len(self.layers) - 1:
+                    x = self.layers[i - 1](x, attention_mask)
+                    x = self.layers[i](x, attention_mask)
+        else:
+            for layer in self.layers:
+                x = layer(x, attention_mask)
+        logits = self.lm_head(self.norm(x))
+        loss = None
+        if labels is not None:
+            shifted_logits = logits[:, :-1, :].contiguous().float()
+            shifted_labels = labels[:, 1:].contiguous()
+            loss = F.cross_entropy(shifted_logits.view(-1, shifted_logits.size(-1)), shifted_labels.view(-1), ignore_index=-100)
+        return CausalLMOutputWithPast(loss=loss, logits=logits, past_key_values=None)
+
+    def prepare_inputs_for_generation(self, input_ids, attention_mask=None, **kwargs):
+        return {"input_ids": input_ids, "attention_mask": attention_mask}
+
+
+def register():
+    from transformers import AutoConfig, AutoModelForCausalLM
+    try:
+        AutoConfig.register("bananaall", BananaAllConfig)
+    except ValueError:
+        pass
+    try:
+        AutoModelForCausalLM.register(BananaAllConfig, BananaAllForCausalLM)
+    except ValueError:
+        pass

pack_ternary.py

@@ -0,0 +1,77 @@
+"""Convert the original checkpoint to ternary weights and 8-bit embeddings.
+
+Run from this directory with: python pack_ternary.py
+The source defaults to model_backup_full.safetensors. Its contents are never edited.
+"""
+
+import json
+import os
+from pathlib import Path
+
+import torch
+from safetensors.torch import load_file, save_file
+
+
+ROOT = Path(__file__).resolve().parent
+SOURCE = ROOT / "model_backup_full.safetensors"
+TARGET = ROOT / "model.safetensors"
+
+
+def pack_matrix(weight):
+    weight = weight.float().contiguous()
+    scale = weight.abs().mean().clamp_min(1e-6)
+    trits = (weight / scale).round().clamp(-1, 1).to(torch.uint8) + 1
+    padding = (-trits.shape[1]) % 5
+    if padding:
+        trits = torch.nn.functional.pad(trits, (0, padding), value=1)
+    packed = trits.reshape(trits.shape[0], -1, 5).to(torch.int32)
+    powers = torch.tensor((1, 3, 9, 27, 81), dtype=torch.int32)
+    packed = (packed * powers).sum(dim=-1).to(torch.uint8)
+    return packed.contiguous(), scale.contiguous()
+
+
+def quantize_embedding(weight):
+    weight = weight.float().contiguous()
+    scale = weight.abs().amax(dim=1, keepdim=True).clamp_min(1e-6) / 127
+    values = (weight / scale).round().clamp(-127, 127).to(torch.int8)
+    return values.contiguous(), scale.contiguous()
+
+
+def main():
+    if not SOURCE.is_file():
+        raise FileNotFoundError(SOURCE)
+    config_path = ROOT / "config.json"
+    config = json.loads(config_path.read_text())
+    if not config.get("ternary"):
+        raise ValueError("This converter requires a ternary model")
+
+    original = load_file(str(SOURCE), device="cpu")
+    converted = {}
+    for key, tensor in original.items():
+        if key == "embed_tokens.weight":
+            values, scale = quantize_embedding(tensor)
+            converted["embed_tokens.quantized_weight"] = values
+            converted["embed_tokens.weight_scale"] = scale
+        elif key.endswith("_proj.weight") or key == "lm_head.weight":
+            packed, scale = pack_matrix(tensor)
+            prefix = key.removesuffix(".weight")
+            converted[prefix + ".packed_weight"] = packed
+            converted[prefix + ".weight_scale"] = scale
+        else:
+            converted[key] = tensor.contiguous()
+
+    output = TARGET.with_suffix(".safetensors.tmp")
+    save_file(converted, str(output), metadata={"format": "pt", "quantization": "ternary_base3_5_per_byte_int8_embeddings"})
+    os.replace(output, TARGET)
+    config["packed_ternary"] = True
+    config["embedding_bits"] = 8
+    # The original file contains different embedding and output matrices.
+    config["tie_word_embeddings"] = False
+    config_path.write_text(json.dumps(config, indent=2) + "\n")
+
+    dense_count = sum(t.numel() for t in original.values())
+    print(f"Packed {len(original)} tensors: {TARGET.stat().st_size:,} bytes, {TARGET.stat().st_size * 8 / dense_count:.3f} bits per original weight")
+
+
+if __name__ == "__main__":
+    main()

tokenizer_config.json

@@ -0,0 +1,12 @@
+{
+  "backend": "tokenizers",
+  "bos_token": "<|bos|>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|eos|>",
+  "is_local": false,
+  "local_files_only": false,
+  "model_max_length": 4096,
+  "pad_token": "<|pad|>",
+  "tokenizer_class": "TokenizersBackend",
+  "unk_token": "<|unk|>"
+}

training_args.bin

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