modeling_rost.py

"""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}