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

原创

小白学大数据

发布于 2026-07-07 16:53:10

调用 from transformers import AutoModel 只要一行代码，但真正理解 Transformer 的人少之又少。多数开发者对 Attention 的认知停留在"Q、K、V 三个矩阵"的口号层面——面试时背得出公式，动手时写不出一个完整的 Block。本文的目标很纯粹：不用任何高层封装库，只用 PyTorch 原生算子，从零实现一个完整的 Decoder-only Transformer 语言模型（GPT 架构，也是 GPT-4 / LLaMA / Qwen 的基础架构）。每一行代码对应论文中的一个公式。一、架构对齐：Decoder-only Transformer

核心实现清单：Token Embedding → RoPE 位置编码 → Multi-Head Attention + Causal Mask → Feed-Forward → Pre-Norm 残差 → 自回归生成。二、基础配置

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_available() else 'cpu'

三、逐层实现3.1 Token Embedding

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):
        # 论文要求：embedding 乘以 sqrt(d_model) 做尺度对齐
        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(QKTdk)VAttention(Q,K,V)=softmax(dkQKT)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

        # 1. 线性投影 → 切分多头
        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)
        # Q/K/V: [B, n_heads, S, d_head]

        # 2. RoPE 注入位置信息（仅作用于 Q 和 K）
        if rope is not None:
            Q = rope(Q)
            K = rope(K)

        # 3. 注意力分数 = Q @ K^T / sqrt(d_head)
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_head)

        # 4. 因果掩码：上三角置 -inf，softmax 后概率为 0
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))

        # 5. Softmax → 加权 V → 合并多头 → 输出投影
        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

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（先 Norm 再 Attention），训练稳定性远优于原始论文的 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):
        # Sub-layer 1: Attention + 残差
        x = x + self.dropout(self.attn(self.norm1(x), rope=rope, mask=mask))
        # Sub-layer 2: FFN + 残差
        x = x + self.dropout(self.ffn(self.norm2(x)))
        return x

3.6 完整模型

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)
        # 权重共享：Embedding 和 LM Head 用同一套权重
        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 训练单步

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 自回归生成

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()

五、模型规模与训练数据工程调整 Config 即可缩放模型规模：

配置	d_model	n_heads	n_layers	参数量	对标
教学版	512	8	6	~38M	—
GPT-2 Small	768	12	12	~85M	GPT-2
GPT-2 Medium	1024	16	24	~350M	GPT-2 Medium

工业级训练的真正瓶颈往往不在模型代码，而在数据吞吐——高质量语料的采集和清洗需要完整的数据管线支撑。以亿牛云的数据采集场景为例，从外部站点批量采集训练语料时，IP 限流和并发频率是首要约束：

import requests
# 亿牛云 API 代理提取出口 IP
api_url = "http://ip.16yun.cn:817/myip/pl/<ORDER_ID>/?s=<SIGN>&u=<USER>&format=json"
proxy_info = requests.get(api_url, timeout=10).json()[0]
proxy = f"http://{proxy_info['ip']}:{proxy_info['port']}"
# 批量采集语料
resp = requests.get("https://corpus.example.com/batch",
                    proxies={"http": proxy, "https": proxy}, timeout=30)

采集返回 403 需检查白名单配置，返回 429 需降低请求频率。稳定的数据采集是 LLM 训练的第一道基础设施保障。六、踩坑速查

陷阱	根因	解法
Loss 变 NaN	梯度爆炸	梯度裁剪 max_norm=1.0 + Pre-Norm
Loss 不下降	学习率不匹配	前 200 步线性预热，再用余弦衰减
生成重复	采样策略单一	Top-K + 温度缩放
显存溢出	batch/seq 过大	减小 batch_size 或用梯度累积
维度报错	d_model % n_heads ≠ 0	确保 d_model 能被 n_heads 整除

七、总结手写一遍 Transformer，收获的是对每个组件的深层理解：Embedding：离散 token → 连续空间，乘 √d_model 做尺度对齐RoPE：旋转编码位置信息到 Q/K，支持相对位置外推Attention：Q·Kᵀ 算相似度 → softmax 归一化 → 加权 V，计算复杂度 O(n²·d)Causal Mask：上三角填 -inf，实现自回归生成Pre-Norm 残差：信息直通通道，让深层网络可训练权重共享：Embedding 和 LM Head 共享权重，参数减半当面试官问"Attention 的计算复杂度为什么是 O(n²)"时，你能立刻回答：因为 Q·Kᵀ 那一步矩阵乘法的维度是 [seq, d] × [d, seq]。这种理解层次，是调用 from_pretrained() 永远无法获得的。

原创声明：本文系作者授权腾讯云开发者社区发表，未经许可，不得转载。

如有侵权，请联系 cloudcommunity@tencent.com 删除。

huggingface-transformers