Update modeling_tinybuddy.py

modeling_tinybuddy.py

@@ -1,46 +1,13 @@
 """
 Tiny GPT-style transformer (~30M params target).
-
-Config:
-  - 6 layers
-  - 8 heads
-  - d_model = 256
-  - vocab_size = 32000 (chosen to push param count up to ~30M, since the
-    transformer blocks themselves only have ~5M params at d_model=256/L=6;
-    the embedding + tied LM head dominates the parameter budget.)
-
-Parameter accounting (approx):
-  Token embedding  : 32000 * 256                    = 8,192,000
-  LM head (untied) : 256   * 32000 + 32000          = 8,224,000
-  Positional emb   : 512   * 256                    =   131,072
-  Per block (x6):
-    attn (qkv+out) : 4 * 256 * 256 + 4*256          =   263,168
-    mlp  (2 linear): 256*1024 + 1024 + 1024*256+256 =   525,568
-    2x LayerNorm   : 4 * 256                        =     1,024
-    block total                                     =   789,760
-  Blocks total     : 6 * 789,760                    = 4,738,560
-  Final LN         : 512
-  ---------------------------------------------------------
-  TOTAL ~ 21.3M (tied) or ~29.5M (untied lm head)  -> ~30M ✓
 """
 
 import math
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from dataclasses import dataclass
-
-
-@dataclass
-class GPTConfig:
-    vocab_size: int = 50000
-    block_size: int = 512      # max context length
-    n_layer: int = 6
-    n_head: int = 8
-    n_embd: int = 256
-    mlp_ratio: int = 4         # hidden = 4 * n_embd
-    dropout: float = 0.0
-    tie_weights: bool = False  # False -> ~30M params; True -> ~21M
+from transformers import PreTrainedModel
+from configuration_tinybuddy import GPTConfig
 
 
 class CausalSelfAttention(nn.Module):
@@ -97,59 +64,59 @@ class Block(nn.Module):
         return x
 
 
-class TinyGPT(nn.Module):
-    def __init__(self, cfg: GPTConfig):
-        super().__init__()
-        self.cfg = cfg
-        self.tok_emb = nn.Embedding(cfg.vocab_size, cfg.n_embd)
-        self.pos_emb = nn.Embedding(cfg.block_size, cfg.n_embd)
-        self.drop = nn.Dropout(cfg.dropout)
-        self.blocks = nn.ModuleList([Block(cfg) for _ in range(cfg.n_layer)])
-        self.ln_f = nn.LayerNorm(cfg.n_embd)
-        self.lm_head = nn.Linear(cfg.n_embd, cfg.vocab_size, bias=False)
-        if cfg.tie_weights:
+class TinyGPT(PreTrainedModel):
+    config_class = GPTConfig
+    
+    def __init__(self, config: GPTConfig):
+        super().__init__(config)
+        self.config = config
+        self.tok_emb = nn.Embedding(config.vocab_size, config.n_embd)
+        self.pos_emb = nn.Embedding(config.block_size, config.n_embd)
+        self.drop = nn.Dropout(config.dropout)
+        self.blocks = nn.ModuleList([Block(config) for _ in range(config.n_layer)])
+        self.ln_f = nn.LayerNorm(config.n_embd)
+        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
+        if config.tie_weights:
             self.lm_head.weight = self.tok_emb.weight
-        self.apply(self._init_weights)
+        self.post_init()
 
-    @staticmethod
-    def _init_weights(m):
-        if isinstance(m, nn.Linear):
-            nn.init.normal_(m.weight, mean=0.0, std=0.02)
-            if m.bias is not None:
-                nn.init.zeros_(m.bias)
-        elif isinstance(m, nn.Embedding):
-            nn.init.normal_(m.weight, mean=0.0, std=0.02)
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range if hasattr(self.config, 'initializer_range') else 0.02)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range if hasattr(self.config, 'initializer_range') else 0.02)
 
     def num_params(self, non_embedding=False):
         n = sum(p.numel() for p in self.parameters())
         if non_embedding:
             n -= self.tok_emb.weight.numel() + self.pos_emb.weight.numel()
-            if not self.cfg.tie_weights:
+            if not self.config.tie_weights:
                 n -= self.lm_head.weight.numel()
         return n
 
-    def forward(self, idx, targets=None):
-        B, T = idx.shape
-        assert T <= self.cfg.block_size, f"sequence length {T} > block_size {self.cfg.block_size}"
-        pos = torch.arange(T, device=idx.device)
-        x = self.tok_emb(idx) + self.pos_emb(pos)[None, :, :]
+    def forward(self, input_ids, attention_mask=None, labels=None, **kwargs):
+        B, T = input_ids.shape
+        assert T <= self.config.block_size, f"sequence length {T} > block_size {self.config.block_size}"
+        pos = torch.arange(T, device=input_ids.device)
+        x = self.tok_emb(input_ids) + self.pos_emb(pos)[None, :, :]
         x = self.drop(x)
         for blk in self.blocks:
             x = blk(x)
         x = self.ln_f(x)
         logits = self.lm_head(x)
         loss = None
-        if targets is not None:
+        if labels is not None:
             loss = F.cross_entropy(logits.view(-1, logits.size(-1)),
-                                   targets.view(-1), ignore_index=-100)
-        return logits, loss
+                                   labels.view(-1), ignore_index=-100)
+        return (logits,) if loss is None else (logits, loss)
 
-    @torch.no_grad()
     def generate(self, idx, max_new_tokens=100, temperature=1.0, top_k=None):
         self.eval()
         for _ in range(max_new_tokens):
-            idx_cond = idx if idx.size(1) <= self.cfg.block_size else idx[:, -self.cfg.block_size:]
-            logits, _ = self(idx_cond)
+            idx_cond = idx if idx.size(1) <= self.config.block_size else idx[:, -self.config.block_size:]
+            logits = self(idx_cond)[0]
             logits = logits[:, -1, :] / max(temperature, 1e-6)
             if top_k is not None:
                 v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
@@ -161,9 +128,10 @@ class TinyGPT(nn.Module):
 
 
 if __name__ == "__main__":
+    from configuration_tinybuddy import GPTConfig
     cfg = GPTConfig()
     m = TinyGPT(cfg)
     total = m.num_params()
     nonemb = m.num_params(non_embedding=True)
     print(f"Total params       : {total:,}  (~{total/1e6:.2f}M)")
-    print(f"Non-embedding params: {nonemb:,}  (~{nonemb/1e6:.2f}M)")
+    print(f"Non-embedding params: {nonemb:,}  (~{nonemb/1e6:.2f}M)")