Upload folder using huggingface_hub

__init__.py

@@ -0,0 +1,5 @@
+from .configuration_nanogpt import NanoGPTConfig
+from .modeling_nanogpt import NanoGPTModel
+from .tokenizer_nanogpt import NanoGPTTokenizer
+
+

config.json

@@ -0,0 +1,17 @@
+{
+    "model_type": "nanogpt",
+    "architectures": [
+        "NanoGPTModel"
+    ],
+    "auto_map": {
+        "AutoConfig": "configuration_nanogpt.NanoGPTConfig",
+        "AutoModel": "modeling_nanogpt.NanoGPTModel",
+        "AutoTokenizer": "tokenizer_nanogpt.NanoGPTTokenizer"
+    },
+    "sequence_len": 2048,
+    "vocab_size": 65536,
+    "n_layer": 20,
+    "n_head": 10,
+    "n_kv_head": 10,
+    "n_embd": 1280
+}

configuration_nanogpt.py

@@ -0,0 +1,26 @@
+from transformers import PretrainedConfig
+
+
+class NanoGPTConfig(PretrainedConfig):
+    model_type = "nanogpt"
+
+    def __init__(
+        self,
+        sequence_len: int = 1024,
+        vocab_size: int = 50304,
+        n_layer: int = 12,
+        n_head: int = 6,
+        n_kv_head: int = 6,
+        n_embd: int = 768,
+        **kwargs,
+    ):
+        self.sequence_len = sequence_len
+        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
+        super().__init__(**kwargs)
+
+
+

modeling_nanogpt.py

@@ -0,0 +1,174 @@
+import math
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PreTrainedModel
+
+from .configuration_nanogpt import NanoGPTConfig
+
+
+def _rms_norm(x: torch.Tensor) -> torch.Tensor:
+    return F.rms_norm(x, (x.size(-1),))
+
+
+def _apply_rotary_emb(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+    assert x.ndim == 4
+    d = x.shape[3] // 2
+    x1, x2 = x[..., :d], x[..., d:]
+    y1 = x1 * cos + x2 * sin
+    y2 = x1 * (-sin) + x2 * cos
+    out = torch.cat([y1, y2], 3)
+    return out.to(x.dtype)
+
+
+def _repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
+    if n_rep == 1:
+        return x
+    bs, n_kv_heads, slen, head_dim = x.shape
+    return (
+        x[:, :, None, :, :]
+        .expand(bs, n_kv_heads, n_rep, slen, head_dim)
+        .reshape(bs, n_kv_heads * n_rep, slen, head_dim)
+    )
+
+
+class CausalSelfAttention(nn.Module):
+    def __init__(self, config: NanoGPTConfig, layer_idx: int):
+        super().__init__()
+        self.layer_idx = layer_idx
+        self.n_head = config.n_head
+        self.n_kv_head = config.n_kv_head
+        self.n_embd = config.n_embd
+        self.head_dim = self.n_embd // self.n_head
+        assert self.n_embd % self.n_head == 0
+        assert self.n_kv_head <= self.n_head and self.n_head % self.n_kv_head == 0
+        self.c_q = nn.Linear(self.n_embd, self.n_head * self.head_dim, bias=False)
+        self.c_k = nn.Linear(self.n_embd, self.n_kv_head * self.head_dim, bias=False)
+        self.c_v = nn.Linear(self.n_embd, self.n_kv_head * self.head_dim, bias=False)
+        self.c_proj = nn.Linear(self.n_embd, self.n_embd, bias=False)
+
+    def forward(self, x: torch.Tensor, cos_sin, kv_cache=None) -> torch.Tensor:
+        B, T, C = 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)
+        cos, sin = cos_sin
+        q, k = _apply_rotary_emb(q, cos, sin), _apply_rotary_emb(k, cos, sin)
+        q, k = _rms_norm(q), _rms_norm(k)
+        q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
+        Tq = q.size(2)
+        Tk = k.size(2)
+        nrep = self.n_head // self.n_kv_head
+        k, v = _repeat_kv(k, nrep), _repeat_kv(v, nrep)
+        if Tq == Tk:
+            y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+        elif Tq == 1:
+            y = F.scaled_dot_product_attention(q, k, v, is_causal=False)
+        else:
+            attn_mask = torch.zeros((Tq, Tk), dtype=torch.bool, device=q.device)
+            prefix_len = Tk - Tq
+            if prefix_len > 0:
+                attn_mask[:, :prefix_len] = True
+            attn_mask[:, prefix_len:] = torch.tril(torch.ones((Tq, Tq), dtype=torch.bool, device=q.device))
+            y = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)
+        y = y.transpose(1, 2).contiguous().view(B, T, -1)
+        y = self.c_proj(y)
+        return y
+
+
+class MLP(nn.Module):
+    def __init__(self, config: NanoGPTConfig):
+        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: torch.Tensor) -> torch.Tensor:
+        x = self.c_fc(x)
+        x = F.relu(x).square()
+        x = self.c_proj(x)
+        return x
+
+
+class Block(nn.Module):
+    def __init__(self, config: NanoGPTConfig, layer_idx: int):
+        super().__init__()
+        self.attn = CausalSelfAttention(config, layer_idx)
+        self.mlp = MLP(config)
+
+    def forward(self, x: torch.Tensor, cos_sin, kv_cache=None) -> torch.Tensor:
+        x = x + self.attn(_rms_norm(x), cos_sin, kv_cache)
+        x = x + self.mlp(_rms_norm(x))
+        return x
+
+
+class NanoGPTModel(PreTrainedModel):
+    config_class = NanoGPTConfig
+
+    def __init__(self, config: NanoGPTConfig):
+        super().__init__(config)
+        self.transformer = nn.ModuleDict({
+            "wte": nn.Embedding(config.vocab_size, config.n_embd),
+            "h": nn.ModuleList([Block(config, layer_idx) for layer_idx in range(config.n_layer)]),
+        })
+        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
+        self.rotary_seq_len = config.sequence_len * 10
+        head_dim = config.n_embd // config.n_head
+        cos, sin = self._precompute_rotary_embeddings(self.rotary_seq_len, head_dim)
+        self.register_buffer("cos", cos, persistent=False)
+        self.register_buffer("sin", sin, persistent=False)
+        # ensure fp32 activations
+        self.transformer.wte.to(dtype=torch.bfloat16)
+
+        # following HF API expectations
+        self.post_init()
+
+    def _init_weights(self, module: nn.Module):
+        if isinstance(module, nn.Linear):
+            fan_out = module.weight.size(0)
+            fan_in = module.weight.size(1)
+            std = 1.0 / math.sqrt(fan_in) * min(1.0, math.sqrt(fan_out / fan_in))
+            torch.nn.init.normal_(module.weight, mean=0.0, std=std)
+            if module.bias is not None:
+                torch.nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=1.0)
+
+    def _precompute_rotary_embeddings(self, seq_len: int, head_dim: int, base: int = 10000, device=None):
+        if device is None:
+            device = self.transformer.wte.weight.device
+            # Handle meta device case - use CPU as fallback
+            if device.type == 'meta':
+                device = torch.device('cpu')
+        channel_range = torch.arange(0, head_dim, 2, dtype=torch.float32, device=device)
+        inv_freq = 1.0 / (base ** (channel_range / head_dim))
+        t = torch.arange(seq_len, dtype=torch.float32, device=device)
+        freqs = torch.outer(t, inv_freq)
+        cos, sin = freqs.cos(), freqs.sin()
+        cos, sin = cos.bfloat16(), sin.bfloat16()
+        cos, sin = cos[None, :, None, :], sin[None, :, None, :]
+        return cos, sin
+
+    def forward(self, input_ids: torch.Tensor, labels=None, **kwargs):
+        idx = input_ids
+        B, T = idx.size()
+        T0 = 0
+        cos_sin = self.cos[:, T0:T0+T], self.sin[:, T0:T0+T]
+        x = self.transformer.wte(idx)
+        x = x.float()
+        x = _rms_norm(x)
+        for block in self.transformer.h:
+            x = block(x, cos_sin, None)
+        x = _rms_norm(x)
+
+        softcap = 15
+        logits = self.lm_head(x)
+        logits = softcap * torch.tanh(logits / softcap)
+        loss = None
+        if labels is not None:
+            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), labels.view(-1), ignore_index=-1, reduction='mean')
+        return {"loss": loss, "logits": logits}
+
+
+

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:574670de45c667a98fe72c2d145f24007ecc1778721e2481e41500842dd3ba13
+size 2076230219

token_bytes.pt

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

tokenizer.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8467414b90511a50c4dac438af25c075817e9d62d799a5ef613b186c977f5d1b
+size 846518

tokenizer_config.json

@@ -0,0 +1,9 @@
+{
+  "auto_map": {
+    "AutoTokenizer": [
+      "tokenizer_nanogpt.NanoGPTTokenizer",
+      null
+    ]
+  },
+  "tokenizer_class": "NanoGPTTokenizer"
+}

tokenizer_nanogpt.py

@@ -0,0 +1,42 @@
+import os
+import pickle
+
+
+class NanoGPTTokenizer:
+    """Lightweight wrapper over a tiktoken Encoding stored in tokenizer.pkl.
+
+    Provides minimal encode/decode needed for inference and a from_pretrained
+    constructor so it can be loaded via AutoTokenizer with trust_remote_code.
+    """
+
+    def __init__(self, enc):
+        self.enc = enc
+        self.bos_token_id = enc.encode_single_token("<|bos|>")
+
+    @classmethod
+    def register_for_auto_class(cls, auto_class="AutoTokenizer"):
+        """Required for AutoTokenizer registration."""
+        pass
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs):
+        tok_path = os.path.join(pretrained_model_name_or_path, "tokenizer.pkl")
+        with open(tok_path, "rb") as f:
+            enc = pickle.load(f)
+        return cls(enc)
+
+    def encode(self, text, prepend=None):
+        ids = self.enc.encode_ordinary(text)
+        if prepend is not None:
+            prepend_id = prepend if isinstance(prepend, int) else self.enc.encode_single_token(prepend)
+            ids.insert(0, prepend_id)
+        return ids
+
+    def decode(self, ids):
+        return self.enc.decode(ids)
+
+    def get_bos_token_id(self):
+        return self.bos_token_id
+
+
+