「从零开始学大模型」手搓GPT

class BigramLanguageModel(nn.Module, Generate):
    def __init__(self, vocab_size):
        super().__init__()
        self.token_embedding_table = nn.Embedding(vocab_size, vocab_size)

    def forward(self, idx, targets=None):
        # idx (Batch, Time) -> logits (Batch, Time, vocab_size)
        logits = self.token_embedding_table(idx)  
        if targets is None:
            loss = None
        else:
            B, T, C = logits.shape
            logits = logits.view(B*T, C)
            targets = targets.view(B*T)
            loss = F.cross_entropy(logits, targets)
        
        return logits, loss

d_model = 32
class NanoGPTv1(nn.Module, Generate):
    def __init__(self, vocab_size):
        super().__init__()
        # token_idx -> emb_idx
        self.tok_emb_table = nn.Embedding(vocab_size, d_model)
        # position -> emb_pos
        self.pos_emb_table = nn.Embedding(block_size, d_model)

        # Language Model Head
        # 可以理解为 Un-embedding，将 emb_vec 转换回 token_idx
        self.lm_head = nn.Linear(d_model, vocab_size)

    def forward(self, idx, targets=None):
        B, T = idx.shape

        tok_emb = self.tok_emb_table(idx)  # (B, T, d_model)
        pos_emb = self.pos_emb_table(torch.arange(T, device=device))  # (T, d_model)
        x = tok_emb + pos_emb  # (B, T, d_model)
        # 取前面所有的 emb_vec 的平均
        W_f3 = torch.tril(torch.ones((T, T), device=device))  # (T, T) 下三角矩阵
        W_f3 = W_f3 / torch.sum(W_f3, dim=1, keepdim=True)
        x = W_f3 @ x # (B, T, d_model)
        logits = self.lm_head(x)  # (B, T, vocab_size)

        if targets is None:
            loss = None
        else:
            B, T, C = logits.shape
            logits = logits.view(B*T, C)
            targets = targets.view(B*T)
            loss = F.cross_entropy(logits, targets)
        
        return logits, loss

class NanoGPTv2(nn.Module, Generate):
    def __init__(self, vocab_size):
        super().__init__()
        # token_idx -> emb_idx
        self.tok_emb_table = nn.Embedding(vocab_size, d_model)
        # position -> emb_pos
        self.pos_emb_table = nn.Embedding(block_size, d_model)

        self.W_f5 = nn.Parameter(torch.tril(torch.ones((block_size, block_size), device=device)))  # (T, T) 下三角矩阵
        # Language Model Head
        # 可以理解为 Un-embedding，将 emb_vec 转换回 token_idx
        self.lm_head = nn.Linear(d_model, vocab_size)

    def forward(self, idx, targets=None):
        B, T = idx.shape

        tok_emb = self.tok_emb_table(idx)  # (B, T, d_model)
        pos_emb = self.pos_emb_table(torch.arange(T, device=device))  # (T, d_model)
        x = tok_emb + pos_emb  # (B, T, d_model)
        tril = torch.tril(torch.ones((T, T), device=device))
        wei = self.W_f5[:T, :T].masked_fill(tril == 0, float('-inf'))
        wei = F.softmax(wei, dim=1)
        x = wei @ x  # (B, T, d_model)
        logits = self.lm_head(x)  # (B, T, vocab_size)

        if targets is None:
            loss = None
        else:
            B, T, C = logits.shape
            logits = logits.view(B*T, C)
            targets = targets.view(B*T)
            loss = F.cross_entropy(logits, targets)
        
        return logits, loss

d_model = 384
d_k = 64
d_v = 64
class NanoGPTv3(nn.Module, Generate):
    def __init__(self, vocab_size):
        super().__init__()
        self.tok_emb_table = nn.Embedding(vocab_size, d_model)
        self.pos_emb_table = nn.Embedding(block_size, d_model)
        self.W_Q = nn.Linear(d_model, d_k, bias=False)
        self.W_K = nn.Linear(d_model, d_k, bias=False)
        self.W_V = nn.Linear(d_model, d_v, bias=False)
        self.lm_head = nn.Linear(d_v, vocab_size)

    def forward(self, idx, targets=None):
        B, T = idx.shape

        tok_emb = self.tok_emb_table(idx)  # (B, T, d_model)
        pos_emb = self.pos_emb_table(torch.arange(T, device=device))  # (T, d_model)
        x = tok_emb + pos_emb  # (B, T, d_model)
        
        Q = self.W_Q(x)  # (B, T, d_k)
        K = self.W_K(x)  # (B, T, d_k)
        V = self.W_V(x)  # (B, T, d_v)
        W = Q @ K.transpose(-2, -1) * (d_k ** -0.5)
        
        tril = torch.tril(torch.ones((T, T), device=device))
        W = W.masked_fill(tril == 0, float('-inf'))
        W = F.softmax(W, dim=-1)

        y = W @ V  # (B, T, d_v)
        logits = self.lm_head(y)  # (B, T, vocab_size)

        if targets is None:
            loss = None
        else:
            B, T, C = logits.shape
            logits = logits.view(B*T, C)
            targets = targets.view(B*T)
            loss = F.cross_entropy(logits, targets)
        
        return logits, loss

d_model = 384
d_k = 64
d_v = 64

class Head(nn.Module):
    def __init__(self):
        super().__init__()
        self.W_Q = nn.Linear(d_model, d_k, bias=False)
        self.W_K = nn.Linear(d_model, d_k, bias=False)
        self.W_V = nn.Linear(d_model, d_v, bias=False)

    def forward(self, x):
        B, T, C = x.shape
        Q = self.W_Q(x)  # (B, T, d_k)
        K = self.W_K(x)  # (B, T, d_k)
        V = self.W_V(x)  # (B, T, d_v)
        W = Q @ K.transpose(-2, -1) * (d_k ** -0.5)
        
        tril = torch.tril(torch.ones((T, T), device=device))
        W = W.masked_fill(tril == 0, float('-inf'))
        W = F.softmax(W, dim=-1)

        y = W @ V  # (B, T, d_v)
        return y

class MultiHead(nn.Module):
    def __init__(self, num_heads):
        super().__init__()
        self.heads = nn.ModuleList([Head() for _ in range(num_heads)])
        self.proj = nn.Linear(num_heads * d_v, d_model)

    def forward(self, x):
        head_outputs = [head(x) for head in self.heads]  # List of (B, T, d_v) * num_heads
        concatenated = torch.cat(head_outputs, dim=-1)  # (B, T, num_heads * d_v) == (B, T, d_model)
        output = self.proj(concatenated)  # (B, T, d_model)
        return output

d_ff = 4 * d_model
class FFN(nn.Module):
    def __init__(self, d_model, d_ff):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.ReLU(),
            nn.Linear(d_ff, d_model)
        )

    def forward(self, x):
        return self.net(x)

class TransformerBlock(nn.Module):
    def __init__(self):
        super().__init__()
        self.ln1 = nn.LayerNorm(d_model)
        self.heads = MultiHead(num_heads = d_model // d_v)
        self.ln2 = nn.LayerNorm(d_model)
        self.ffn = FFN(d_model, d_ff)
        self.dropout = nn.Dropout(0.1)

    def forward(self, x):
        x = self.dropout(self.heads(self.ln1(x))) + x  # (B, T, d_model)
        x = self.dropout(self.ffn(self.ln2(x))) + x     # (B, T, d_model)
        return x

class NanoGPT(nn.Module):
    def __init__(self):
        # ...
        self.blocks = nn.Sequential(*[TransformerBlock() for _ in range(num_blocks)])
        # ...

class Attention(nn.Module):
    def __init__(self, config: Hyperparameters):
        super().__init__()
        # query, key, value projections for all heads, but in a batch
        self.qkv = nn.Linear(config.d_model, 3 * config.d_model, bias=False)
        # output projection
        self.proj = nn.Linear(config.d_model, config.d_model)
        # regularization
        self.attn_dropout = nn.Dropout(config.dropout)
        self.resid_dropout = nn.Dropout(config.dropout)
        self.d_h = config.d_model // config.d_k
        self.d_model = config.d_model
        self.dropout = config.dropout
        # TODO: flash attention

        self.register_buffer(
            "tril",
            torch.tril(
                torch.ones(config.block_size, config.block_size)
            ).view(
                1, 1, config.block_size, config.block_size
            )
        )
    
    def forward(self, x):
        B, T, C = x.size()  # batch size, sequence length, embedding dimensionality (d_model)

        q, k, v = self.qkv(x).split(self.d_model, dim=2)  # (B, T, d_model) x 3
        # split into heads
        q = q.view(B, T, self.d_h, C // self.d_h).transpose(1, 2)  # (B, nh, T, hs)
        k = k.view(B, T, self.d_h, C // self.d_h).transpose(1, 2)  # (B, nh, T, hs)
        v = v.view(B, T, self.d_h, C // self.d_h).transpose(1, 2)  # (B, nh, T, hs)

        att = (q @ k.transpose(-2, -1)) * (C // self.d_h) ** -0.5  # (B, nh, T, T)
        att = att.masked_fill(self.tril[:, :, :T, :T] == 0, float('-inf'))  # (B, nh, T, T)
        att = F.softmax(att, dim=-1)  # (B, nh, T, T)
        att = self.attn_dropout(att)
        y = att @ v  # (B, nh, T, hs)
        y = y.transpose(1, 2).contiguous().view(B, T, C)  # re-assemble all head outputs side by side
        y = self.proj(y)  # (B, T, C)
        y = self.resid_dropout(y)
        return y

# 数据集获取、处理
#!wget https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt
with open('input.txt', 'r', encoding='utf-8') as f:
    text = f.read()

chars = sorted(list(set(text)))  # 以字符作为 token 
vocab_size = len(chars)          # 根据字符获取字典大小


# 构建字符到索引的映射，encoder、decoder
stoi = {ch: i for i, ch in enumerate(chars)}
itos = {i: ch for i, ch in enumerate(chars)}
encode = lambda s: [stoi[c] for c in s]           # 字符    -> 字典索引
decode = lambda l: ''.join([itos[i] for i in l])  # 字典索引 -> 字符


# 区分训练集和测试集
data = torch.tensor(encode(text), dtype=torch.long)
n = int(0.9 * len(data))
train_data = data[:n]
val_data = data[n:]


# 辅助函数
def get_batch(split):
    data = train_data if split == 'train' else val_data
    ix = torch.randint(len(data) - block_size, (batch_size,))
    x = torch.stack([data[i:i+block_size] for i in ix])
    y = torch.stack([data[i+1:i+block_size+1] for i in ix])
    x, y = x.to(device), y.to(device)
    return x, y

@torch.no_grad()
def estimate_loss():
    out = {}
    model.eval()
    for split in ['train', 'val']:
        losses = torch.zeros(eval_iters)
        for k in range(eval_iters):
            X, Y = get_batch(split)
            logits, loss = model(X, Y)
            losses[k] = loss.item()
        out[split] = losses.mean()
    model.train()
    return out


model = NanoGPTv1(vocab_size).to(device)

optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)

for iter in range(train_iters):
    if iter % eval_interval == 0:
        losses = estimate_loss()
        print(f"step {iter}: train loss {losses['train']:.4f}, val loss {losses['val']:.4f}")

    xb, yb = get_batch('train')
    xb, yb = xb.to(device), yb.to(device)

    logits, loss = model(xb, yb)
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    optimizer.step()

class Generate():
    @torch.no_grad()  # 生成文本时不需要计算梯度，会节省内存和计算资源
    def generate(self, idx, max_new_tokens, temperature=1.0):
        for _ in range(max_new_tokens):
            # 获取输入的最后 block_size 个字符
            idx_cond = idx[:, -block_size:]
            # 前向传播调用 forward()，获取下一个字符的概率分布
            logits, _ = self(idx_cond)
            # 获取最后一个时间步的输出
            logits = logits[:, -1, :]
            # 将 logits 转换为概率分布
            probs = F.softmax(logits / temperature, dim=-1)
            # 从概率分布中采样下一个字符的索引
            next_idx = torch.multinomial(probs, num_samples=1)
            # 将新生成的字符索引添加到输入序列中
            idx = torch.cat((idx, next_idx), dim=1)
        return idx