Upload pico-30 checkpoint

README.md

@@ -0,0 +1,67 @@
+---
+language:
+- en
+library_name: transformers
+pipeline_tag: text-generation
+datasets:
+- epfml/FineWeb-HQ
+- HuggingFaceTB/smollm-corpus
+tags:
+- causal-lm
+- base-model
+- muon
+- xsa-refresh
+- custom-code
+- trust-remote-code
+---
+
+# BananaMind 2 Pico Test - 30%
+
+This is the **30% checkpoint** of a 900,002-parameter
+base causal language model. It is not instruction tuned.
+
+## Architecture
+
+| Field | Value |
+|---|---:|
+| Parameters | 900,002 |
+| Layers / hidden size | 6 / 96 |
+| SwiGLU intermediate size | 380 |
+| Query / KV heads | 6 / 2 |
+| Head dimension | 16 |
+| Context | 4,096 |
+| Vocabulary | 384, tied |
+| Refresh layers | 4 and 6 |
+| Refresh kernel | Causal depthwise, width 9 |
+
+The selective XSA refresh gate reads detached attention output as its signal,
+reinjects the original input embedding as its value, and carries convolution
+history alongside the K/V cache. Its learned residual scalar starts at zero.
+
+## Training
+
+| Field | Value |
+|---|---:|
+| Progress | 30% |
+| Tokens seen | 60,001,615,872 |
+| Target tokens | 200,000,000,000 |
+| Hardware | 4 x NVIDIA H200 |
+| Matrix optimizer | Stock `torch.optim.Muon` |
+| Muon peak LR | 0.07 |
+| Embedding/control optimizer | AdamW, LR 0.004 |
+| Precision | bfloat16 autocast |
+
+| Token range | FineWeb-HQ | Cosmopedia v2 |
+|---|---:|---:|
+| 0.00B-100.00B | 80% | 20% |
+| 100.00B-200.00B | 60% | 40% |
+
+## Usage
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+model_id = "Banaxi-Tech/pico-test"
+tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True)
+```

checkpoint_metadata.json

@@ -0,0 +1,72 @@
+{
+  "parameters": 900002,
+  "architecture": {
+    "vocab_size": 384,
+    "hidden_size": 96,
+    "num_hidden_layers": 6,
+    "num_attention_heads": 6,
+    "num_key_value_heads": 2,
+    "head_dim": 16,
+    "intermediate_size": 380,
+    "refresh_layers": [
+      3,
+      5
+    ],
+    "refresh_kernel_size": 9,
+    "refresh_alpha_init": 0.0,
+    "max_position_embeddings": 4096,
+    "rope_theta": 100000.0,
+    "rms_norm_eps": 1e-06,
+    "initializer_range": 0.02,
+    "use_cache": true,
+    "transformers_version": "5.15.0",
+    "architectures": null,
+    "output_hidden_states": false,
+    "return_dict": true,
+    "dtype": null,
+    "chunk_size_feed_forward": 0,
+    "is_encoder_decoder": false,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "problem_type": null,
+    "_name_or_path": "",
+    "tie_word_embeddings": true,
+    "bos_token_id": 1,
+    "eos_token_id": 2,
+    "pad_token_id": 0,
+    "unk_token_id": 3,
+    "model_type": "bananamind2_pico",
+    "output_attentions": false
+  },
+  "training_percent": 30,
+  "step": 6732,
+  "total_steps": 22439,
+  "tokens_seen": 60001615872,
+  "target_tokens": 200000000000,
+  "trained_token_target": 199996473344,
+  "tokens_per_step": 8912896,
+  "local_batch": 544,
+  "effective_global_batch": 2176,
+  "world_size": 4,
+  "gpu_name": "NVIDIA H200",
+  "optimizer": "Muon + AdamW",
+  "muon_peak_lr": 0.07,
+  "adamw_peak_lr": 0.004,
+  "architecture_revision": "ee5739850035c878ce4d9fc29fc38b075987e03b",
+  "tokenizer_revision": "ee5739850035c878ce4d9fc29fc38b075987e03b",
+  "dataset_revisions": {
+    "fineweb_hq": "e58199cdd52438d94405df1a4d8630cc5f13bf84",
+    "cosmopedia_v2": "3ba9d605774198c5868892d7a8deda78031a781f"
+  },
+  "source_tokens": {
+    "fineweb_hq": 48004857856,
+    "cosmopedia_v2": 11996758016
+  },
+  "elapsed_seconds": 8633.815947532654
+}

config.json

@@ -0,0 +1,51 @@
+{
+  "vocab_size": 384,
+  "hidden_size": 96,
+  "num_hidden_layers": 6,
+  "num_attention_heads": 6,
+  "num_key_value_heads": 2,
+  "head_dim": 16,
+  "intermediate_size": 380,
+  "refresh_layers": [
+    3,
+    5
+  ],
+  "refresh_kernel_size": 9,
+  "refresh_alpha_init": 0.0,
+  "max_position_embeddings": 4096,
+  "rope_theta": 100000.0,
+  "rms_norm_eps": 1e-06,
+  "initializer_range": 0.02,
+  "use_cache": true,
+  "transformers_version": "5.15.0",
+  "architectures": [
+    "BananaMind2PicoForCausalLM"
+  ],
+  "output_hidden_states": false,
+  "return_dict": true,
+  "dtype": null,
+  "chunk_size_feed_forward": 0,
+  "is_encoder_decoder": false,
+  "id2label": {
+    "0": "LABEL_0",
+    "1": "LABEL_1"
+  },
+  "label2id": {
+    "LABEL_0": 0,
+    "LABEL_1": 1
+  },
+  "problem_type": null,
+  "_name_or_path": "Banaxi-Tech/pico-test",
+  "tie_word_embeddings": true,
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "pad_token_id": 0,
+  "unk_token_id": 3,
+  "model_type": "bananamind2_pico",
+  "output_attentions": false,
+  "auto_map": {
+    "AutoConfig": "configuration_bananamind2pico.BananaMind2PicoConfig",
+    "AutoModelForCausalLM": "modeling_bananamind2pico.BananaMind2PicoForCausalLM"
+  },
+  "torch_dtype": "float32"
+}

configuration_bananamind2pico.py

@@ -0,0 +1,111 @@
+"""Configuration for BananaMind 2 Pico."""
+
+from __future__ import annotations
+
+from transformers import PretrainedConfig
+
+
+class BananaMind2PicoConfig(PretrainedConfig):
+    model_type = "bananamind2_pico"
+
+    def __init__(
+        self,
+        vocab_size: int = 384,
+        hidden_size: int = 96,
+        num_hidden_layers: int = 6,
+        num_attention_heads: int = 6,
+        num_key_value_heads: int = 2,
+        head_dim: int = 16,
+        intermediate_size: int = 380,
+        refresh_layers: tuple[int, ...] | list[int] = (3, 5),
+        refresh_kernel_size: int = 9,
+        refresh_alpha_init: float = 0.0,
+        max_position_embeddings: int = 4096,
+        rope_theta: float = 100000.0,
+        rms_norm_eps: float = 1e-6,
+        initializer_range: float = 0.02,
+        tie_word_embeddings: bool = True,
+        use_cache: bool = True,
+        bos_token_id: int = 1,
+        eos_token_id: int = 2,
+        pad_token_id: int = 0,
+        unk_token_id: int = 3,
+        **kwargs,
+    ):
+        if hidden_size != num_attention_heads * head_dim:
+            raise ValueError("hidden_size must equal num_attention_heads * head_dim")
+        if num_attention_heads % num_key_value_heads:
+            raise ValueError("num_attention_heads must be divisible by num_key_value_heads")
+        if intermediate_size <= 0:
+            raise ValueError("intermediate_size must be positive")
+        if refresh_kernel_size <= 0 or refresh_kernel_size % 2 == 0:
+            raise ValueError("refresh_kernel_size must be a positive odd integer")
+        refresh_layers = tuple(int(index) for index in refresh_layers)
+        if len(refresh_layers) != len(set(refresh_layers)):
+            raise ValueError("refresh_layers must not contain duplicates")
+        if any(index < 0 or index >= num_hidden_layers for index in refresh_layers):
+            raise ValueError("refresh_layers contains an out-of-range zero-based index")
+
+        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.refresh_layers = list(refresh_layers)
+        self.refresh_kernel_size = refresh_kernel_size
+        self.refresh_alpha_init = refresh_alpha_init
+        self.max_position_embeddings = max_position_embeddings
+        self.rope_theta = rope_theta
+        self.rms_norm_eps = rms_norm_eps
+        self.initializer_range = initializer_range
+        self.use_cache = use_cache
+        super().__init__(
+            tie_word_embeddings=tie_word_embeddings,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            pad_token_id=pad_token_id,
+            unk_token_id=unk_token_id,
+            **kwargs,
+        )
+
+    def parameter_counts(self) -> dict[str, int]:
+        """Return an analytical parameter breakdown for drift checks."""
+        hidden = self.hidden_size
+        query_width = self.num_attention_heads * self.head_dim
+        kv_width = self.num_key_value_heads * self.head_dim
+        embedding = self.vocab_size * hidden
+        attention = hidden * query_width + 2 * hidden * kv_width + query_width * hidden
+        qk_norms = 2 * self.head_dim
+        mlp = 3 * hidden * self.intermediate_size
+        block_norms = 2 * hidden
+        refresh = (
+            3 * hidden * hidden
+            + hidden * self.refresh_kernel_size
+            + 3 * hidden
+            + 1
+        )
+        base_layer = attention + qk_norms + mlp + block_norms
+        total = (
+            embedding
+            + self.num_hidden_layers * base_layer
+            + len(self.refresh_layers) * refresh
+            + hidden
+        )
+        if not self.tie_word_embeddings:
+            total += embedding
+        return {
+            "total": total,
+            "embedding": embedding,
+            "base_layer": base_layer,
+            "attention_per_layer": attention + qk_norms,
+            "mlp_per_layer": mlp,
+            "refresh_per_selected_layer": refresh,
+            "refresh_layer_count": len(self.refresh_layers),
+        }
+
+
+BananaMind2PicoConfig.register_for_auto_class("AutoConfig")
+
+__all__ = ["BananaMind2PicoConfig"]

generation_config.json

@@ -0,0 +1,7 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "pad_token_id": 0,
+  "transformers_version": "5"
+}

model.safetensors

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

modeling_bananamind2pico.py

@@ -0,0 +1,456 @@
+"""BananaMind 2 Pico causal LM with selective XSA refresh gates."""
+
+from __future__ import annotations
+
+import math
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PreTrainedModel
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.generation.utils import GenerationMixin
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+try:
+    from .configuration_bananamind2pico import BananaMind2PicoConfig
+except ImportError:  # Standalone training jobs import the mounted/downloaded files.
+    from configuration_bananamind2pico import BananaMind2PicoConfig
+
+
+class BananaMind2PicoRMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        states = hidden_states.float()
+        states = states * torch.rsqrt(states.square().mean(-1, keepdim=True) + self.eps)
+        return (states * self.weight.float()).to(hidden_states.dtype)
+
+
+def _rope_cos_sin(
+    head_dim: int,
+    positions: torch.Tensor,
+    theta: float,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    inv_freq = 1.0 / (
+        theta
+        ** (
+            torch.arange(0, head_dim, 2, dtype=torch.float32, device=positions.device)
+            / head_dim
+        )
+    )
+    frequencies = torch.outer(positions.float(), inv_freq)
+    return frequencies.cos(), frequencies.sin()
+
+
+def _apply_rope(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    cosine: torch.Tensor,
+    sine: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    query_dtype, key_dtype = query.dtype, key.dtype
+    cosine = cosine[None, None, :, :]
+    sine = sine[None, None, :, :]
+    query_pairs = query.float().reshape(*query.shape[:-1], -1, 2)
+    key_pairs = key.float().reshape(*key.shape[:-1], -1, 2)
+    query_even, query_odd = query_pairs.unbind(-1)
+    key_even, key_odd = key_pairs.unbind(-1)
+    query = torch.stack(
+        (query_even * cosine - query_odd * sine, query_even * sine + query_odd * cosine),
+        dim=-1,
+    ).flatten(-2)
+    key = torch.stack(
+        (key_even * cosine - key_odd * sine, key_even * sine + key_odd * cosine),
+        dim=-1,
+    ).flatten(-2)
+    return query.to(query_dtype), key.to(key_dtype)
+
+
+class BananaMind2PicoCache(DynamicCache):
+    """K/V cache plus causal-convolution history for selected refresh layers."""
+
+    def __init__(self, config: BananaMind2PicoConfig):
+        try:
+            super().__init__(config=config)
+        except TypeError:
+            super().__init__()
+        self.refresh_states: list[torch.Tensor | None] = [
+            None for _ in range(config.num_hidden_layers)
+        ]
+
+    def refresh_input(
+        self,
+        layer_idx: int,
+        current: torch.Tensor,
+        history_size: int,
+    ) -> torch.Tensor:
+        history = self.refresh_states[layer_idx]
+        if history is None:
+            history = current.new_zeros(current.size(0), current.size(1), history_size)
+        elif history.size(-1) < history_size:
+            history = F.pad(history, (history_size - history.size(-1), 0))
+        convolution_input = torch.cat((history, current), dim=-1)
+        self.refresh_states[layer_idx] = convolution_input[..., -history_size:]
+        return convolution_input
+
+    def reorder_cache(self, beam_idx: torch.LongTensor):
+        super().reorder_cache(beam_idx)
+        self.refresh_states = [
+            None if state is None else state.index_select(0, beam_idx.to(state.device))
+            for state in self.refresh_states
+        ]
+
+    def batch_repeat_interleave(self, repeats: int):
+        super().batch_repeat_interleave(repeats)
+        self.refresh_states = [
+            None if state is None else state.repeat_interleave(repeats, dim=0)
+            for state in self.refresh_states
+        ]
+
+    def batch_select_indices(self, indices: torch.Tensor):
+        super().batch_select_indices(indices)
+        self.refresh_states = [
+            None if state is None else state.index_select(0, indices.to(state.device))
+            for state in self.refresh_states
+        ]
+
+    def crop(self, max_length: int):
+        current_length = self.get_seq_length()
+        target_length = current_length + max_length if max_length < 0 else max_length
+        if target_length < current_length:
+            raise NotImplementedError(
+                "BananaMind2PicoCache cannot roll back causal refresh state"
+            )
+        super().crop(max_length)
+
+
+class BananaMind2PicoAttention(nn.Module):
+    def __init__(self, config: BananaMind2PicoConfig, layer_idx: int):
+        super().__init__()
+        self.layer_idx = layer_idx
+        self.num_heads = config.num_attention_heads
+        self.num_kv_heads = config.num_key_value_heads
+        self.head_dim = config.head_dim
+        self.num_kv_groups = self.num_heads // self.num_kv_heads
+        self.q_proj = nn.Linear(config.hidden_size, self.num_heads * self.head_dim, bias=False)
+        self.k_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
+        self.o_proj = nn.Linear(self.num_heads * self.head_dim, config.hidden_size, bias=False)
+        self.o_proj.BANANAMIND_SCALE_INIT = True
+        self.q_norm = BananaMind2PicoRMSNorm(self.head_dim, config.rms_norm_eps)
+        self.k_norm = BananaMind2PicoRMSNorm(self.head_dim, config.rms_norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        cosine: torch.Tensor,
+        sine: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        past_key_values: Cache | None = None,
+    ) -> torch.Tensor:
+        batch_size, query_length, _ = hidden_states.shape
+        query = self.q_proj(hidden_states).view(
+            batch_size, query_length, self.num_heads, self.head_dim
+        ).transpose(1, 2)
+        key = self.k_proj(hidden_states).view(
+            batch_size, query_length, self.num_kv_heads, self.head_dim
+        ).transpose(1, 2)
+        value = self.v_proj(hidden_states).view(
+            batch_size, query_length, self.num_kv_heads, self.head_dim
+        ).transpose(1, 2)
+        query = self.q_norm(query)
+        key = self.k_norm(key)
+        query, key = _apply_rope(query, key, cosine, sine)
+
+        past_length = 0
+        if past_key_values is not None:
+            past_length = past_key_values.get_seq_length(self.layer_idx)
+            key, value = past_key_values.update(key, value, self.layer_idx)
+
+        key_length = key.size(-2)
+        key = key.repeat_interleave(self.num_kv_groups, dim=1)
+        value = value.repeat_interleave(self.num_kv_groups, dim=1)
+        is_causal = query_length > 1 and past_length == 0 and attention_mask is None
+        sdpa_mask = None
+        if not is_causal and query_length > 1:
+            query_positions = past_length + torch.arange(query_length, device=query.device)
+            key_positions = torch.arange(key_length, device=query.device)
+            sdpa_mask = (key_positions[None, :] <= query_positions[:, None])[None, None]
+        if attention_mask is not None:
+            key_padding = attention_mask.to(torch.bool)
+            if key_padding.size(-1) < key_length:
+                key_padding = F.pad(
+                    key_padding,
+                    (key_length - key_padding.size(-1), 0),
+                    value=True,
+                )
+            else:
+                key_padding = key_padding[:, -key_length:]
+            key_padding = key_padding[:, None, None, :]
+            sdpa_mask = key_padding if sdpa_mask is None else sdpa_mask & key_padding
+            is_causal = False
+
+        output = F.scaled_dot_product_attention(
+            query,
+            key,
+            value,
+            attn_mask=sdpa_mask,
+            is_causal=is_causal,
+        )
+        output = output.transpose(1, 2).contiguous().view(
+            batch_size, query_length, self.num_heads * self.head_dim
+        )
+        return self.o_proj(output)
+
+
+class BananaMind2PicoRefreshGate(nn.Module):
+    def __init__(self, config: BananaMind2PicoConfig, layer_idx: int):
+        super().__init__()
+        hidden = config.hidden_size
+        self.layer_idx = layer_idx
+        self.kernel_size = config.refresh_kernel_size
+        self.attention_norm = BananaMind2PicoRMSNorm(hidden, config.rms_norm_eps)
+        self.embedding_norm = BananaMind2PicoRMSNorm(hidden, config.rms_norm_eps)
+        self.output_norm = BananaMind2PicoRMSNorm(hidden, config.rms_norm_eps)
+        self.gate_proj = nn.Linear(hidden, hidden, bias=False)
+        self.value_proj = nn.Linear(hidden, hidden, bias=False)
+        self.out_proj = nn.Linear(hidden, hidden, bias=False)
+        self.out_proj.BANANAMIND_SCALE_INIT = True
+        # A real 2D parameter keeps the depthwise kernel eligible for stock Muon.
+        self.depthwise_kernel = nn.Parameter(torch.empty(hidden, self.kernel_size))
+        nn.init.normal_(self.depthwise_kernel, mean=0.0, std=config.initializer_range)
+        self.alpha = nn.Parameter(torch.tensor(float(config.refresh_alpha_init)))
+
+    def _causal_depthwise_conv(
+        self,
+        attention_signal: torch.Tensor,
+        past_key_values: BananaMind2PicoCache | None,
+    ) -> torch.Tensor:
+        signal = attention_signal.transpose(1, 2)
+        history_size = self.kernel_size - 1
+        if past_key_values is None:
+            convolution_input = F.pad(signal, (history_size, 0))
+        else:
+            convolution_input = past_key_values.refresh_input(
+                self.layer_idx,
+                signal,
+                history_size,
+            )
+        convolved = F.conv1d(
+            convolution_input,
+            self.depthwise_kernel.unsqueeze(1),
+            groups=signal.size(1),
+        )
+        return convolved.transpose(1, 2)
+
+    def forward(
+        self,
+        attention_output: torch.Tensor,
+        original_embedding: torch.Tensor,
+        past_key_values: BananaMind2PicoCache | None,
+    ) -> torch.Tensor:
+        attention_signal = self.attention_norm(attention_output.detach())
+        embedding_value = self.embedding_norm(original_embedding)
+        gate = self.gate_proj(attention_signal) + self._causal_depthwise_conv(
+            attention_signal,
+            past_key_values,
+        )
+        value = self.value_proj(embedding_value)
+        refreshed = self.output_norm(self.out_proj(F.silu(gate) * value))
+        return self.alpha * refreshed
+
+
+class BananaMind2PicoMLP(nn.Module):
+    def __init__(self, config: BananaMind2PicoConfig):
+        super().__init__()
+        self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
+        self.down_proj.BANANAMIND_SCALE_INIT = True
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.down_proj(
+            F.silu(self.gate_proj(hidden_states)) * self.up_proj(hidden_states)
+        )
+
+
+class BananaMind2PicoBlock(nn.Module):
+    def __init__(self, config: BananaMind2PicoConfig, layer_idx: int):
+        super().__init__()
+        self.input_norm = BananaMind2PicoRMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.attention = BananaMind2PicoAttention(config, layer_idx)
+        self.refresh = (
+            BananaMind2PicoRefreshGate(config, layer_idx)
+            if layer_idx in config.refresh_layers
+            else None
+        )
+        self.post_attention_norm = BananaMind2PicoRMSNorm(
+            config.hidden_size,
+            config.rms_norm_eps,
+        )
+        self.mlp = BananaMind2PicoMLP(config)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        original_embedding: torch.Tensor,
+        cosine: torch.Tensor,
+        sine: torch.Tensor,
+        attention_mask: torch.Tensor | None,
+        past_key_values: BananaMind2PicoCache | None,
+    ) -> torch.Tensor:
+        attention_output = self.attention(
+            self.input_norm(hidden_states),
+            cosine,
+            sine,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+        )
+        hidden_states = hidden_states + attention_output
+        if self.refresh is not None:
+            hidden_states = hidden_states + self.refresh(
+                attention_output,
+                original_embedding,
+                past_key_values,
+            )
+        return hidden_states + self.mlp(self.post_attention_norm(hidden_states))
+
+
+class BananaMind2PicoPreTrainedModel(PreTrainedModel):
+    config_class = BananaMind2PicoConfig
+    base_model_prefix = "transformer"
+    supports_gradient_checkpointing = False
+    _no_split_modules = ["BananaMind2PicoBlock"]
+    _supports_sdpa = True
+    _supports_cache_class = True
+
+    def _init_weights(self, module: nn.Module):
+        std = self.config.initializer_range
+        if hasattr(module, "BANANAMIND_SCALE_INIT"):
+            std *= (2 * self.config.num_hidden_layers) ** -0.5
+        if isinstance(module, nn.Linear):
+            nn.init.normal_(module.weight, mean=0.0, std=std)
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(module.weight, mean=0.0, std=std)
+
+
+class BananaMind2PicoForCausalLM(BananaMind2PicoPreTrainedModel, GenerationMixin):
+    _tied_weights_keys = {"lm_head.weight": "transformer.wte.weight"}
+
+    @classmethod
+    def _supports_default_dynamic_cache(cls) -> bool:
+        return False
+
+    def __init__(self, config: BananaMind2PicoConfig):
+        super().__init__(config)
+        self.transformer = nn.ModuleDict(
+            {
+                "wte": nn.Embedding(config.vocab_size, config.hidden_size),
+                "h": nn.ModuleList(
+                    [
+                        BananaMind2PicoBlock(config, index)
+                        for index in range(config.num_hidden_layers)
+                    ]
+                ),
+                "ln_f": BananaMind2PicoRMSNorm(
+                    config.hidden_size,
+                    config.rms_norm_eps,
+                ),
+            }
+        )
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.embedding_scale = math.sqrt(config.hidden_size)
+        self.post_init()
+        if config.tie_word_embeddings:
+            self.tie_weights()
+
+    def get_input_embeddings(self):
+        return self.transformer["wte"]
+
+    def set_input_embeddings(self, value):
+        self.transformer["wte"] = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, value):
+        self.lm_head = value
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        use_cache: Optional[bool] = None,
+        **kwargs,
+    ) -> CausalLMOutputWithPast:
+        if use_cache is None:
+            use_cache = self.config.use_cache and labels is None
+        if use_cache and past_key_values is None:
+            past_key_values = BananaMind2PicoCache(self.config)
+        if use_cache and not isinstance(past_key_values, BananaMind2PicoCache):
+            raise TypeError("BananaMind 2 Pico requires its cache for refresh state")
+        if not use_cache:
+            past_key_values = None
+
+        past_length = past_key_values.get_seq_length() if past_key_values is not None else 0
+        sequence_length = input_ids.size(1)
+        total_length = past_length + sequence_length
+        if total_length > self.config.max_position_embeddings:
+            raise ValueError(
+                f"Sequence length {total_length} exceeds "
+                f"{self.config.max_position_embeddings}"
+            )
+
+        original_embedding = self.transformer["wte"](input_ids) * self.embedding_scale
+        hidden_states = original_embedding
+        positions = torch.arange(
+            past_length,
+            total_length,
+            dtype=torch.float32,
+            device=input_ids.device,
+        )
+        cosine, sine = _rope_cos_sin(
+            self.config.head_dim,
+            positions,
+            self.config.rope_theta,
+        )
+        for block in self.transformer["h"]:
+            hidden_states = block(
+                hidden_states,
+                original_embedding,
+                cosine,
+                sine,
+                attention_mask,
+                past_key_values,
+            )
+        hidden_states = self.transformer["ln_f"](hidden_states)
+        logits = self.lm_head(hidden_states)
+
+        loss = None
+        if labels is not None:
+            loss = F.cross_entropy(
+                logits[..., :-1, :].float().reshape(-1, logits.size(-1)),
+                labels[..., 1:].reshape(-1),
+                ignore_index=-100,
+            )
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=past_key_values,
+        )
+
+
+BananaMind2PicoForCausalLM.register_for_auto_class("AutoModelForCausalLM")
+
+__all__ = [
+    "BananaMind2PicoCache",
+    "BananaMind2PicoForCausalLM",
+    "BananaMind2PicoPreTrainedModel",
+]

special_tokens_map.json

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

tokenizer.json

@@ -0,0 +1,991 @@
+{
+  "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": {
+    "type": "NFKC"
+  },
+  "pre_tokenizer": {
+    "type": "Sequence",
+    "pretokenizers": [
+      {
+        "type": "Split",
+        "pattern": {
+          "Regex": "\\d"
+        },
+        "behavior": "Isolated",
+        "invert": false
+      },
+      {
+        "type": "ByteLevel",
+        "add_prefix_space": false,
+        "trim_offsets": true,
+        "use_regex": true
+      }
+    ]
+  },
+  "post_processor": {
+    "type": "TemplateProcessing",
+    "single": [
+      {
+        "Sequence": {
+          "id": "A",
+          "type_id": 0
+        }
+      }
+    ],
+    "pair": [
+      {
+        "Sequence": {
+          "id": "A",
+          "type_id": 0
+        }
+      },
+      {
+        "Sequence": {
+          "id": "B",
+          "type_id": 1
+        }
+      }
+    ],
+    "special_tokens": {}
+  },
+  "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,
+      "Ġt": 260,
+      "Ġa": 261,
+      "in": 262,
+      "he": 263,
+      "re": 264,
+      "on": 265,
+      "Ġthe": 266,
+      "er": 267,
+      "Ġs": 268,
+      "at": 269,
+      "Ġo": 270,
+      "en": 271,
+      "Ġc": 272,
+      "Ġw": 273,
+      "es": 274,
+      "nd": 275,
+      "it": 276,
+      "ing": 277,
+      "is": 278,
+      "or": 279,
+      "ou": 280,
+      "al": 281,
+      "Ġp": 282,
+      "ed": 283,
+      "Ġf": 284,
+      "an": 285,
+      "Ġb": 286,
+      "Ġin": 287,
+      "ar": 288,
+      "Ġof": 289,
+      "Ġm": 290,
+      "Ġand": 291,
+      "Ġto": 292,
+      "Ġd": 293,
+      "ion": 294,
+      "ic": 295,
+      "le": 296,
+      "om": 297,
+      "Ġh": 298,
+      "as": 299,
+      "Ġth": 300,
+      "ent": 301,
+      "Ġe": 302,
+      "Ġre": 303,
+      "ct": 304,
+      "il": 305,
+      "ro": 306,
+      "Ġl": 307,
+      "st": 308,
+      "ve": 309,
+      "Ġn": 310,
+      "ly": 311,
+      "et": 312,
+      "Ġy": 313,
+      "Ġg": 314,
+      "id": 315,
+      "im": 316,
+      "ĠT": 317,
+      "se": 318,
+      "ation": 319,
+      "Ġbe": 320,
+      "ol": 321,
+      "ce": 322,
+      "ur": 323,
+      "ut": 324,
+      "ow": 325,
+      "ra": 326,
+      "ot": 327,
+      "ig": 328,
+      "Ġyou": 329,
+      "Ġis": 330,
+      "Ġu": 331,
+      "ĠS": 332,
+      "ĠA": 333,
+      "ĠI": 334,
+      "Ġfor": 335,
+      "ch": 336,
+      "ver": 337,
+      "Ġon": 338,
+      "Ġst": 339,
+      "am": 340,
+      "ay": 341,
+      "ĠC": 342,
+      "Ġthat": 343,
+      "ir": 344,
+      "ith": 345,
+      "el": 346,
+      "Ġhe": 347,
+      "ul": 348,
+      "ad": 349,
+      "âĢ": 350,
+      "Ġcon": 351,
+      "Ġwith": 352,
+      "Ġas": 353,
+      "ter": 354,
+      "if": 355,
+      "Ġwh": 356,
+      "Ġan": 357,
+      "Ġit": 358,
+      "ĠM": 359,
+      "Ġor": 360,
+      "Ġpro": 361,
+      "Ġal": 362,
+      "her": 363,
+      "ers": 364,
+      "em": 365,
+      "od": 366,
+      "ess": 367,
+      "Ġex": 368,
+      "ate": 369,
+      "ĠB": 370,
+      "ag": 371,
+      "ist": 372,
+      "Ġcom": 373,
+      "ri": 374,
+      "ĠP": 375,
+      "us": 376,
+      "ill": 377,
+      "res": 378,
+      "pp": 379,
+      "ab": 380,
+      "ke": 381,
+      "th": 382,
+      "ĠH": 383
+    },
+    "merges": [
+      [
+        "Ġ",
+        "t"
+      ],
+      [
+        "Ġ",
+        "a"
+      ],
+      [
+        "i",
+        "n"
+      ],
+      [
+        "h",
+        "e"
+      ],
+      [
+        "r",
+        "e"
+      ],
+      [
+        "o",
+        "n"
+      ],
+      [
+        "Ġt",
+        "he"
+      ],
+      [
+        "e",
+        "r"
+      ],
+      [
+        "Ġ",
+        "s"
+      ],
+      [
+        "a",
+        "t"
+      ],
+      [
+        "Ġ",
+        "o"
+      ],
+      [
+        "e",
+        "n"
+      ],
+      [
+        "Ġ",
+        "c"
+      ],
+      [
+        "Ġ",
+        "w"
+      ],
+      [
+        "e",
+        "s"
+      ],
+      [
+        "n",
+        "d"
+      ],
+      [
+        "i",
+        "t"
+      ],
+      [
+        "in",
+        "g"
+      ],
+      [
+        "i",
+        "s"
+      ],
+      [
+        "o",
+        "r"
+      ],
+      [
+        "o",
+        "u"
+      ],
+      [
+        "a",
+        "l"
+      ],
+      [
+        "Ġ",
+        "p"
+      ],
+      [
+        "e",
+        "d"
+      ],
+      [
+        "Ġ",
+        "f"
+      ],
+      [
+        "a",
+        "n"
+      ],
+      [
+        "Ġ",
+        "b"
+      ],
+      [
+        "Ġ",
+        "in"
+      ],
+      [
+        "a",
+        "r"
+      ],
+      [
+        "Ġo",
+        "f"
+      ],
+      [
+        "Ġ",
+        "m"
+      ],
+      [
+        "Ġa",
+        "nd"
+      ],
+      [
+        "Ġt",
+        "o"
+      ],
+      [
+        "Ġ",
+        "d"
+      ],
+      [
+        "i",
+        "on"
+      ],
+      [
+        "i",
+        "c"
+      ],
+      [
+        "l",
+        "e"
+      ],
+      [
+        "o",
+        "m"
+      ],
+      [
+        "Ġ",
+        "h"
+      ],
+      [
+        "a",
+        "s"
+      ],
+      [
+        "Ġt",
+        "h"
+      ],
+      [
+        "en",
+        "t"
+      ],
+      [
+        "Ġ",
+        "e"
+      ],
+      [
+        "Ġ",
+        "re"
+      ],
+      [
+        "c",
+        "t"
+      ],
+      [
+        "i",
+        "l"
+      ],
+      [
+        "r",
+        "o"
+      ],
+      [
+        "Ġ",
+        "l"
+      ],
+      [
+        "s",
+        "t"
+      ],
+      [
+        "v",
+        "e"
+      ],
+      [
+        "Ġ",
+        "n"
+      ],
+      [
+        "l",
+        "y"
+      ],
+      [
+        "e",
+        "t"
+      ],
+      [
+        "Ġ",
+        "y"
+      ],
+      [
+        "Ġ",
+        "g"
+      ],
+      [
+        "i",
+        "d"
+      ],
+      [
+        "i",
+        "m"
+      ],
+      [
+        "Ġ",
+        "T"
+      ],
+      [
+        "s",
+        "e"
+      ],
+      [
+        "at",
+        "ion"
+      ],
+      [
+        "Ġb",
+        "e"
+      ],
+      [
+        "o",
+        "l"
+      ],
+      [
+        "c",
+        "e"
+      ],
+      [
+        "u",
+        "r"
+      ],
+      [
+        "u",
+        "t"
+      ],
+      [
+        "o",
+        "w"
+      ],
+      [
+        "r",
+        "a"
+      ],
+      [
+        "o",
+        "t"
+      ],
+      [
+        "i",
+        "g"
+      ],
+      [
+        "Ġy",
+        "ou"
+      ],
+      [
+        "Ġ",
+        "is"
+      ],
+      [
+        "Ġ",
+        "u"
+      ],
+      [
+        "Ġ",
+        "S"
+      ],
+      [
+        "Ġ",
+        "A"
+      ],
+      [
+        "Ġ",
+        "I"
+      ],
+      [
+        "Ġf",
+        "or"
+      ],
+      [
+        "c",
+        "h"
+      ],
+      [
+        "v",
+        "er"
+      ],
+      [
+        "Ġ",
+        "on"
+      ],
+      [
+        "Ġs",
+        "t"
+      ],
+      [
+        "a",
+        "m"
+      ],
+      [
+        "a",
+        "y"
+      ],
+      [
+        "Ġ",
+        "C"
+      ],
+      [
+        "Ġth",
+        "at"
+      ],
+      [
+        "i",
+        "r"
+      ],
+      [
+        "it",
+        "h"
+      ],
+      [
+        "e",
+        "l"
+      ],
+      [
+        "Ġ",
+        "he"
+      ],
+      [
+        "u",
+        "l"
+      ],
+      [
+        "a",
+        "d"
+      ],
+      [
+        "â",
+        "Ģ"
+      ],
+      [
+        "Ġc",
+        "on"
+      ],
+      [
+        "Ġw",
+        "ith"
+      ],
+      [
+        "Ġa",
+        "s"
+      ],
+      [
+        "t",
+        "er"
+      ],
+      [
+        "i",
+        "f"
+      ],
+      [
+        "Ġw",
+        "h"
+      ],
+      [
+        "Ġa",
+        "n"
+      ],
+      [
+        "Ġ",
+        "it"
+      ],
+      [
+        "Ġ",
+        "M"
+      ],
+      [
+        "Ġo",
+        "r"
+      ],
+      [
+        "Ġp",
+        "ro"
+      ],
+      [
+        "Ġa",
+        "l"
+      ],
+      [
+        "he",
+        "r"
+      ],
+      [
+        "er",
+        "s"
+      ],
+      [
+        "e",
+        "m"
+      ],
+      [
+        "o",
+        "d"
+      ],
+      [
+        "es",
+        "s"
+      ],
+      [
+        "Ġe",
+        "x"
+      ],
+      [
+        "at",
+        "e"
+      ],
+      [
+        "Ġ",
+        "B"
+      ],
+      [
+        "a",
+        "g"
+      ],
+      [
+        "is",
+        "t"
+      ],
+      [
+        "Ġc",
+        "om"
+      ],
+      [
+        "r",
+        "i"
+      ],
+      [
+        "Ġ",
+        "P"
+      ],
+      [
+        "u",
+        "s"
+      ],
+      [
+        "il",
+        "l"
+      ],
+      [
+        "re",
+        "s"
+      ],
+      [
+        "p",
+        "p"
+      ],
+      [
+        "a",
+        "b"
+      ],
+      [
+        "k",
+        "e"
+      ],
+      [
+        "t",
+        "h"
+      ],
+      [
+        "Ġ",
+        "H"
+      ]
+    ]
+  }
+}

tokenizer_config.json

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

tokenizer_training_manifest.json

@@ -0,0 +1,29 @@
+{
+  "format": "bananamind2_pico_tokenizer_v1",
+  "stage": "tokenizer_complete",
+  "vocab_size": 384,
+  "model_max_length": 8192,
+  "target_gib": 1.0,
+  "sources": [
+    {
+      "name": "fineweb_hq",
+      "dataset_id": "epfml/FineWeb-HQ",
+      "config": null,
+      "text_field": "text",
+      "byte_share": 0.75,
+      "revision": "e58199cdd52438d94405df1a4d8630cc5f13bf84",
+      "target_bytes": 805306368,
+      "consumed_bytes": 805307468
+    },
+    {
+      "name": "cosmopedia_v2",
+      "dataset_id": "HuggingFaceTB/smollm-corpus",
+      "config": "cosmopedia-v2",
+      "text_field": "text",
+      "byte_share": 0.25,
+      "revision": "3ba9d605774198c5868892d7a8deda78031a781f",
+      "target_bytes": 268435456,
+      "consumed_bytes": 268438475
+    }
+  ]
+}