Upload 9 files

config.json

@@ -0,0 +1,21 @@
+{
+  "architectures": [
+    "MiniLlama"
+  ],
+  "auto_map": {
+    "AutoConfig": "config_minillama.MiniLlamaConfig",
+    "AutoModel": "modeling_minillama.MiniLlama"
+  },
+  "dim": 768,
+  "dropout": 0.1,
+  "dtype": "float32",
+  "max_seq_len": 1024,
+  "model_type": "mini-llama",
+  "multiple_of": 256,
+  "n_heads": 12,
+  "n_kv_heads": 12,
+  "n_layers": 24,
+  "norm_eps": 1e-05,
+  "transformers_version": "4.57.3",
+  "vocab_size": 50257
+}

config_minillama.py

@@ -0,0 +1,33 @@
+from transformers import AutoConfig, PretrainedConfig
+
+
+class MiniLlamaConfig(PretrainedConfig):
+    model_type = "mini-llama"
+
+    def __init__(
+        self,
+        *,
+        vocab_size: int = 32000,
+        dim: int = 512,  # Embedding dimension
+        n_layers: int = 8,  # Number of transformer blocks
+        n_heads: int = 8,  # Attention heads
+        n_kv_heads: int = 8,  # Key/Value heads (for Grouped Query Attention)
+        multiple_of: int = 256,  # For SwiGLU hidden layer size
+        norm_eps: float = 1e-5,
+        max_seq_len: int = 1024,
+        dropout: float = 0.0,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.vocab_size = vocab_size
+        self.dim = dim
+        self.n_layers = n_layers
+        self.n_heads = n_heads
+        self.n_kv_heads = n_kv_heads
+        self.multiple_of = multiple_of
+        self.norm_eps = norm_eps
+        self.max_seq_len = max_seq_len
+        self.dropout = dropout
+
+
+MiniLlamaConfig.register_for_auto_class(AutoConfig)

model.safetensors

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

modeling_minillama.py

@@ -0,0 +1,170 @@
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+# We use relative references here because this file is a training artifact
+# Hard references will prevent it from being dynamically loaded correctly
+# ideally it should only reference the config and nothing else in our project
+from .config_minillama import MiniLlamaConfig
+from transformers import AutoModel, PreTrainedModel
+
+
+def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0):
+    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
+    t = torch.arange(end, device=freqs.device)
+    freqs = torch.outer(t, freqs).float()
+    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)  # complex64
+    return freqs_cis
+
+
+def apply_rotary_emb(xq, xk, freqs_cis):
+    # Reshape as complex numbers
+    xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
+    xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
+
+    # Align shapes for broadcasting
+    freqs_cis = freqs_cis.view(1, xq_.size(1), 1, xq_.size(-1))
+
+    # Rotate!
+    xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)
+    xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)
+    return xq_out.type_as(xq), xk_out.type_as(xk)
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def _norm(self, x):
+        # rsqrt = reciprocal square root (1 / sqrt(x))
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+    def forward(self, x):
+        return self._norm(x.float()).type_as(x) * self.weight
+
+
+class Attention(nn.Module):
+    def __init__(self, args: MiniLlamaConfig):
+        super().__init__()
+        self.n_heads = args.n_heads
+        self.n_kv_heads = args.n_kv_heads
+        if args.dim % args.n_heads != 0:
+            raise ValueError("Model dimension must be divisible by number of heads.")
+        head_dim = args.dim // args.n_heads
+        if head_dim % 2 != 0:
+            raise ValueError("Head dimension must be even to apply rotary embeddings.")
+        if args.n_kv_heads > args.n_heads:
+            raise ValueError("n_kv_heads must be less than or equal to n_heads.")
+        if args.n_heads % args.n_kv_heads != 0:
+            raise ValueError("n_heads must be divisible by n_kv_heads for GQA.")
+        self.head_dim = head_dim
+
+        self.wq = nn.Linear(args.dim, args.n_heads * self.head_dim, bias=False)
+        self.wk = nn.Linear(args.dim, self.n_kv_heads * self.head_dim, bias=False)
+        self.wv = nn.Linear(args.dim, self.n_kv_heads * self.head_dim, bias=False)
+        self.wo = nn.Linear(args.n_heads * self.head_dim, args.dim, bias=False)
+
+    def forward(self, x, freqs_cis, mask=None):
+        bsz, seqlen, _ = x.shape
+        xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)
+
+        # Reshape for heads
+        xq = xq.view(bsz, seqlen, self.n_heads, self.head_dim)
+        xk = xk.view(bsz, seqlen, self.n_kv_heads, self.head_dim)
+        xv = xv.view(bsz, seqlen, self.n_kv_heads, self.head_dim)
+
+        # Apply RoPE
+        xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
+
+        # GQA: Repeat keys/values if n_kv_heads < n_heads
+        if self.n_kv_heads < self.n_heads:
+            xk = torch.repeat_interleave(xk, self.n_heads // self.n_kv_heads, dim=2)
+            xv = torch.repeat_interleave(xv, self.n_heads // self.n_kv_heads, dim=2)
+
+        # Attention Calculation
+        xq, xk, xv = xq.transpose(1, 2), xk.transpose(1, 2), xv.transpose(1, 2)
+        scores = torch.matmul(xq, xk.transpose(2, 3)) / math.sqrt(self.head_dim)
+        if mask is not None:
+            scores = scores + mask
+        output = torch.matmul(F.softmax(scores.float(), dim=-1).type_as(xq), xv)
+
+        return self.wo(output.transpose(1, 2).contiguous().view(bsz, seqlen, -1))
+
+
+class FeedForward(nn.Module):
+    def __init__(self, dim, hidden_dim, multiple_of):
+        super().__init__()
+        hidden_dim = int(2 * hidden_dim / 3)
+        hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)
+
+        self.w1 = nn.Linear(dim, hidden_dim, bias=False)
+        self.w2 = nn.Linear(hidden_dim, dim, bias=False)
+        self.w3 = nn.Linear(dim, hidden_dim, bias=False)
+
+    def forward(self, x):
+        return self.w2(F.silu(self.w1(x)) * self.w3(x))
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, layer_id, args):
+        super().__init__()
+        self.layer_id = layer_id
+        self.attention = Attention(args)
+        self.feed_forward = FeedForward(args.dim, 4 * args.dim, args.multiple_of)
+        self.attention_norm = RMSNorm(args.dim, eps=args.norm_eps)
+        self.ffn_norm = RMSNorm(args.dim, eps=args.norm_eps)
+        self.dropout = nn.Dropout(args.dropout)
+
+    def forward(self, x, freqs_cis, mask):
+        h = x + self.dropout(self.attention(self.attention_norm(x), freqs_cis, mask))
+        out = h + self.dropout(self.feed_forward(self.ffn_norm(h)))
+        return out
+
+
+class MiniLlama(PreTrainedModel):
+    config_class = MiniLlamaConfig
+
+    def __init__(self, config: MiniLlamaConfig):
+        super().__init__(config)
+        self.config = config
+        self.tok_embeddings = nn.Embedding(config.vocab_size, config.dim)
+        self.layers = nn.ModuleList(
+            [TransformerBlock(i, config) for i in range(config.n_layers)]
+        )
+        self.norm = RMSNorm(config.dim, eps=config.norm_eps)
+        self.output = nn.Linear(config.dim, config.vocab_size, bias=False)
+
+        # Precompute RoPE
+        self.freqs_cis = precompute_freqs_cis(
+            config.dim // config.n_heads, config.max_seq_len * 2
+        )
+
+    def forward(self, tokens, targets=None):
+        bsz, seqlen = tokens.shape
+        h = self.tok_embeddings(tokens)
+        freqs_cis = self.freqs_cis[:seqlen].to(h.device)
+
+        mask = None
+        if seqlen > 1:
+            mask = torch.full((seqlen, seqlen), float("-inf"), device=tokens.device)
+            # Causal mask
+            mask = torch.triu(mask, diagonal=1).unsqueeze(0).unsqueeze(0)
+
+        for layer in self.layers:
+            h = layer(h, freqs_cis, mask)
+
+        logits = self.output(self.norm(h))
+
+        if targets is not None:
+            # Add loss calculation to make post training easy.
+            # You don't have to do this
+            return logits, F.cross_entropy(
+                logits.view(-1, self.config.vocab_size),
+                targets.view(-1),
+            )
+        return logits, None
+
+
+MiniLlama.register_for_auto_class(AutoModel)

special_tokens_map.json

@@ -0,0 +1,6 @@
+{
+  "bos_token": "<|endoftext|>",
+  "eos_token": "<|endoftext|>",
+  "pad_token": "<|endoftext|>",
+  "unk_token": "<|endoftext|>"
+}

tokenizer_config.json

@@ -0,0 +1,21 @@
+{
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "50256": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "<|endoftext|>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|endoftext|>",
+  "extra_special_tokens": {},
+  "model_max_length": 1024,
+  "pad_token": "<|endoftext|>",
+  "tokenizer_class": "GPT2Tokenizer",
+  "unk_token": "<|endoftext|>"
+}