从零实现 LLaMA 架构：一步步构建轻量级大语言模型
大语言模型（LLM）的爆发式发展让 LLaMA 系列模型成为开源社区的焦点 ——Meta 推出的 LLaMA 以简洁的架构设计和高效的性能，成为很多自研大语言模型的基准。不同于传统 Transformer，LLaMA 做了诸多关键优化：用 RMSNorm 替代 LayerNorm、SwiGLU 激活的 FeedForward、旋转位置编码（RoPE）、Pre-Norm 架构等。
本文将从零开始，拆解 LLaMA 的核心设计，并通过可运行的代码实现一个轻量级的 LLaMA-like 模型，帮助你理解大模型的底层原理。
目录
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可以简洁、类型安全地管理这些参数，方便后续扩展和修改：
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，它去掉了均值中心化和偏置项，计算效率更高：
# 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 通过 “门控机制” 提升模型的非线性表达能力：
# 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）的输出都与输入相加，保证梯度能有效传播到浅层。
生成逻辑解读：
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 架构：一步步构建轻量级大语言模型
如需源码，请在评论区下留言，作者会逐个回复，制作不易，请各位看官老爷点个赞和收藏！！！
