File size: 3,609 Bytes
b615ad6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 | """Layer 1: the embedding - id -> vector.
wte: token table (vocab_size, d_model) "what"
wpe: position table (block_size, d_model) "where"
x = wte(ids) + wpe(pos), then optional dropout.
Independence: imports only torch. Knows nothing about layer 0 (data)
or anything above it. Exposes tie_lm_head for the (later) output layer.
"""
from dataclasses import dataclass
import torch
import torch.nn as nn
@dataclass
class Layer1Config:
vocab_size: int = 109
block_size: int = 1024
d_model: int = 512
init_std: float = 0.02
dropout: float = 0.0
class Layer1(nn.Module):
def __init__(self, cfg: Layer1Config):
super().__init__()
self.cfg = cfg
self.wte = nn.Embedding(cfg.vocab_size, cfg.d_model)
self.wpe = nn.Embedding(cfg.block_size, cfg.d_model)
self.drop = nn.Dropout(cfg.dropout)
self.init_weights()
def init_weights(self):
nn.init.normal_(self.wte.weight, mean=0.0, std=self.cfg.init_std)
nn.init.normal_(self.wpe.weight, mean=0.0, std=self.cfg.init_std)
def forward(self, ids, pos_offset=0):
"""ids: (B, T) long -> (B, T, d_model). pos_offset: position embedding
start, used by stateful generation so wpe advances with the stream."""
B, T = ids.shape
assert pos_offset + T <= self.cfg.block_size, \
f"pos_offset+T={pos_offset + T} exceeds block_size {self.cfg.block_size}"
pos = torch.arange(T, device=ids.device) + pos_offset
return self.drop(self.wte(ids) + self.wpe(pos))
def tie_lm_head(self, head):
"""Share wte with the (later) output head: head.weight IS wte.weight."""
head.weight = self.wte.weight
return head
@property
def n_params(self):
return sum(p.numel() for p in self.parameters())
def _self_test():
torch.manual_seed(0)
cfg = Layer1Config()
l1 = Layer1(cfg)
B, T = 4, 16
ids = torch.randint(0, cfg.vocab_size, (B, T))
x = l1(ids)
assert x.shape == (B, T, cfg.d_model), "embed shape"
# same id everywhere -> the wte part is constant across positions
same = torch.full((1, T), 42, dtype=torch.long)
xs = l1(same)
for p in range(T - 1):
d = (xs[0, p] - l1.wpe.weight[p]) - (xs[0, p + 1] - l1.wpe.weight[p + 1])
assert d.abs().max().item() < 1e-6, "wte part varies with position"
# same position, different ids -> the wpe part is constant
for b in range(B - 1):
d = (x[b, 0] - l1.wte(ids[b, 0])) - (x[b + 1, 0] - l1.wte(ids[b + 1, 0]))
assert d.abs().max().item() < 1e-6, "wpe part varies with token"
# gradients flow through both tables
x.sum().backward()
assert l1.wte.weight.grad is not None and l1.wte.weight.grad.abs().sum().item() > 0
assert l1.wpe.weight.grad is not None and l1.wpe.weight.grad.abs().sum().item() > 0
# param count: 109*256 + 256*256
assert l1.n_params == cfg.vocab_size * cfg.d_model + cfg.block_size * cfg.d_model == 93440
# weight tying is a true share, not a copy
head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
l1.tie_lm_head(head)
assert head.weight is l1.wte.weight, "tie is not a shared tensor"
# block_size guard
try:
l1(torch.zeros((1, cfg.block_size + 1), dtype=torch.long))
raise SystemExit("expected assertion for T > block_size")
except AssertionError:
pass
print(f"wte {tuple(l1.wte.weight.shape)} wpe {tuple(l1.wpe.weight.shape)} "
f"params {l1.n_params:,}")
print("layer1 self-test: PASS")
if __name__ == "__main__":
_self_test()
|