chore: restore custom modeling

modeling_nanochat.py

@@ -0,0 +1,365 @@
+"""
+Hugging Face-compatible nanochat Transformer implementation.
+
+This file mirrors the architecture used during training (RoPE, RMSNorm,
+multi-query attention, relu^2 MLP, untied embeddings, logits softcap) while
+presenting the familiar `PreTrainedModel` interface so that checkpoints can be
+served directly from the Hugging Face Hub.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers.modeling_outputs import CausalLMOutputWithPast
+from transformers.modeling_utils import PreTrainedModel
+from transformers.utils import logging
+from transformers import AutoConfig, AutoModelForCausalLM
+
+logger = logging.get_logger(__name__)
+
+
+class NanoChatConfig(PretrainedConfig):
+    model_type = "nanochat"
+
+    def __init__(
+        self,
+        vocab_size=65536,
+        sequence_len=2048,
+        n_layer=20,
+        n_head=10,
+        n_kv_head=10,
+        n_embd=1280,
+        rotary_dim=None,
+        activation_function="relu_squared",
+        use_rope=True,
+        use_qk_norm=True,
+        tie_word_embeddings=False,
+        softcap=15.0,
+        bos_token_id=1,
+        eos_token_id=1,
+        pad_token_id=None,
+        **kwargs,
+    ):
+        super().__init__(
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            pad_token_id=pad_token_id,
+            **kwargs,
+        )
+        self.vocab_size = vocab_size
+        self.sequence_len = sequence_len
+        self.n_layer = n_layer
+        self.n_head = n_head
+        self.n_kv_head = n_kv_head
+        self.n_embd = n_embd
+        self.rotary_dim = rotary_dim or (n_embd // n_head)
+        self.activation_function = activation_function
+        self.use_rope = use_rope
+        self.use_qk_norm = use_qk_norm
+        self.tie_word_embeddings = tie_word_embeddings
+        self.softcap = softcap
+        
+        # Aliases for transformers compatibility
+        self.num_hidden_layers = n_layer
+        self.hidden_size = n_embd
+        self.num_attention_heads = n_head
+        self.num_key_value_heads = n_kv_head
+
+
+def rms_norm(x: Tensor) -> Tensor:
+    return F.rms_norm(x, (x.size(-1),))
+
+
+def relu_squared(x: Tensor) -> Tensor:
+    return F.relu(x) ** 2
+
+
+def rotate_half(x: Tensor) -> Tensor:
+    x1, x2 = x.chunk(2, dim=-1)
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_emb(q: Tensor, k: Tensor, cos: Tensor, sin: Tensor) -> Tuple[Tensor, Tensor]:
+    q = (q * cos) + (rotate_half(q) * sin)
+    k = (k * cos) + (rotate_half(k) * sin)
+    return q, k
+
+
+def repeat_kv(x: Tensor, n_rep: int) -> Tensor:
+    if n_rep == 1:
+        return x
+    b, n_kv_heads, seq_len, head_dim = x.shape
+    x = x[:, :, None, :, :].expand(b, n_kv_heads, n_rep, seq_len, head_dim)
+    return x.reshape(b, n_kv_heads * n_rep, seq_len, head_dim)
+
+
+class NanoChatAttention(nn.Module):
+    def __init__(self, config: NanoChatConfig):
+        super().__init__()
+        self.config = config
+        self.n_head = config.n_head
+        self.n_kv_head = config.n_kv_head
+        self.head_dim = config.n_embd // config.n_head
+        if config.n_embd % config.n_head != 0:
+            raise ValueError("Embedding dimension must be divisible by number of heads")
+
+        self.q_proj = nn.Linear(config.n_embd, self.n_head * self.head_dim, bias=False)
+        self.k_proj = nn.Linear(config.n_embd, self.n_kv_head * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(config.n_embd, self.n_kv_head * self.head_dim, bias=False)
+        self.out_proj = nn.Linear(config.n_embd, config.n_embd, bias=False)
+
+    def forward(
+        self,
+        hidden_states: Tensor,
+        cos: Tensor,
+        sin: Tensor,
+        past_key_value: Optional[Tuple[Tensor, Tensor]] = None,
+        use_cache: bool = False,
+    ) -> Tuple[Tensor, Optional[Tuple[Tensor, Tensor]]]:
+        bsz, q_len, _ = hidden_states.shape
+
+        query = self.q_proj(hidden_states)
+        key = self.k_proj(hidden_states)
+        value = self.v_proj(hidden_states)
+
+        query = query.view(bsz, q_len, self.n_head, self.head_dim).transpose(1, 2)
+        key = key.view(bsz, q_len, self.n_kv_head, self.head_dim).transpose(1, 2)
+        value = value.view(bsz, q_len, self.n_kv_head, self.head_dim).transpose(1, 2)
+
+        query, key = apply_rotary_emb(query, key, cos, sin)
+        if self.config.use_qk_norm:
+            query = rms_norm(query)
+            key = rms_norm(key)
+
+        if past_key_value is not None:
+            past_k, past_v = past_key_value
+            if past_k is not None and past_v is not None:
+                key = torch.cat([past_k, key], dim=2)
+                value = torch.cat([past_v, value], dim=2)
+
+        present = (key, value) if use_cache else None
+
+        key_for_scores = repeat_kv(key, self.n_head // self.n_kv_head)
+        value_for_scores = repeat_kv(value, self.n_head // self.n_kv_head)
+
+        attn_scores = torch.matmul(query, key_for_scores.transpose(-1, -2)) / math.sqrt(self.head_dim)
+        attn_scores = attn_scores.to(torch.float32)
+
+        # causal mask that accounts for the prefix introduced by past key values
+        if attn_scores.size(-1) != q_len:
+            total_k = attn_scores.size(-1)
+            past_len = total_k - q_len
+            mask = torch.arange(total_k, device=attn_scores.device)
+            causal = mask.unsqueeze(0) <= (mask.new_tensor(past_len) + torch.arange(q_len, device=mask.device).unsqueeze(1))
+            attn_scores = attn_scores.masked_fill(~causal, torch.finfo(attn_scores.dtype).min)
+        else:
+            mask = torch.triu(torch.ones((q_len, q_len), device=attn_scores.device, dtype=torch.bool), diagonal=1)
+            attn_scores = attn_scores.masked_fill(mask, torch.finfo(attn_scores.dtype).min)
+
+        attn_weights = F.softmax(attn_scores, dim=-1, dtype=torch.float32)
+        attn_output = torch.matmul(attn_weights, value_for_scores).to(value_for_scores.dtype)
+
+        attn_output = attn_output.transpose(1, 2).contiguous().view(bsz, q_len, -1)
+        attn_output = self.out_proj(attn_output)
+
+        return attn_output, present
+
+
+class NanoChatMLP(nn.Module):
+    def __init__(self, config: NanoChatConfig):
+        super().__init__()
+        hidden_dim = config.n_embd * 4
+        self.fc = nn.Linear(config.n_embd, hidden_dim, bias=False)
+        self.proj = nn.Linear(hidden_dim, config.n_embd, bias=False)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.proj(relu_squared(self.fc(x)))
+
+
+class NanoChatBlock(nn.Module):
+    def __init__(self, config: NanoChatConfig):
+        super().__init__()
+        self.attn = NanoChatAttention(config)
+        self.mlp = NanoChatMLP(config)
+
+    def forward(
+        self,
+        x: Tensor,
+        cos: Tensor,
+        sin: Tensor,
+        past_key_value: Optional[Tuple[Tensor, Tensor]] = None,
+        use_cache: bool = False,
+    ) -> Tuple[Tensor, Optional[Tuple[Tensor, Tensor]]]:
+        residual = x
+        attn_input = rms_norm(x)
+        attn_output, present = self.attn(attn_input, cos, sin, past_key_value, use_cache)
+        x = residual + attn_output
+        mlp_input = rms_norm(x)
+        x = x + self.mlp(mlp_input)
+        return x, present
+
+
+class NanoChatModel(nn.Module):
+    def __init__(self, config: NanoChatConfig):
+        super().__init__()
+        self.config = config
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.n_embd)
+        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
+        self.blocks = nn.ModuleList([NanoChatBlock(config) for _ in range(config.n_layer)])
+
+        self.softcap = config.softcap
+        self._rope_cache: Optional[Tuple[Tensor, Tensor]] = None
+        self._rope_cache_length = 0
+
+    def _build_rope_cache(self, seq_len: int, device: torch.device, dtype: torch.dtype) -> Tuple[Tensor, Tensor]:
+        if self._rope_cache is not None and self._rope_cache_length >= seq_len and self._rope_cache[0].device == device:
+            return self._rope_cache
+
+        head_dim = self.config.n_embd // self.config.n_head
+        theta = 10000.0 ** (-torch.arange(0, head_dim, 2, device=device, dtype=torch.float32) / head_dim)
+        position_ids = torch.arange(seq_len, device=device, dtype=torch.float32)
+        freqs = torch.einsum("i,j->ij", position_ids, theta)
+        cos = freqs.cos()[None, None, :, :]
+        sin = freqs.sin()[None, None, :, :]
+        # Expand to full head_dim (from head_dim/2 to head_dim)
+        cos = torch.repeat_interleave(cos, repeats=2, dim=-1)
+        sin = torch.repeat_interleave(sin, repeats=2, dim=-1)
+        cos = cos.to(dtype=dtype)
+        sin = sin.to(dtype=dtype)
+
+        self._rope_cache = (cos, sin)
+        self._rope_cache_length = seq_len
+        return cos, sin
+
+    def forward(
+        self,
+        input_ids: Tensor,
+        past_key_values: Optional[Tuple[Tuple[Tensor, Tensor], ...]] = None,
+        attention_mask: Optional[Tensor] = None,
+        labels: Optional[Tensor] = None,
+        use_cache: bool = False,
+    ) -> Tuple[Tensor, Optional[Tuple[Tuple[Tensor, Tensor], ...]]]:
+        del attention_mask  # attention masking is handled implicitly via causal masking
+        bsz, seq_len = input_ids.shape
+        device = input_ids.device
+        dtype = self.embed_tokens.weight.dtype
+
+        inputs_embeds = self.embed_tokens(input_ids)
+        x = inputs_embeds
+
+        past_key_values = past_key_values or tuple([None] * len(self.blocks))
+        # Handle DynamicCache which may have (None, None) tuples
+        past_length = 0
+        if past_key_values and past_key_values[0] is not None:
+            if past_key_values[0][0] is not None:
+                past_length = past_key_values[0][0].size(2)
+
+        cos_full, sin_full = self._build_rope_cache(seq_len + past_length, device, dtype)
+        cos = cos_full[:, :, past_length:, :]
+        sin = sin_full[:, :, past_length:, :]
+        new_past_key_values = [] if use_cache else None
+
+        for layer, block in enumerate(self.blocks):
+            past = past_key_values[layer] if past_key_values[layer] is not None else None
+            x, present = block(x, cos, sin, past, use_cache)
+            if use_cache:
+                new_past_key_values.append(present)
+
+        x = rms_norm(x)
+        logits = self.lm_head(x)
+
+        if self.softcap is not None and self.softcap > 0:
+            logits = self.softcap * torch.tanh(logits / self.softcap)
+
+        loss = None
+        if labels is not None:
+            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), labels.view(-1), ignore_index=-1)
+
+        return logits, loss, tuple(new_past_key_values) if use_cache else None
+
+
+class NanoChatForCausalLM(PreTrainedModel):
+    config_class = NanoChatConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = False
+
+    def __init__(self, config: NanoChatConfig):
+        super().__init__(config)
+        self.model = NanoChatModel(config)
+        if config.tie_word_embeddings:
+            self.tie_weights()
+
+    def get_input_embeddings(self) -> nn.Embedding:
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value: nn.Embedding) -> None:
+        self.model.embed_tokens = value
+
+    def get_output_embeddings(self) -> nn.Linear:
+        return self.model.lm_head
+
+    def set_output_embeddings(self, new_embeddings: nn.Linear) -> None:
+        self.model.lm_head = new_embeddings
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids: Tensor,
+        past_key_values: Optional[Tuple[Tuple[Tensor, Tensor], ...]] = None,
+        **kwargs,
+    ):
+        if past_key_values:
+            input_ids = input_ids[:, -1:]
+        return {"input_ids": input_ids, "past_key_values": past_key_values, "use_cache": kwargs.get("use_cache", True)}
+
+    def _reorder_cache(self, past_key_values, beam_idx):
+        reordered = []
+        for layer_past in past_key_values:
+            reordered.append(
+                (
+                    layer_past[0].index_select(0, beam_idx),
+                    layer_past[1].index_select(0, beam_idx),
+                )
+            )
+        return tuple(reordered)
+
+    def forward(
+        self,
+        input_ids: Tensor,
+        attention_mask: Optional[Tensor] = None,
+        past_key_values: Optional[Tuple[Tuple[Tensor, Tensor], ...]] = None,
+        labels: Optional[Tensor] = None,
+        use_cache: bool = False,
+        **kwargs,
+    ) -> CausalLMOutputWithPast:
+        logits, loss, new_past = self.model(
+            input_ids=input_ids,
+            past_key_values=past_key_values,
+            attention_mask=attention_mask,
+            labels=labels,
+            use_cache=use_cache,
+        )
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=new_past,
+        )
+
+
+try:
+    AutoConfig.register("nanochat", NanoChatConfig)
+except ValueError:
+    # Transformers build already provides this registration (e.g., nanochat branch); reuse it.
+    pass
+
+try:
+    AutoModelForCausalLM.register(NanoChatConfig, NanoChatForCausalLM)
+except ValueError:
+    pass