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