大模型底层架构与AI源码深度剖析

风骏时光牛马
16 阅读3分钟

大模型底层架构与AI源码深度剖析

一、大模型的核心架构

现代大语言模型（LLM）普遍基于 Transformer 架构，其核心思想是"注意力机制"——让模型在处理每个词时，能够动态权衡序列中其他词的重要性。整体结构可分为三大模块：

  1. 编码器：将输入文本映射为上下文相关的向量表示
  2. 解码器：基于上下文向量逐步生成输出
  3. 输出层：将隐状态映射为词表概率分布

以 GPT 系列为代表的"仅解码器"架构，通过自回归方式逐 token 生成，其内部由多层堆叠的 Decoder Block 组成，每层又包含多头自注意力（Multi-Head Self-Attention）和前馈网络（FFN）两个子层，并配以残差连接与层归一化。

二、自注意力机制源码实现

自注意力是整个架构的基石。以下用 PyTorch 演示单头注意力的核心计算流程：

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

class SelfAttention(nn.Module):
    def __init__(self, d_model, d_k):
        super().__init__()
        self.W_q = nn.Linear(d_model, d_k)
        self.W_k = nn.Linear(d_model, d_k)
        self.W_v = nn.Linear(d_model, d_k)
        self.d_k = d_k

    def forward(self, x, mask=None):
        Q = self.W_q(x)   # [B, T, d_k]
        K = self.W_k(x)
        V = self.W_v(x)
        # 缩放点积注意力: softmax(QK^T / sqrt(d_k)) * V
        scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.d_k ** 0.5)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))
        attn = F.softmax(scores, dim=-1)
        out = torch.matmul(attn, V)
        return out

x = torch.randn(2, 10, 64)  # batch=2, 序列长=10, 维度=64
attn = SelfAttention(d_model=64, d_k=64)
print(attn(x).shape)  # torch.Size([2, 10, 64])

代码要点：Q·K^T 计算两两 token 的相似度，除以 sqrt(d_k) 防止梯度消失，masked_fill 用于屏蔽未来信息（因果掩码），最后与 V 加权求和得到输出。

三、多头注意力与位置编码

单头只能关注单一子空间，因此实际模型将注意力拆分为多个头并行计算：

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        self.heads = nn.ModuleList(
            [SelfAttention(d_model, d_model // n_heads) for _ in range(n_heads)]
        )
        self.out = nn.Linear(d_model, d_model)

    def forward(self, x):
        heads = [h(x) for h in self.heads]
        return self.out(torch.cat(heads, dim=-1))

# 绝对位置编码: 用正弦/余弦函数注入位置信息
def positional_encoding(max_len, d_model):
    pe = torch.zeros(max_len, d_model)
    pos = torch.arange(max_len).unsqueeze(1).float()
    div = torch.exp(torch.arange(0, d_model, 2).float() * (-torch.log(torch.tensor(10000.0)) / d_model))
    pe[:, 0::2] = torch.sin(pos * div)
    pe[:, 1::2] = torch.cos(pos * div)
    return pe

由于注意力本身无序，位置编码决定了词的先后关系，正弦式编码能让模型泛化到更长序列。

四、前馈网络与残差归一化

每层注意力之后是两层的全连接网络，配合残差连接稳定深层训练：

class FeedForward(nn.Module):
    def __init__(self, d_model, d_ff):
        super().__init__()
        self.fc1 = nn.Linear(d_model, d_ff)
        self.fc2 = nn.Linear(d_ff, d_model)

    def forward(self, x):
        return self.fc2(F.gelu(self.fc1(x)))

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

    def forward(self, x):
        x = self.norm1(x + self.attn(x))  # 残差 + 层归一化
        return self.norm2(x + self.ffn(x))

残差连接让梯度能直达浅层，LayerNorm 将每层输出归一化，二者共同解决了深网络的训练不稳定问题。

五、训练与推理的关键机制

  • 训练：采用交叉熵损失 + 下一个 token 预测，配合混合精度、梯度累积和分布式并行加速。
  • 推理：KV Cache 缓存历史 K/V，避免重复计算，是吞吐优化的核心；贪心解码每次取概率最高 token，束搜索保留多个候选路径以提升质量。
  • 参数规模：从数亿（如 GPT-2）到千亿级（如 GPT-4 级别），参数量与训练数据量共同决定模型能力上限。

六、小结

理解大模型底层架构，本质是理解"注意力如何建模序列依赖、堆叠如何扩展表达能力、训练与推理如何权衡效率与质量"。上述源码演示了从单头注意力到完整 Decoder Block 的最小实现，掌握了这条主线，便能在源码层面读透主流开源模型的每一层结构。

海量精选技术文档和实战案例持续更新，敬请关注【风骏时光少年】

avatar
文章
阅读
粉丝