手写大模型核心：从零实现一个最小生成式Transformer

import torch
import torch.nn as nn
import torch.nn.functional as F

class CausalAttention(nn.Module):
    def __init__(self, d_model, n_head):
        super().__init__()
        assert d_model % n_head == 0
        self.n_head = n_head
        self.head_dim = d_model // n_head
        
        # 一次性计算 Q、K、V，提升效率
        self.qkv = nn.Linear(d_model, 3 * d_model)
        self.proj = nn.Linear(d_model, d_model)

    def forward(self, x):
        B, T, C = x.shape  # 批次、序列长度、特征维度
        qkv = self.qkv(x).reshape(B, T, 3, self.n_head, self.head_dim)
        qkv = qkv.permute(2, 0, 3, 1, 4)  # 拆出 Q、K、V
        q, k, v = qkv[0], qkv[1], qkv[2]

        # 缩放点积注意力
        att = (q @ k.transpose(-2, -1)) * (self.head_dim ** -0.5)
        
        # 生成下三角掩码，屏蔽未来位置（核心！）
        mask = torch.tril(torch.ones(T, T, device=x.device)).view(1, 1, T, T)
        att = att.masked_fill(mask == 0, float('-inf'))
        att = F.softmax(att, dim=-1)

        y = (att @ v).transpose(1, 2).reshape(B, T, C)
        return self.proj(y)

class FeedForward(nn.Module):
    def __init__(self, d_model):
        super().__init__()
        # 经典扩缩结构：4倍维度膨胀
        self.net = nn.Sequential(
            nn.Linear(d_model, 4 * d_model),
            nn.GELU(),           # 比ReLU更平滑的激活
            nn.Linear(4 * d_model, d_model)
        )
    def forward(self, x): return self.net(x)

class DecoderBlock(nn.Module):
    def __init__(self, d_model, n_head):
        super().__init__()
        self.ln1 = nn.LayerNorm(d_model)
        self.attn = CausalAttention(d_model, n_head)
        self.ln2 = nn.LayerNorm(d_model)
        self.ffwd = FeedForward(d_model)

    def forward(self, x):
        # 残差连接（Pre-Norm结构，现今主流）
        x = x + self.attn(self.ln1(x))
        x = x + self.ffwd(self.ln2(x))
        return x

@torch.no_grad()
def generate(model, idx, max_new_tokens, temperature=1.0, top_k=None):
    """
    idx: 初始输入索引，形状 (1, T)
    model: 我们手写的Decoder堆叠实例
    """
    model.eval()
    for _ in range(max_new_tokens):
        # 若上下文过长，可截断（此处只取最后1024位，模拟滑动窗口）
        idx_cond = idx[:, -1024:] 
        logits = model(idx_cond)          # 前向传播，得到 (B, T, vocab_size)
        logits = logits[:, -1, :]         # 只取最后一个时间步
        logits = logits / temperature     # 温度控制随机性

        # 可选：Top-K采样，只保留概率最高的K个
        if top_k is not None:
            v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
            logits[logits < v[:, [-1]]] = float('-inf')

        probs = F.softmax(logits, dim=-1)
        idx_next = torch.multinomial(probs, num_samples=1)  # 按概率抽样
        idx = torch.cat((idx, idx_next), dim=1)             # 拼回序列
    return idx