File size: 2,583 Bytes
5c6c847 | 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 | """对照组：标准的 GPT 式 Transformer（因果自注意力），结构上和 ZeroX 一一对应。
位置编码两种可选（《Attention Is All You Need》§3.5 两种都试过，结果几乎一样）：
pos="learned" 可学习的位置向量，只有 block_size 个，处理不了更长的序列
pos="sin" 论文用的正弦编码，任意长度都能算（论文选它就是因为"可能可以外推"）
"""
import numpy as np
from zerocg import Tensor
from zerocg import functional as F
from zerocg import nn
def sinusoidal_encoding(T, d):
"""PE(pos, 2i) = sin(pos / 10000^(2i/d))，PE(pos, 2i+1) = cos(...)。论文 §3.5。"""
pos = np.arange(T)[:, None]
div = 10000 ** (np.arange(0, d, 2) / d)
pe = np.zeros((T, d), dtype=np.float32)
pe[:, 0::2] = np.sin(pos / div)
pe[:, 1::2] = np.cos(pos / div)
return pe
class Block(nn.Module):
def __init__(self, n_embd, n_head, dropout=0.0):
self.ln1 = nn.LayerNorm(n_embd)
self.attn = nn.CausalSelfAttention(n_embd, n_head, dropout)
self.ln2 = nn.LayerNorm(n_embd)
self.fc1 = nn.Linear(n_embd, 4 * n_embd)
self.fc2 = nn.Linear(4 * n_embd, n_embd)
self.drop = nn.Dropout(dropout)
def forward(self, x):
x = x + self.attn(self.ln1(x))
return x + self.drop(self.fc2(F.gelu(self.fc1(self.ln2(x)))))
class TransformerLM(nn.Module):
def __init__(self, vocab_size, block_size, n_embd=64, n_head=4, n_layer=2, dropout=0.0, pos="learned"):
assert pos in ("learned", "sin")
self.block_size = block_size
self.pos_kind = pos
self.tok = nn.Embedding(vocab_size, n_embd)
if pos == "learned":
self.pos = nn.Embedding(block_size, n_embd)
self.blocks = [Block(n_embd, n_head, dropout) for _ in range(n_layer)]
self.ln_f = nn.LayerNorm(n_embd)
self.head = nn.Linear(n_embd, vocab_size)
def forward(self, idx):
T = idx.shape[1]
if self.pos_kind == "learned":
assert T <= self.block_size, f"位置编码只有 {self.block_size} 个，处理不了 {T} 个字"
x = self.tok(idx) + self.pos(np.arange(T))
else:
x = self.tok(idx) + Tensor(sinusoidal_encoding(T, self.tok.weight.shape[1]))
for b in self.blocks:
x = b(x)
return self.head(self.ln_f(x))
def state_floats(self, T):
"""推理到第 T 个字时 KV 缓存的大小（浮点数个数），随 T 线性增长。"""
return sum(2 * T * b.ln1.weight.shape[0] for b in self.blocks)
|