从零实现 LLaMA 架构：构建轻量级大语言模型

从零实现 LLaMA 架构：构建轻量级大语言模型

大语言模型（LLM）的爆发式发展让 LLaMA 系列成为开源社区的焦点。Meta 推出的 LLaMA 凭借简洁的架构设计和高效的性能，成为了很多自研大模型的基准。不同于传统 Transformer，LLaMA 做了一系列关键优化：用 RMSNorm 替代 LayerNorm、采用 SwiGLU 激活的 FeedForward、引入旋转位置编码（RoPE）以及 Pre-Norm 架构等。

我们将从零开始，拆解 LLaMA 的核心设计，并通过可运行的代码实现一个轻量级的 LLaMA-like 模型，帮你理解大模型的底层原理。

一、LLaMA 核心设计亮点

先梳理 LLaMA 相对于经典 Transformer 的核心改进，这也是后续代码实现的重点：

优化点	传统 Transformer	LLaMA 设计	优势
归一化层	LayerNorm（含均值中心化 + 偏置）	RMSNorm（仅均方根归一化）	计算更快，训练稳定性相当
前馈网络激活	ReLU/GELU + 单线性层	SwiGLU（门控激活）	提升模型表达能力
位置编码	绝对位置编码	旋转位置编码（RoPE）	更好的长序列泛化能力
归一化位置	Post-Norm（注意力 / FFN 后）	Pre-Norm（注意力 / FFN 前）	训练更稳定，梯度传播更顺畅
线性层偏置	带 bias	无 bias	减少参数规模，提升推理速度

二、代码架构总览

我们将模型拆解为 5 个职责清晰的核心文件，从基础组件到完整模型再到测试，层层递进：

文件名称	核心功能
config.py	模型超参数管理（类型安全的 dataclass）
layers.py	基础层实现（RMSNorm、SwiGLU FeedForward、RoPE）
attention.py	因果自注意力层（集成 RoPE+Flash Attention）
model.py	Transformer 块封装 + 完整 LLM 模型
main.py	前向传播测试 + 自回归文本生成

三、逐模块解析代码

3.1 配置模块：config.py

模型超参数是大模型的'骨架'，用 dataclass 可以简洁、类型安全地管理这些参数，方便后续扩展和修改：

from dataclasses import dataclass

@dataclass
class LLMConfig:
    vocab_size: int = 32000          # 词表大小（原版 LLaMA 为 32000）
    hidden_size: int = 1024          # 隐藏层核心维度 (dim)
    num_layers: int = 12             # Transformer 层数
    num_heads: int = 16              # 注意力头数
    intermediate_size: int = 2816    # FFN 中间层维度（通常是 hidden_size 的 8/3 倍）
    max_seq_len: int = 2048          # 最大上下文长度
    rms_norm_eps: float = 1e-5       # RMSNorm 防止除零的极小值
    dropout: float = 0.1             # Dropout 概率

关键参数解读：

• intermediate_size：FFN 中间层维度选择 hidden_size * 8/3 是 LLaMA 的经验值，平衡模型表达能力和参数量；
• rms_norm_eps：极小值（1e-5）避免均方根计算时除以 0；
• max_seq_len：决定模型能处理的最长文本长度，也影响 RoPE 频率矩阵的预计算范围。

3.2 基础层模块：layers.py

这是模型的'基础组件库'，实现了 LLaMA 最核心的三个基础层：RMSNorm、SwiGLU FeedForward、RoPE。

3.2.1 均方根归一化（RMSNorm）

RMSNorm 是 LLaMA 的核心优化之一，数学公式为 $y = \frac{x}{\sqrt{E[x^2] + \epsilon}} \times \gamma$。相比 LayerNorm，它去掉了均值中心化和偏置项，计算效率更高：

# RMSNorm 实现
import torch
import torch.nn as nn

class RMSNorm(nn.Module):
    """均方根归一化 (Root Mean Square Normalization)"""
    def __init__(self, dim: int, eps: float = 1e-5):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))  # 可学习的缩放权重 γ

    def _norm(self, x):
        # 计算最后一维的均方根，keepdim 保证广播维度匹配
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)

    def forward(self, x):
        # 转为 float 计算避免精度问题，再转回原类型
        output = self._norm(x.float()).type_as(x)
        return output * self.weight

3.2.2 SwiGLU 前馈网络

LLaMA 的 FeedForward 层使用 SwiGLU 激活（替代传统 GELU），公式为 $SwiGLU(x) = Swish(xW_1) \otimes (xW_3)W_2$。其中 $\otimes$ 是逐元素相乘，SwiGLU 通过'门控机制'提升模型的非线性表达能力：

# FeedForward (SwiGLU) 实现
class FeedForward(nn.Module):
    """采用 SwiGLU 激活的基于门控的前馈神经网络"""
    def __init__(self, config):
        super().__init__()
        self.w1 = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.w2 = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
        self.w3 = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)

    def forward(self, x):
        # SwiGLU 核心逻辑：(Swish(xW1) * xW3) W2
        return self.w2(F.silu(self.w1(x)) * self.w3(x))

设计细节：所有线性层都去掉了 bias，这是 LLaMA 的核心设计之一，减少参数的同时提升训练稳定性。

3.2.3 旋转位置编码（RoPE）

RoPE 的核心是将位置信息编码为复数旋转角度，让 Query/Key 在注意力计算时随位置'旋转'，既保留绝对位置信息，又具备相对位置的泛化能力。

# RoPE 实现
def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0):
    """预计算 RoPE 的频率矩阵（复数形式）"""
    # 计算基础频率：1 / theta^(2i/dim)，i 为维度索引
    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
    # 生成时间步（序列位置）t 的外积，shape: [end, dim//2]
    t = torch.arange(end, device=freqs.device)
    freqs = torch.outer(t, freqs).float()
    # 转为复数（极坐标）：模为 1，角度为 freqs
    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)
    return freqs_cis

def apply_rotary_emb(xq: torch.Tensor, xk: torch.Tensor, freqs_cis: torch.Tensor):
    """将 RoPE 应用到 Query/Key 上"""
    # 将 Q/K 重塑为复数形式：[B, T, n_heads, head_dim] → [B, T, n_heads, head_dim//2, 2]
    xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
    xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
    # 广播频率矩阵到 Q/K 的形状：[T, dim//2] → [1, T, 1, dim//2]
    freqs_cis = freqs_cis.unsqueeze(0).unsqueeze(2)
    # 复数乘法 = 旋转操作，再转回实数形式并展平
    xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)
    xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)
    return xq_out.type_as(xq), xk_out.type_as(xk)

核心逻辑：

• precompute_freqs_cis：提前计算所有位置的旋转频率（复用性高，无需每次前向都计算）；
• apply_rotary_emb：将 Q/K 按两个维度为一组拆分为复数，与频率矩阵相乘实现'旋转'，再转回实数。

3.3 注意力模块：attention.py

因果自注意力是 Transformer 的核心，LLaMA 的注意力层做了两大优化：QKV 合并映射（工程高效）、集成 Flash Attention（PyTorch 2.0 + 内置）。

import math
import torch
import torch.nn as nn
from layers import apply_rotary_emb

class CausalSelfAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        assert config.hidden_size % config.num_heads == 0
        self.n_heads = config.num_heads
        self.head_dim = config.hidden_size // config.num_heads
        # 合并 Q/K/V 的线性映射（工程高效，也可拆分支持 GQA/MQA）
        self.wqkv = nn.Linear(config.hidden_size, 3 * config.hidden_size, bias=False)
        self.wo = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
        self.attn_dropout = nn.Dropout(config.dropout)
        self.resid_dropout = nn.Dropout(config.dropout)

    def forward(self, x, freqs_cis, mask=None):
        B, T, C = x.size()  # B: batch_size, T: seq_len, C: hidden_size
        # 合并计算 Q/K/V，再拆分
        qkv = self.wqkv(x)
        q, k, v = qkv.split(C, dim=2)
        # 重塑为按头划分的形状：[B, T, n_heads, head_dim]
        q = q.view(B, T, self.n_heads, self.head_dim)
        k = k.view(B, T, self.n_heads, self.head_dim)
        v = v.view(B, T, self.n_heads, self.head_dim)
        # 应用 RoPE 位置编码
        q, k = apply_rotary_emb(q, k, freqs_cis[:T])
        # 转置为 [B, n_heads, T, head_dim]，适配 PyTorch 的 scaled_dot_product_attention
        q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
        # 调用 PyTorch 内置的缩放点积注意力（集成 Flash Attention，速度/显存优化）
        y = torch.nn.functional.scaled_dot_product_attention(
            q, k, v,
            attn_mask=mask,
            dropout_p=self.attn_dropout.p if self.training else 0.0,
            is_causal=True if mask is None else False  # 因果掩码，防止看到未来 token
        )
        # 拼接各头结果，转回 [B, T, C]
        y = y.transpose(1, 2).contiguous().view(B, T, C)
        return self.resid_dropout(self.wo(y))

关键优化：

• scaled_dot_product_attention：PyTorch 2.0 + 内置接口，自动启用 Flash Attention，大幅降低显存占用、提升计算速度；
• is_causal=True：自动生成因果掩码，避免手动构造掩码矩阵，代码更简洁。

3.4 Transformer 块与主模型：model.py

将注意力层和 FFN 层组合成 Transformer 块，再堆叠为完整的 LLM 模型，核心是 Pre-Norm 架构和残差连接。

import torch
import torch.nn as nn
from config import LLMConfig
from layers import RMSNorm, FeedForward, precompute_freqs_cis
from attention import CausalSelfAttention

class TransformerBlock(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.attention = CausalSelfAttention(config)
        self.feed_forward = FeedForward(config)
        # LLaMA 核心：Pre-Norm（归一化在注意力/FFN 之前）
        self.attention_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.ffn_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

    def forward(self, x, freqs_cis):
        # 注意力层：残差连接 + Pre-Norm
        h = x + self.attention(self.attention_norm(x), freqs_cis)
        # FFN 层：残差连接 + Pre-Norm
        out = h + self.feed_forward(self.ffn_norm(h))
        return out

class LLM(nn.Module):
    def __init__(self, config: LLMConfig):
        super().__init__()
        self.config = config
        self.vocab_size = config.vocab_size
        # 词嵌入层（无 bias，对齐 LLaMA 设计）
        self.tok_embeddings = nn.Embedding(config.vocab_size, config.hidden_size)
        # 堆叠 Transformer 块
        self.layers = nn.ModuleList([TransformerBlock(config) for _ in range(config.num_layers)])
        # 输出归一化 + 分类头（无 bias）
        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.output = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        # 预计算 RoPE 频率矩阵，注册为 buffer（不参与梯度更新）
        freqs_cis = precompute_freqs_cis(config.hidden_size // config.num_heads, config.max_seq_len)
        self.register_buffer("freqs_cis", freqs_cis, persistent=False)

    def forward(self, tokens, targets=None):
        B, T = tokens.size()
        # 词嵌入：[B, T] → [B, T, C]
        h = self.tok_embeddings(tokens)
        # 逐层前向传播
        for layer in self.layers:
            h = layer(h, self.freqs_cis[:T])
        # 输出处理
        h = self.norm(h)
        logits = self.output(h)  # [B, T, vocab_size]
        # 计算损失（如果传入 targets）
        loss = None
        if targets is not None:
            loss = nn.functional.cross_entropy(logits.view(-1, self.vocab_size), targets.view(-1))
        return logits, loss

核心设计：

• register_buffer：将 RoPE 频率矩阵注册为非训练参数，避免每次前向都重新计算；
• Pre-Norm：归一化层在注意力 / FFN 之前，相比 Post-Norm，训练时梯度更稳定，无需额外的初始化技巧；
• 残差连接：每个子层（注意力 / FFN）的输出都与输入相加，保证梯度能有效传播到浅层。

四、实战运行与结果解读

将所有代码文件放在同一目录，运行 main.py，输出示例如下：

正在初始化 LLM 模型 (类 LLaMA 架构)...
模型参数量：0.85 M
前向传播测试：Loss = 6.9078, Logits Shape = torch.Size([1, 5, 1000])
开始生成文本...
原始输入：[10, 20, 30, 40, 50]
生成结果：[10, 20, 30, 40, 50, 88, 123, 45, 789, 23, 90, 111, 56, 89]

结果分析：

1. 模型参数量约 0.85M，属于轻量级，可在 CPU 上快速测试；
2. 初始 Loss≈6.9，符合预期（随机初始化的模型，Loss 接近 ln(vocab_size)=ln(1000)≈6.9）；
3. 生成的 token 序列是随机的（模型未训练），但验证了自回归生成逻辑的正确性。

五、总结

本文从 LLaMA 的核心设计出发，拆解并实现了一个轻量级的 LLaMA-like 模型，覆盖了 RMSNorm、SwiGLU、RoPE、因果自注意力等关键组件。

大模型看似复杂，但本质是'简单组件的有序组合'。掌握这些核心设计，就能理解大模型的底层逻辑，为后续的模型训练、优化和部署打下基础。

如果你想进一步深入，可以尝试：

1. 训练模型：用小数据集（如 WikiText）训练模型，观察 Loss 的下降趋势；
2. 扩展参数：将 hidden_size 调至 4096、num_layers 调至 32，实现原版 LLaMA 7B 的架构；
3. 部署推理：将模型导出为 ONNX/TensorRT，提升推理速度。