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

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

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

[一、LLaMA 核心设计亮点](#一、LLaMA 核心设计亮点)

二、代码架构总览

三、逐模块解析代码

[3.1 配置模块：config.py](#3.1 配置模块：config.py)

[3.2 基础层模块：layers.py](#3.2 基础层模块：layers.py)

[3.2.1 均方根归一化（RMSNorm）](#3.2.1 均方根归一化（RMSNorm）)

[3.2.2 SwiGLU 前馈网络](#3.2.2 SwiGLU 前馈网络)

[3.2.3 旋转位置编码（RoPE）](#3.2.3 旋转位置编码（RoPE）)

[3.3 注意力模块：attention.py](#3.3 注意力模块：attention.py)

[3.4 Transformer 块与主模型：model.py](#3.4 Transformer 块与主模型：model.py)

四、实战运行与结果解读

五、总结

一、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可以简洁、类型安全地管理这些参数，方便后续扩展和修改：

python 复制代码

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=E[x2]+ϵx×γ相比 LayerNorm，它去掉了均值中心化 和偏置项，计算效率更高：

python 复制代码

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

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(xW1)⊗(xW3)W2其中⊗是逐元素相乘，SwiGLU 通过 "门控机制" 提升模型的非线性表达能力：

python 复制代码

#  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 在注意力计算时随位置 "旋转"，既保留绝对位置信息，又具备相对位置的泛化能力。

python 复制代码

#  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 + 内置）。

python 复制代码

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 架构 和残差连接。

python 复制代码

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）的输出都与输入相加，保证梯度能有效传播到浅层。

生成逻辑解读：

temperature：温度越高，生成的随机性越强（logits 除以 temperature 后，概率分布更平缓）；
torch.multinomial：多项式采样（相比 argmax 的 "贪心采样"，生成结果更丰富）；
序列裁剪：每次生成前裁剪序列到max_seq_len，避免超出模型的上下文长度限制。

四、实战运行与结果解读

将所有代码文件放在同一目录，运行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]

结果分析：

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

五、总结

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

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

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

训练模型：用小数据集（如 WikiText）训练模型，观察 Loss 的下降趋势；
扩展参数：将hidden_size调至 4096、num_layers调至 32，实现原版 LLaMA 7B 的架构；
部署推理：将模型导出为 ONNX/TensorRT，提升推理速度。# 从零实现 LLaMA 架构：一步步构建轻量级大语言模型

如需源码，请在评论区下留言，作者会逐个回复，制作不易，请各位看官老爷点个赞和收藏！！！