无库无捷径 PyTorch手写完整Transformer大语言模型

只需一行代码 from transformers import AutoModel 就能加载预训练模型，但真正透彻理解 Transformer 底层原理的开发者却寥寥无几。许多人对 Attention 的认知止步于“Q、K、V 三个矩阵”的口诀——面试时能流畅背诵公式，真要动手实现一个完整的 Block 就寸步难行。本文的目标极为纯粹：不借助任何高层封装库，仅依靠 PyTorch 原生算子，从零手写一个完整的 Decoder-only Transformer 语言模型（即 GPT 架构，也是 GPT-4 / LLaMA / Qwen 的基础）。你看到的每一行代码都能在经典论文中找到对应的公式。通过这场实战，你将从“调包侠”蜕变为真正的模型建造者。

一、架构对齐：Decoder-only Transformer

下面是我们将要逐个攻破的核心模块清单：

• Token Embedding → 将离散 token 映射到连续向量空间
• RoPE 位置编码 → 注入相对位置信息，支持外推
• Multi-Head Attention + Causal Mask → 自回归注意力计算
• Feed-Forward Network → 非线性特征变换
• Pre-Norm 残差连接 → 稳定深层训练
• 自回归生成 → 逐个 token 预测，实现文本生成

二、基础配置与超参数

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

class Config:
    vocab_size = 10000   # 词表大小
    d_model = 512        # 模型维度
    n_heads = 8          # 注意力头数
    n_layers = 6         # Block 层数
    d_ff = 2048          # FFN 中间维度
    max_seq_len = 512    # 最大序列长度
    dropout = 0.1
    pad_id = 0

config = Config()
device = 'cuda' if torch.cuda.is_a vailable() else 'cpu'

三、逐层实现详解

3.1 Token Embedding

将每个离散 token 映射为 d_model 维的连续向量，并按照 Transformer 原始论文的要求，乘以 √d_model 完成尺度对齐，使嵌入向量与后续位置编码的尺度匹配。

class TokenEmbedding(nn.Module):
    def __init__(self, vocab_size, d_model):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.d_model = d_model

    def forward(self, x):
        return self.embedding(x) * math.sqrt(self.d_model)

3.2 旋转位置编码（RoPE）

RoPE 是 LLaMA、Qwen 等主流 LLM 广泛采用的位置编码方案。它通过旋转操作将位置信息编码到 Q 和 K 向量中，天然支持相对位置外推，效果远超传统绝对位置编码，也是位置编码领域的热门选择。

class RotaryPositionalEmbedding(nn.Module):
    def __init__(self, d_model, max_seq_len):
        super().__init__()
        d_half = d_model // 2
        # 频率：θ_i = 10000^(-2i/d)
        freqs = 1.0 / (10000 ** (torch.arange(0, d_half).float() / d_half))
        positions = torch.arange(max_seq_len).float()
        angles = torch.outer(positions, freqs)  # [seq_len, d_half]
        self.register_buffer('cos', angles.cos())
        self.register_buffer('sin', angles.sin())

    def forward(self, x):
        seq_len = x.size(1)
        cos = self.cos[:seq_len].unsqueeze(0)
        sin = self.sin[:seq_len].unsqueeze(0)
        d_half = x.size(-1) // 2
        x1, x2 = x[..., :d_half], x[..., d_half:]
        # 旋转：(x1 + i*x2) * (cos + i*sin)
        return torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)

3.3 Multi-Head Causal Self-Attention

核心公式铭记于心：Attention(Q,K,V) = softmax(QK^T / √d_k) V。下面代码严格复现了这一计算，并加入了因果掩码以实现自回归。

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, n_heads, dropout=0.1):
        super().__init__()
        assert d_model % n_heads == 0
        self.n_heads = n_heads
        self.d_head = d_model // n_heads
        self.W_q = nn.Linear(d_model, d_model, bias=False)
        self.W_k = nn.Linear(d_model, d_model, bias=False)
        self.W_v = nn.Linear(d_model, d_model, bias=False)
        self.W_o = nn.Linear(d_model, d_model, bias=False)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, rope=None, mask=None):
        B, S, _ = x.shape
        # 线性投影 → 切分多头
        Q = self.W_q(x).view(B, S, self.n_heads, self.d_head).transpose(1, 2)
        K = self.W_k(x).view(B, S, self.n_heads, self.d_head).transpose(1, 2)
        V = self.W_v(x).view(B, S, self.n_heads, self.d_head).transpose(1, 2)
        # 注入 RoPE（仅作用于 Q 和 K）
        if rope is not None:
            Q = rope(Q)
            K = rope(K)
        # 注意力分数
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_head)
        # 因果掩码：上三角置 -inf
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))
        attn = self.dropout(F.softmax(scores, dim=-1))
        out = torch.matmul(attn, V)  # [B, n_heads, S, d_head]
        out = out.transpose(1, 2).contiguous().view(B, S, -1)
        return self.W_o(out)

3.4 Feed-Forward Network

经典的 MLP 结构，先升维再降维，通过 ReLU 引入非线性，增强模型表达能力。

class FeedForward(nn.Module):
    def __init__(self, d_model, d_ff, dropout=0.1):
        super().__init__()
        self.w1 = nn.Linear(d_model, d_ff)
        self.w2 = nn.Linear(d_ff, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        return self.w2(self.dropout(F.relu(self.w1(x))))

3.5 Transformer Block（Pre-Norm 残差结构）

现代 LLM 全部采用 Pre-Norm 设计——先做 LayerNorm 再进入 Attention / FFN，训练稳定性远超原始论文中的 Post-Norm，这也是深度网络得以收敛的关键。

class TransformerBlock(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.attn = MultiHeadAttention(config.d_model, config.n_heads, config.dropout)
        self.ffn = FeedForward(config.d_model, config.d_ff, config.dropout)
        self.norm1 = nn.LayerNorm(config.d_model)
        self.norm2 = nn.LayerNorm(config.d_model)
        self.dropout = nn.Dropout(config.dropout)

    def forward(self, x, rope=None, mask=None):
        x = x + self.dropout(self.attn(self.norm1(x), rope=rope, mask=mask))
        x = x + self.dropout(self.ffn(self.norm2(x)))
        return x

3.6 完整模型组装

将所有组件拼接成完整的 Decoder-only 语言模型，特别地，Embedding 与 LM Head 共享权重（参数减半），这是 GPT 系列常用的参数节约技巧。

class TransformerLLM(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.token_embedding = TokenEmbedding(config.vocab_size, config.d_model)
        self.rope = RotaryPositionalEmbedding(config.d_model, config.max_seq_len)
        self.blocks = nn.ModuleList([TransformerBlock(config) for _ in range(config.n_layers)])
        self.final_norm = nn.LayerNorm(config.d_model)
        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
        # 权重共享
        self.lm_head.weight = self.token_embedding.embedding.weight
        self.apply(self._init_weights)

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(self, input_ids):
        B, S = input_ids.shape
        # 因果掩码：下三角矩阵
        mask = torch.tril(torch.ones(S, S, device=input_ids.device)).unsqueeze(0).unsqueeze(0)
        x = self.token_embedding(input_ids)
        for block in self.blocks:
            x = block(x, rope=self.rope, mask=mask)
        x = self.final_norm(x)
        return self.lm_head(x)   # [B, S, vocab_size]

3.7 前向传播验证

用随机输入测试维度是否正确，确保模型搭建无误。

model = TransformerLLM(config).to(device)
n_params = sum(p.numel() for p in model.parameters())
print(f"参数量: {n_params/1e6:.1f}M")

dummy = torch.randint(1, config.vocab_size, (2, 64)).to(device)
logits = model(dummy)
print(f"输出: {logits.shape}")   # [2, 64, 10000]

四、训练与推理实战

4.1 训练单步实现

采用标准交叉熵损失，注意输入与目标错位（shift），并使用梯度裁剪防止梯度爆炸。

def train_step(model, batch, optimizer, config):
    model.train()
    input_ids = batch[:, :-1].to(device)   # 去掉最后一个 token
    targets = batch[:, 1:].to(device)      # 左移一位
    logits = model(input_ids)
    loss = F.cross_entropy(logits.reshape(-1, config.vocab_size),
                           targets.reshape(-1),
                           ignore_index=config.pad_id)
    optimizer.zero_grad()
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
    optimizer.step()
    return loss.item()

4.2 自回归生成流程

逐个 token 采样，每次取概率最高的 token 拼接，直到达到目标长度或遇到结束符。这是 GPT 推理的标准方式。

@torch.no_grad()
def generate(model, start_tokens, max_new_tokens, config):
    model.eval()
    input_ids = torch.tensor([start_tokens]).to(device)
    for _ in range(max_new_tokens):
        logits = model(input_ids[:, -config.max_seq_len:])   # 裁剪到最大长度
        next_token = logits[:, -1, :].argmax(-1).unsqueeze(-1)
        input_ids = torch.cat([input_ids, next_token], dim=-1)
    return input_ids[0].tolist()

五、模型规模与数据工程

通过调整 Config 即可随意缩放模型：

工业级训练的真正瓶颈往往不在模型代码本身，而在数据吞吐——高质量语料的采集和清洗需要完整的数据管线支撑，这部分工作量常被低估。

六、踩坑速查手册

七、总结与核心收获

手写一遍 Transformer 后，你将获得每个组件的深层理解：

• Embedding：离散 token → 连续空间，乘 √d_model 完成尺度对齐
• RoPE：用旋转编码将位置信息注入 Q/K，天然支持相对位置外推
• Attention：Q·K^T 计算相似度 → softmax 归一化 → 加权融合 V，复杂度 O(n²·d)
• Causal Mask：上三角填入 -inf，保证自回归单向信息流
• Pre-Norm 残差：信息直通通道，让深层网络得以稳定训练
• 权重共享：Embedding 与 LM Head 共享权重，参数量减半

当面试官问“Attention 的计算复杂度为什么是 O(n²)”时，你能立刻回答：因为 Q·K^T 矩阵乘法的维度是 [seq, d] × [d, seq]，每一对位置都要计算内积。这种底层视角，是仅仅调用 from_pretrained() 永远无法获得的。