Upload 3 files

config.json

@@ -0,0 +1,15 @@
+{
+  "architectures": [
+    "WorkshopGPTForCausalLM"
+  ],
+  "block_size": 1024,
+  "dtype": "float32",
+  "model_type": "workshop_gpt",
+  "n_embd": 768,
+  "n_head": 12,
+  "n_inner": 3072,
+  "n_layer": 12,
+  "rope_theta": 10000.0,
+  "transformers_version": "5.0.0",
+  "vocab_size": 50304
+}

model.safetensors

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

modeling_workshop_gpt.py

@@ -0,0 +1,96 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PretrainedConfig, PreTrainedModel
+
+class WorkshopGPTConfig(PretrainedConfig):
+    model_type = "workshop_gpt"
+    def __init__(self, n_layer=12, n_head=12, n_embd=768, vocab_size=50304,
+                 block_size=1024, n_inner=3072, rope_theta=10000.0, **kwargs):
+        super().__init__(**kwargs)
+        self.n_layer, self.n_head, self.n_embd = n_layer, n_head, n_embd
+        self.vocab_size, self.block_size = vocab_size, block_size
+        self.n_inner, self.rope_theta = n_inner, rope_theta
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim, eps=1e-6):
+        super().__init__()
+        self.scale = nn.Parameter(torch.ones(dim))
+        self.eps = eps
+    def forward(self, x):
+        return x * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps).type_as(x) * self.scale
+
+class RotaryPositionalEmbeddings(nn.Module):
+    def __init__(self, dim, max_seq_len=1024, base=10000.0):
+        super().__init__()
+        self.dim, self.max_seq_len, self.base = dim, max_seq_len, base
+        theta = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
+        self.register_buffer("theta", theta, persistent=False)
+        self._build_cache(max_seq_len)
+    def _build_cache(self, seq_len):
+        seq = torch.arange(seq_len, device=self.theta.device)
+        freqs = torch.outer(seq, self.theta)
+        self.register_buffer("cache", torch.stack([freqs.cos(), freqs.sin()], dim=-1), persistent=False)
+    def forward(self, x, *, input_pos=None):
+        seq_len = x.shape[-2]
+        if seq_len > self.cache.shape[0]: self._build_cache(seq_len)
+        cache = self.cache[:seq_len] if input_pos is None else self.cache[input_pos]
+        x1, x2 = x.float().unflatten(-1, (-1, 2)).unbind(-1)
+        cos, sin = cache.unbind(-1)
+        shape = [1] * (x.ndim - 2) + list(cos.shape)
+        cos, sin = cos.view(*shape), sin.view(*shape)
+        return torch.stack([x1*cos - x2*sin, x1*sin + x2*cos], dim=-1).flatten(-2).type_as(x)
+
+class ReluSquaredMLP(nn.Module):
+    def __init__(self, dim, hidden_dim):
+        super().__init__()
+        self.fc_in  = nn.Linear(dim, hidden_dim, bias=False)
+        self.fc_out = nn.Linear(hidden_dim, dim, bias=False)
+    def forward(self, x):
+        return self.fc_out(F.relu(self.fc_in(x)).square())
+
+class CausalSelfAttention(nn.Module):
+    def __init__(self, n_embd, n_head, head_dim, rope):
+        super().__init__()
+        self.n_head, self.head_dim = n_head, head_dim
+        self.q_proj = nn.Linear(n_embd, n_embd, bias=False)
+        self.k_proj = nn.Linear(n_embd, n_embd, bias=False)
+        self.v_proj = nn.Linear(n_embd, n_embd, bias=False)
+        self.output_proj = nn.Linear(n_embd, n_embd, bias=False)
+        self.rope = rope
+    def forward(self, x):
+        B, T, C = x.shape
+        q = self.q_proj(x).view(B, T, self.n_head, self.head_dim).transpose(1, 2)
+        k = self.k_proj(x).view(B, T, self.n_head, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x).view(B, T, self.n_head, self.head_dim).transpose(1, 2)
+        q, k = self.rope(q), self.rope(k)
+        y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+        return self.output_proj(y.transpose(1, 2).contiguous().view(B, T, C))
+
+class TransformerBlock(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        hd = config.n_embd // config.n_head
+        rope = RotaryPositionalEmbeddings(hd, config.block_size, config.rope_theta)
+        self.sa_norm  = RMSNorm(config.n_embd)
+        self.attn     = CausalSelfAttention(config.n_embd, config.n_head, hd, rope)
+        self.mlp_norm = RMSNorm(config.n_embd)
+        self.mlp      = ReluSquaredMLP(config.n_embd, config.n_inner)
+    def forward(self, x):
+        x = x + self.attn(self.sa_norm(x))
+        return x + self.mlp(self.mlp_norm(x))
+
+class WorkshopGPTForCausalLM(PreTrainedModel):
+    config_class = WorkshopGPTConfig
+    def __init__(self, config):
+        super().__init__(config)
+        self.tok_embeddings = nn.Embedding(config.vocab_size, config.n_embd)
+        self.layers = nn.ModuleList([TransformerBlock(config) for _ in range(config.n_layer)])
+        self.norm = RMSNorm(config.n_embd)
+        self.output = nn.Linear(config.n_embd, config.vocab_size, bias=False)
+        self.output.weight = nn.Parameter(self.tok_embeddings.weight.clone())
+    def forward(self, input_ids, **kwargs):
+        x = self.tok_embeddings(input_ids)
+        for layer in self.layers:
+            x = layer(x)
+        return type("Out", (), {"logits": self.output(self.norm(x))})()