RoST-1B-Instruct-v2: frozen SFT-v2 arm C, step 365

Exported from checkpoint model_000365.pt (sha256 397dcd5264e2e0dafadd20da3726d8105cafb0becf43ced607d1cd721b0df121), verified in fp32 against the source checkpoint with a maximum logit difference of 0.0, then cast to bfloat16. Bitwise identical on all 175 tensors to the export that produced the published benchmark numbers.

config.json

@@ -0,0 +1,19 @@
+{
+  "vocab_size": 32768,
+  "n_layer": 24,
+  "n_head": 12,
+  "n_kv_head": 12,
+  "n_embd": 1536,
+  "sequence_len": 4096,
+  "rope_base": 100000,
+  "window_pattern": "SSSL",
+  "model_type": "rost",
+  "architectures": [
+    "RostForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_rost.RostConfig",
+    "AutoModelForCausalLM": "modeling_rost.RostForCausalLM"
+  },
+  "torch_dtype": "bfloat16"
+}

configuration_rost.py

@@ -0,0 +1,83 @@
+"""HuggingFace config for RoST.
+
+Ships inside the published model repository and runs on the downloader's
+machine, so it must not import anything from `nanochat`.
+
+Field names mirror `nanochat.gpt.GPTConfig` exactly rather than being renamed
+to Llama's vocabulary. A rename would need a mapping table that nothing checks,
+and a silently wrong mapping produces a model that loads and computes the wrong
+thing -- the one failure mode this whole export has to avoid.
+"""
+from transformers.configuration_utils import PretrainedConfig
+
+
+class RostConfig(PretrainedConfig):
+    model_type = "rost"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        vocab_size=32768,
+        n_layer=24,
+        n_head=12,
+        n_kv_head=12,
+        n_embd=1536,
+        sequence_len=4096,
+        rope_base=100000,
+        window_pattern="SSSL",
+        pad_vocab_size_to=64,
+        logit_softcap=15.0,
+        attention_scale=1.2,
+        ve_gate_channels=12,
+        smear_gate_channels=24,
+        bos_token_id=None,
+        eos_token_id=None,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.n_layer = n_layer
+        self.n_head = n_head
+        self.n_kv_head = n_kv_head
+        self.n_embd = n_embd
+        self.sequence_len = sequence_len
+        self.rope_base = rope_base
+        self.window_pattern = window_pattern
+        self.pad_vocab_size_to = pad_vocab_size_to
+        # Constants in the training code, carried as config so a checkpoint
+        # trained under different ones cannot be served under these.
+        self.logit_softcap = logit_softcap
+        self.attention_scale = attention_scale
+        self.ve_gate_channels = ve_gate_channels
+        self.smear_gate_channels = smear_gate_channels
+        super().__init__(bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs)
+
+    @property
+    def padded_vocab_size(self):
+        pad = self.pad_vocab_size_to
+        return ((self.vocab_size + pad - 1) // pad) * pad
+
+    @property
+    def head_dim(self):
+        return self.n_embd // self.n_head
+
+    # Aliases so generic HuggingFace code (generation, device maps, pipelines)
+    # finds what it expects without the weights being renamed.
+    @property
+    def hidden_size(self):
+        return self.n_embd
+
+    @property
+    def num_attention_heads(self):
+        return self.n_head
+
+    @property
+    def num_key_value_heads(self):
+        return self.n_kv_head
+
+    @property
+    def num_hidden_layers(self):
+        return self.n_layer
+
+    @property
+    def max_position_embeddings(self):
+        return self.sequence_len

model.safetensors

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

modeling_rost.py

@@ -0,0 +1,325 @@
+"""HuggingFace modelling code for RoST.
+
+Ships inside the published model repository and runs on the downloader's
+machine, so it imports nothing from `nanochat` and uses no FlashAttention-3.
+
+This is a transcription of `nanochat/gpt.py`, not a reimplementation. Parameter
+names, the order of operations and every constant are kept identical, because
+the only thing that makes an export trustworthy is that it computes the same
+function -- `tests/test_hf_export.py` asserts that against the source model.
+
+RoST is not a Llama variant. It carries nine components with no equivalent in
+standard architectures: smear, per-layer resid/x0 lambdas, gated value
+embeddings on alternating layers, backout, QK-norm with double 1.2 scaling,
+relu-squared MLP, parameter-free RMSNorm, logit softcap and a tiled sliding
+window. Each is transcribed below with the reason it exists.
+"""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers.cache_utils import DynamicCache
+from transformers.generation.utils import GenerationMixin
+from transformers.modeling_outputs import CausalLMOutputWithPast
+from transformers.modeling_utils import PreTrainedModel
+
+from .configuration_rost import RostConfig
+
+
+def norm(x):
+    """RMSNorm with NO learnable scale. RoST has no norm parameters at all."""
+    return F.rms_norm(x, (x.size(-1),))
+
+
+def has_ve(layer_idx, n_layer):
+    """Value embeddings sit on alternating layers, last layer always included."""
+    return layer_idx % 2 == (n_layer - 1) % 2
+
+
+def apply_rotary_emb(x, cos, sin):
+    # Rotates by -theta, the transpose of the textbook convention. Only the
+    # relative q/k rotation matters so it is functionally equivalent, but it
+    # must be transcribed as-is or the loaded weights mean something else.
+    d = x.shape[3] // 2
+    x1, x2 = x[..., :d], x[..., d:]
+    y1 = x1 * cos + x2 * sin
+    y2 = x1 * (-sin) + x2 * cos
+    return torch.cat([y1, y2], 3)
+
+
+def compute_window_sizes(config):
+    """Per-layer left-attention span, tiled from `window_pattern`.
+
+    S is a quarter of the context rounded up to 128; L is the full context. The
+    final layer is always L. Mirrors `GPT._compute_window_sizes`.
+    """
+    pattern = config.window_pattern.upper()
+    long_window = config.sequence_len
+    short_window = -(-long_window // 4 // 128) * 128
+    sizes = [long_window if pattern[i % len(pattern)] == "L" else short_window
+             for i in range(config.n_layer)]
+    sizes[-1] = long_window
+    return sizes
+
+
+class RostAttention(nn.Module):
+    def __init__(self, config, layer_idx):
+        super().__init__()
+        self.layer_idx = layer_idx
+        self.n_head = config.n_head
+        self.n_kv_head = config.n_kv_head
+        self.head_dim = config.head_dim
+        self.attention_scale = config.attention_scale
+        self.c_q = nn.Linear(config.n_embd, self.n_head * self.head_dim, bias=False)
+        self.c_k = nn.Linear(config.n_embd, self.n_kv_head * self.head_dim, bias=False)
+        self.c_v = nn.Linear(config.n_embd, self.n_kv_head * self.head_dim, bias=False)
+        self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=False)
+        self.ve_gate_channels = config.ve_gate_channels
+        self.ve_gate = (nn.Linear(self.ve_gate_channels, self.n_kv_head, bias=False)
+                        if has_ve(layer_idx, config.n_layer) else None)
+
+    def forward(self, x, ve, cos, sin, attn_mask, cache, layer_idx):
+        B, T, _ = x.size()
+        q = self.c_q(x).view(B, T, self.n_head, self.head_dim)
+        k = self.c_k(x).view(B, T, self.n_kv_head, self.head_dim)
+        v = self.c_v(x).view(B, T, self.n_kv_head, self.head_dim)
+
+        # Value residual (ResFormer): a per-token, per-kv-head gate in (0, 3)
+        # mixes a learned per-layer value embedding into v.
+        if ve is not None:
+            ve = ve.view(B, T, self.n_kv_head, self.head_dim)
+            gate = 3 * torch.sigmoid(self.ve_gate(x[..., :self.ve_gate_channels]))
+            v = v + gate.unsqueeze(-1) * ve
+
+        q, k = apply_rotary_emb(q, cos, sin), apply_rotary_emb(k, cos, sin)
+        q, k = norm(q), norm(k)  # QK norm
+        # Sharper attention: the 1.2 is applied to BOTH q and k, so the effective
+        # logit scale is 1.44x the usual 1/sqrt(head_dim).
+        q = q * self.attention_scale
+        k = k * self.attention_scale
+
+        # (B, T, H, D) -> (B, H, T, D) for SDPA
+        q = q.transpose(1, 2)
+        k = k.transpose(1, 2)
+        v = v.transpose(1, 2)
+
+        # Append through the cache's own API rather than concatenating tensors
+        # by hand: `generate()` owns the cache object and expects to be the one
+        # tracking its length.
+        if cache is not None:
+            k, v = cache.update(k, v, layer_idx)
+
+        if self.n_kv_head != self.n_head:
+            repeat = self.n_head // self.n_kv_head
+            k = k.repeat_interleave(repeat, dim=1)
+            v = v.repeat_interleave(repeat, dim=1)
+
+        y = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)
+        y = y.transpose(1, 2).contiguous().view(B, T, -1)
+        return self.c_proj(y)
+
+
+class RostMLP(nn.Module):
+    """relu-squared at 4x expansion, not SwiGLU at 8/3x."""
+
+    def __init__(self, config):
+        super().__init__()
+        self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd, bias=False)
+        self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=False)
+
+    def forward(self, x):
+        return self.c_proj(F.relu(self.c_fc(x)).square())
+
+
+class RostBlock(nn.Module):
+    def __init__(self, config, layer_idx):
+        super().__init__()
+        self.attn = RostAttention(config, layer_idx)
+        self.mlp = RostMLP(config)
+
+    def forward(self, x, ve, cos, sin, attn_mask, cache, layer_idx):
+        x = x + self.attn(norm(x), ve, cos, sin, attn_mask, cache, layer_idx)
+        x = x + self.mlp(norm(x))
+        return x
+
+
+class RostCache(DynamicCache):
+    """A KV cache that also carries smear's previous-token embedding.
+
+    Smear mixes the previous token's embedding into the current one. During
+    incremental decoding that embedding is not in `input_ids`, and it is not a
+    key or a value, so there is nowhere in the standard cache to put it. It
+    rides along as an attribute here.
+
+    `generate()` builds its own `DynamicCache` rather than this subclass, so the
+    forward pass reads the attribute defensively with `getattr` and sets it on
+    whatever cache object it was handed. That works because a plain
+    `DynamicCache` accepts attribute assignment -- and it must keep working,
+    because the alternative failure is silent: without the previous embedding
+    every decoded token is smeared against nothing.
+    """
+
+    prev_embedding = None
+
+
+class RostPreTrainedModel(PreTrainedModel):
+    config_class = RostConfig
+    base_model_prefix = "transformer"
+    supports_gradient_checkpointing = False
+    _no_split_modules = ["RostBlock"]
+
+
+class RostForCausalLM(RostPreTrainedModel, GenerationMixin):
+    # GenerationMixin after PreTrainedModel, or `generate` is unavailable from
+    # transformers 4.50 onward.
+    def __init__(self, config):
+        super().__init__(config)
+        padded = config.padded_vocab_size
+        self.transformer = nn.ModuleDict({
+            "wte": nn.Embedding(padded, config.n_embd),
+            "h": nn.ModuleList([RostBlock(config, i) for i in range(config.n_layer)]),
+        })
+        self.lm_head = nn.Linear(config.n_embd, padded, bias=False)
+        # Per-layer scalars from modded-nanogpt: resid_lambdas rescales the
+        # residual stream, x0_lambdas blends the initial embedding back in.
+        self.resid_lambdas = nn.Parameter(torch.ones(config.n_layer))
+        self.x0_lambdas = nn.Parameter(torch.zeros(config.n_layer))
+        # Smear: mixes the previous token's embedding into the current one.
+        self.smear_gate = nn.Linear(config.smear_gate_channels, 1, bias=False)
+        self.smear_lambda = nn.Parameter(torch.zeros(1))
+        # Backout: removes the mid-layer residual before the logit projection.
+        self.backout_lambda = nn.Parameter(0.2 * torch.ones(1))
+        kv_dim = config.n_kv_head * config.head_dim
+        self.value_embeds = nn.ModuleDict({
+            str(i): nn.Embedding(padded, kv_dim)
+            for i in range(config.n_layer) if has_ve(i, config.n_layer)})
+
+        self.window_sizes = compute_window_sizes(config)
+        # Rotary tables are built on first use, not in __init__.
+        #
+        # They are derived from config, so they are absent from the checkpoint.
+        # `from_pretrained` initializes on the meta device and materializes only
+        # tensors the checkpoint supplies, so buffers registered here would stay
+        # meta and the model would return NaN -- silently, and only after a
+        # round trip through disk, which is exactly how a published model breaks
+        # while every in-memory test passes.
+        self._rotary_cache = None
+        self.post_init()
+
+    def _rotary(self, device, dtype, length):
+        cached = self._rotary_cache
+        if (cached is not None and cached[0].device == device
+                and cached[0].dtype == dtype and cached[0].size(1) >= length):
+            return cached
+        head_dim = self.config.head_dim
+        # Table length mirrors nanochat's 10x over-compute, so a sequence longer
+        # than the trained context still has rotations available rather than
+        # tripping an index error at serving time.
+        size = max(length, self.config.sequence_len * 10)
+        channel_range = torch.arange(0, head_dim, 2, dtype=torch.float32, device=device)
+        inv_freq = 1.0 / (self.config.rope_base ** (channel_range / head_dim))
+        t = torch.arange(size, dtype=torch.float32, device=device)
+        freqs = torch.outer(t, inv_freq)
+        cos = freqs.cos()[None, :, None, :].to(dtype)
+        sin = freqs.sin()[None, :, None, :].to(dtype)
+        self._rotary_cache = (cos, sin)
+        return self._rotary_cache
+
+    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 _window_mask(self, window, q_len, kv_len, offset, device):
+        """Causal mask restricted to a left-window, matching FA3's semantics.
+
+        FA3's `window_size=(left, 0)` attends to keys in `[i - left, i]`
+        inclusive. A mask that dropped the `i - left` position, or that used the
+        window as a count rather than a span, would change what 18 of 24 layers
+        can see -- quietly, and only on long inputs.
+        """
+        q_pos = torch.arange(offset, offset + q_len, device=device).unsqueeze(1)
+        k_pos = torch.arange(kv_len, device=device).unsqueeze(0)
+        allowed = (k_pos <= q_pos) & (k_pos >= q_pos - window)
+        return allowed.unsqueeze(0).unsqueeze(0)
+
+    def forward(self, input_ids, attention_mask=None, past_key_values=None,
+                use_cache=None, labels=None, return_dict=True, **kwargs):
+        B, T = input_ids.size()
+        device = input_ids.device
+        use_cache = True if use_cache is None else use_cache
+        if use_cache and past_key_values is None:
+            past_key_values = RostCache()
+
+        # Position of this chunk in the sequence. Read from the cache rather
+        # than tracked separately: `generate()` supplies its own cache object,
+        # and a private counter would silently desynchronise from it.
+        offset = past_key_values.get_seq_length() if past_key_values is not None else 0
+
+        x = self.transformer["wte"](input_ids)
+        cos_table, sin_table = self._rotary(device, x.dtype, offset + T)
+        cos, sin = cos_table[:, offset:offset + T], sin_table[:, offset:offset + T]
+        x = norm(x)
+
+        # Smear. During incremental decoding the previous token's embedding is
+        # not in `input_ids`, so it is carried in the cache. HuggingFace's cache
+        # API has no slot for non-KV state, which is why the cache here is a
+        # plain dict rather than a `Cache` subclass.
+        prev = getattr(past_key_values, "prev_embedding", None)
+        gate_channels = self.config.smear_gate_channels
+        # Stored BEFORE smear is applied, matching nanochat, where
+        # `kv_cache.prev_embedding = x[:, -1:, :]` is assigned on the post-norm
+        # pre-smear activation.
+        new_prev = x[:, -1:, :]
+        if T > 1:
+            # Position 0 is left unsmeared even when a previous embedding
+            # exists. nanochat's prefill branch does the same; carrying `prev`
+            # in here would make a two-call prefill differ from a one-call one.
+            gate = self.smear_lambda.to(x.dtype) * torch.sigmoid(
+                self.smear_gate(x[:, 1:, :gate_channels]))
+            x = torch.cat([x[:, :1], x[:, 1:] + gate * x[:, :-1]], dim=1)
+        elif prev is not None:
+            gate = self.smear_lambda.to(x.dtype) * torch.sigmoid(
+                self.smear_gate(x[:, :, :gate_channels]))
+            x = x + gate * prev
+
+        x0 = x
+        n_layer = self.config.n_layer
+        backout_layer = n_layer // 2
+        x_backout = None
+        for i, block in enumerate(self.transformer["h"]):
+            x = self.resid_lambdas[i] * x + self.x0_lambdas[i] * x0
+            ve = (self.value_embeds[str(i)](input_ids).to(x.dtype)
+                  if str(i) in self.value_embeds else None)
+            mask = self._window_mask(self.window_sizes[i], T, offset + T, offset, device)
+            x = block(x, ve, cos, sin, mask, past_key_values, i)
+            if i == backout_layer:
+                x_backout = x
+        if x_backout is not None:
+            x = x - self.backout_lambda.to(x.dtype) * x_backout
+        x = norm(x)
+
+        logits = self.lm_head(x)[..., :self.config.vocab_size].float()
+        softcap = self.config.logit_softcap
+        logits = softcap * torch.tanh(logits / softcap)
+
+        loss = None
+        if labels is not None:
+            loss = F.cross_entropy(logits[:, :-1].reshape(-1, logits.size(-1)),
+                                   labels[:, 1:].reshape(-1), ignore_index=-1)
+
+        if past_key_values is not None:
+            past_key_values.prev_embedding = new_prev
+        return CausalLMOutputWithPast(loss=loss, logits=logits,
+                                      past_key_values=past_key_values if use_cache else None)
+
+    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, **kwargs):
+        # Feed only the new tokens once the cache holds the prefix.
+        if past_key_values is not None and past_key_values.get_seq_length() > 0:
+            input_ids = input_ids[:, past_key_values.get_seq_length():]
+        return {"input_ids": input_ids, "past_key_values": past_key_values, "use_cache": True}

tokenizer_config.json

@@ -0,0 +1,8 @@
+{
+  "tokenizer_class": "PreTrainedTokenizerFast",
+  "bos_token": "<|bos|>",
+  "eos_token": "<|assistant_end|>",
+  "model_max_length": 4096,
+  "clean_up_tokenization_spaces": false,
+  "chat_template": "{{- bos_token -}}{%- set ns = namespace(system='') -%}{%- for message in messages -%}{%- if message['role'] == 'system' -%}{%- set ns.system = message['content'] -%}{%- elif message['role'] == 'user' -%}{{- '<|user_start|>' -}}{%- if ns.system -%}{{- ns.system + '\n\n' -}}{%- set ns.system = '' -%}{%- endif -%}{{- message['content'] + '<|user_end|>' -}}{%- elif message['role'] == 'assistant' -%}{{- '<|assistant_start|>' + message['content'] + '<|assistant_end|>' -}}{%- endif -%}{%- endfor -%}{%- if add_generation_prompt -%}{{- '<|assistant_start|>' -}}{%- endif -%}"
+}