Views
No views yet
| d_model | 256 |
| layers | 200 |
| heads | 8 |
| d_ff | 1024 |
| vocab | 16,384 (word-level, custom) |
| seq_len | 512 |
| params | 166M |
| dtype | bfloat16 |
AutoTokenizer compatible), has no chat/instruct
format, and was trained on a relatively small corpus (~270K words). It is a research
artifact for the Narrow-vs-Wide comparison, not a production-ready model.1import torch
2import torch.nn as nn
3import torch.nn.functional as F
4
5VOCAB_SIZE, SEQ_LEN = 16384, 512
6
7class Block(nn.Module):
8 def __init__(self, d, n_heads, d_ff):
9 super().__init__()
10 self.norm1, self.norm2 = nn.LayerNorm(d), nn.LayerNorm(d)
11 self.Wq = nn.Linear(d, d, bias=False)
12 self.Wk = nn.Linear(d, d, bias=False)
13 self.Wv = nn.Linear(d, d, bias=False)
14 self.Wo = nn.Linear(d, d, bias=False)
15 self.ff1 = nn.Linear(d, d_ff, bias=False)
16 self.ff2 = nn.Linear(d_ff, d, bias=False)
17 self.n_heads, self.d_head = n_heads, d // n_heads
18
19 def forward(self, x):
20 B, T, D = x.shape
21 h = self.norm1(x)
22 q = self.Wq(h).reshape(B, T, self.n_heads, self.d_head).transpose(1, 2)
23 k = self.Wk(h).reshape(B, T, self.n_heads, self.d_head).transpose(1, 2)
24 v = self.Wv(h).reshape(B, T, self.n_heads, self.d_head).transpose(1, 2)
25 out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
26 x = x + self.Wo(out.transpose(1, 2).reshape(B, T, D))
27 return x + self.ff2(F.gelu(self.ff1(self.norm2(x))))
28
29class NarrowLM(nn.Module):
30 def __init__(self, vocab=VOCAB_SIZE, d=256, n_heads=8, n_layers=200, d_ff=1024):
31 super().__init__()
32 self.embed = nn.Embedding(vocab, d)
33 self.pos = nn.Embedding(SEQ_LEN, d)
34 self.blocks = nn.ModuleList([Block(d, n_heads, d_ff) for _ in range(n_layers)])
35 self.norm = nn.LayerNorm(d)
36 self.lm_head = nn.Linear(d, vocab, bias=False)
37
38 def forward(self, idx):
39 B, T = idx.shape
40 x = self.embed(idx) + self.pos(torch.arange(T, device=idx.device))
41 for blk in self.blocks:
42 x = blk(x)
43 return self.lm_head(self.norm(x))
44
45model = NarrowLM()
46model.load_state_dict(torch.load("Narrow_distill_1B.pt", map_location="cpu", weights_only=False))
47model.eval()