Software 2.0

代码学习: GPT-2 vs LLaMA 3

admin -

学习LLM工作原理有两种方法。第一种是阅读"Attention Is All You Need"，盯着编码器-解码器图看30分钟直到你放弃。第二种是打开Karpathy的nanoGPT/model.py——300行PyTorch代码——从头到尾阅读，理解每一行，然后打开LLaMA 3的model.py，看看究竟改了什么以及为什么。

第二种方法更好的原因很简单：在计算机科学中，人脑不记忆散文，它记忆结构。当你阅读真实代码时，你内化结构。当你只读散文时，你相信自己一周后还记得——而通常你不会。

这篇文章是对两个经典实现的注释式阅读：

  • nanoGPT/model.py，由Karpathy编写，用约300行干净的PyTorch忠实地复现了GPT-2（2019年）。
  • meta-llama/llama/model.py，LLaMA的参考实现，遵循原始论文以及2022至2024年间分批引入的架构变更。

我不会从头重新推导注意力机制——我假设你已经见过softmax(QK^T / sqrt(d_k)) V公式。本文的重点是忠实GPT-2与现代transformer之间发生了什么变化，为什么变化，以及这对想要部署、优化或定制这些模型的人意味着什么。

这是LLM内部系列4篇文章中的第一篇：

  1. 内部模型架构：理论与实践结合代码
  2. 部署真实LLM：生产实践
  3. 高级技术（KV缓存、分页注意力、量化、推测解码）：对比分析
  4. Transformer的演进（以及未来方向）：架构/概念层面

留出30分钟，打开一个标签页看nanoGPT代码，另一个看LLaMA，让我们开始吧。

1、骨架：两者相同

在看差异之前，值得注意2019年到2026年间没有改变的东西：decoder-only transformer块的总体骨架。它仍然是：

input → embedding
loop N times:
    block = norm → attention → residual
            norm → mlp → residual
output = final norm → vocabulary projection

这个"norm → sub-block → residual"模式如此基础，以至于每个现代变体都遵循它。改变的是哪种norm、哪种attention和哪种MLP。

组织本文其余部分的结构图：

四个刻意变更，每一个都是针对具体瓶颈的响应：

  1. LayerNorm → RMSNorm — 训练稳定性和推理速度。
  2. 位置嵌入 → RoPE — 更长上下文的可扩展性。
  3. MHA → GQA — KV缓存减少（8倍）。
  4. GELU → SwiGLU — 门控提升MLP表达能力。

让我们逐一走过这四个变更，每个都附带前后代码对比。

2、LayerNorm vs. RMSNorm

这是没人要求但所有人都在用的优化。

2.1 GPT-2代码

在nanoGPT/model.py中，LayerNorm的实现实际上是一个8行的类：

class LayerNorm(nn.Module):
    """ LayerNorm but with an optional bias. PyTorch doesn't support simply bias=False """
    def __init__(self, ndim, bias):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(ndim))
        self.bias = nn.Parameter(torch.zeros(ndim)) if bias else None

    def forward(self, input):
        return F.layer_norm(input, self.weight.shape, self.weight, self.bias, 1e-5)

F.layer_norm的作用：对最后一个维度的每个元素，计算均值（μ）、标准差（σ），归一化为(x - μ) / σ，然后应用weight * normalized + bias。保证输出具有约0均值和约1方差，带有可学习的weight（缩放）和bias（偏移）。

2.2 LLaMA代码

class RMSNorm(torch.nn.Module):
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def _norm(self, x):
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)

    def forward(self, x):
        output = self._norm(x.float()).type_as(x)
        return output * self.weight

差异：

  • 不计算均值。 只计算均方根（rsqrt(mean(x²))）。这就是名称中"RMS"的含义。
  • 没有bias。 只有weight。

2.3 为什么重要

原始理由（RMSNorm论文，Zhang & Sennrich 2019）是经验性和实用性的：

LayerNorm的大部分好处来自重新缩放（除以标准差），而不是重新居中（减去均值）。

如果你能去掉重新居中并保留大部分质量，你可以获得7%到64%的速度提升，具体取决于实现和硬件。在训练中，这会在所有层和所有token上累积。在推理中，它减少了内存带宽使用。

附注：LLaMA还使用了pre-norm（在attention/MLP之前、residual内部进行归一化），而不是post-norm（之后归一化）。Pre-norm不是LLaMA的发明——自从GPT-2以来它已经成为共识做法，nanoGPT中也是这样：

# nanoGPT, same pattern:
def forward(self, x):
    x = x + self.attn(self.ln_1(x))   # pre-norm
    x = x + self.mlp(self.ln_2(x))    # pre-norm
    return x

Pre-norm的优势：梯度通过残差连接流动更顺畅。在32+层的模型中，这是真实的训练稳定性差异。

2.4 运营影响

对于部署LLaMA风格模型的人来说：每一层少了一个操作。在32层模型中，每个token节省了32次均值归约，包括prefill和decode阶段。单独看不惊人，但它与后续所有内容相辅相成——LLaMA的哲学是"去掉没有帮助的东西"。

3、位置嵌入 vs. RoPE

这是影响最深的替换，是一个解锁上下文窗口的变更，也是仅看代码而不理解数学时最不明显的一个。

3.1 GPT-2代码

在nanoGPT/model.py中：

self.transformer = nn.ModuleDict(dict(
    wte = nn.Embedding(config.vocab_size, config.n_embd),  # token embeddings
    wpe = nn.Embedding(config.block_size, config.n_embd),  # position embeddings
    ...
))

# in forward:
pos = torch.arange(0, t, dtype=torch.long, device=device)  # (t,)
tok_emb = self.transformer.wte(idx)  # (b, t, n_embd)
pos_emb = self.transformer.wpe(pos)  # (t, n_embd)
x = tok_emb + pos_emb                # sum!

翻译：GPT-2有一个位置嵌入表，包含block_size个条目（最大序列长度，通常为1024）。每个位置变成一个可学习的向量，简单地加到token嵌入上。

问题：如果你用block_size=1024训练模型，它不知道位置1025是什么。表中没有那个条目。想要扩展上下文到4k或32k？重新训练或者放弃。

3.2 LLaMA代码

def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0):
    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
    t = torch.arange(end, device=freqs.device)
    freqs = torch.outer(t, freqs).float()
    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)  # complex: cos+isin
    return freqs_cis

def apply_rotary_emb(xq, xk, freqs_cis):
    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))
    freqs_cis = reshape_for_broadcast(freqs_cis, xq_)
    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)

不要被复数运算吓到。直觉是简单而美妙的。

3.3 为什么重要

RoPE（旋转位置嵌入，Su等人2021）将Q/K向量中每对相邻维度视为一个复数，并按与token位置成正比的角度旋转该复数。旋转频率在不同维度对之间变化——某些对旋转快（捕获局部模式），其他对旋转慢（捕获全局模式）。

这种构造产生了三个本质属性：

  1. 位置信息被注入到Q和K中，而不是输入中。 不再有"位置表"大小不够的问题。输入只获得token嵌入；位置在注意力计算时进入。
  2. 两个token之间的关系仅取决于相对距离。 当你计算Q_i · K_j（旋转后的点积）时，结果取决于差值i - j，而不是绝对值。这就是从一个简单操作中涌现的相对位置编码。
  3. 外推到更大上下文是可行的。 如果你用theta=10000和4k上下文训练，可以通过调整theta在8k上下文上运行推理（还有YaRN和Position Interpolation等技术可以很好地实现这一点）。LLaMA 3.1将theta提高到500000正是为了启用长上下文窗口（128k token）。在GPT-2中，这根本没有途径实现——表中没有相应条目。

3.4 运营影响

对于部署：RoPE是实现现代应用所需的128k上下文token的关键。没有它，你就被限制在训练大小。对于定制：如果你想扩展微调后LLaMA的上下文窗口，调整theta而不是从头重新训练。主要部署框架（vLLM、TensorRT-LLM、SGLang）都有一流的RoPE支持。

4、多头注意力 vs. 分组查询注意力

这是真正重要的优化。

4.1 GPT-2代码

在nanoGPT/model.py中：

class CausalSelfAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        # K, Q, V projections in a single linear, then split
        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)
        self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)
        self.n_head = config.n_head
        self.n_embd = config.n_embd

    def forward(self, x):
        B, T, C = x.size()
        q, k, v = self.c_attn(x).split(self.n_embd, dim=2)
        # split into heads
        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
        # attention
        y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        y = y.transpose(1, 2).contiguous().view(B, T, C)
        y = self.c_proj(y)
        return y

每个头都有自己的Q、K、V。12个头意味着12个Q矩阵、12个K、12个V。

4.2 LLaMA代码

class Attention(nn.Module):
    def __init__(self, args: ModelArgs):
        super().__init__()
        self.n_kv_heads = args.n_kv_heads if args.n_kv_heads is not None else args.n_heads
        self.n_local_heads = args.n_heads
        self.n_local_kv_heads = self.n_kv_heads
        self.n_rep = self.n_local_heads // self.n_local_kv_heads
        self.head_dim = args.dim // args.n_heads

        self.wq = nn.Linear(args.dim, args.n_heads * self.head_dim, bias=False)
        self.wk = nn.Linear(args.dim, self.n_kv_heads * self.head_dim, bias=False)
        self.wv = nn.Linear(args.dim, self.n_kv_heads * self.head_dim, bias=False)
        self.wo = nn.Linear(args.n_heads * self.head_dim, args.dim, bias=False)

    def forward(self, x, freqs_cis, mask, cache_k, cache_v, start_pos):
        bsz, seqlen, _ = x.shape
        xq = self.wq(x).view(bsz, seqlen, self.n_local_heads, self.head_dim)
        xk = self.wk(x).view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)
        xv = self.wv(x).view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)

        xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs_cis)
        # ... cache, repeat KV for multi-query, attention ...

关键差异：

self.wq = nn.Linear(dim, n_heads * head_dim)         # n_heads queries
self.wk = nn.Linear(dim, n_kv_heads * head_dim)      # n_kv_heads keys (fewer!)
self.wv = nn.Linear(dim, n_kv_heads * head_dim)      # n_kv_heads values (fewer!)

在LLaMA 3 70B中，n_heads=64但n_kv_heads=8。K和V矩阵少了八倍。每8个查询头共享同一对K/V。这就是GQA（分组查询注意力，Ainslie等人2023）。

4.3 为什么重要

要理解其影响，你需要看自回归推理——所有节省都在这里实现。

当逐个token生成时，模型必须记住所有先前token的K和V来计算注意力。自2018年以来的标准解决方案是KV缓存：存储K和V一次，每个新token重复使用。

缓存增长方式：每层(batch, n_kv_heads, seq_len, head_dim)，乘以2（K和V），乘以层数，乘以每个元素的字节数（fp16为2，int8为1）。

对于LLaMA 3 70B，4k token上下文，batch为1，fp16：

  • 没有GQA（n_kv_heads = n_heads = 64）：约10 GB KV缓存。
  • 有GQA（n_kv_heads = 8）：约1.25 GB。少八倍。

在更大batch或更长上下文中，这个差异就是"能放进一个GPU"和"立即OOM"之间的分水岭。

接下来是一个改变你对部署一切认知的真相：

自回归推理是内存带宽受限的，而非计算受限。

在每个decode步骤，GPU读取整个缓存（而且缓存还在增长）。在任何现代GPU上，读取100 GB内存比计算100 GFLOPs慢几个数量级。例如，DGX Spark提供1 PFLOP的FP4算力（惊人的计算能力），但只有273 GB/s的内存带宽（统一LPDDR5x）——正是这种不平衡被LMSYS观察到限制了大模型的吞吐量。

GQA没有让模型更好。它让模型在现实的VRAM和带宽预算下可以部署。

4.4 运营影响

选择要部署的模型时，看n_heads / n_kv_heads比率。在LLaMA 3 8B中为4（32/8），在70B中为8（64/8）。比率越大，KV缓存越小，你能服务更多batch，每个GPU吞吐量越高。这是相同"大小"的现代模型交付截然不同的生产吞吐量的主要原因之一。

5、GELU vs SwiGLU

这是多花50%参数但值得的门控。

5.1 GPT-2代码

class MLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.c_fc    = nn.Linear(config.n_embd, 4 * config.n_embd, bias=config.bias)
        self.gelu    = nn.GELU()
        self.c_proj  = nn.Linear(4 * config.n_embd, config.n_embd, bias=config.bias)

    def forward(self, x):
        x = self.c_fc(x)
        x = self.gelu(x)
        x = self.c_proj(x)
        return x

经典方案：扩展到4x大小，GELU，再缩回来。两个线性层和一个非线性。

5.2 LLaMA代码

class FeedForward(nn.Module):
    def __init__(self, dim, hidden_dim, multiple_of, ffn_dim_multiplier):
        super().__init__()
        hidden_dim = int(2 * hidden_dim / 3)
        if ffn_dim_multiplier is not None:
            hidden_dim = int(ffn_dim_multiplier * hidden_dim)
        hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)

        self.w1 = nn.Linear(dim, hidden_dim, bias=False)  # gate
        self.w2 = nn.Linear(hidden_dim, dim, bias=False)  # down
        self.w3 = nn.Linear(dim, hidden_dim, bias=False)  # up

    def forward(self, x):
        return self.w2(F.silu(self.w1(x)) * self.w3(x))

三个线性层而不是两个。核心操作：

silu(gate(x)) * up(x)  →  down(...)

其中silu(z) = z * sigmoid(z)，也称为Swish。*是逐元素乘法。

5.3 为什么重要

SwiGLU（Shazeer 2020，GLU变体改进Transformer）是一系列门控激活函数。直觉：不是应用一个简单的非线性，而是有一个"门"路径（控制）和一个"上"路径（值），门选择/加权多少值通过。

经验上，SwiGLU比GELU提升困惑度一个小但一致的量（约1-2%）。听起来微不足道，直到你记住在这个行业中1%的困惑度通常意味着数百万美元的计算成本。

代价：三个线性层而不是两个。为了保持相同的MLP总参数量，LLaMA将hidden_dim减少到经典GELU的2/3，然后乘以ffn_dim_multiplier（LLaMA 3中为1.3），并舍入到256的倍数以实现硬件对齐。所有这些账目计算都在上面的代码中。

5.4 运营影响

对于量化模型的人来说：SwiGLU在量化下的行为与GELU不同。具体来说，2023-2024年成为标准的AWQ和GPTQ技术是在LLaMA模型上开发和验证的——SwiGLU的门控是它们利用的一部分。AWQ文档强调LLaMA风格基准测试并非巧合。在经典GELU的GPT-2风格模型中，同样的技术也能工作，但量化SwiGLU损失更少质量，因为门提供了一条"替代"路径来保留动态范围。

这种细节就是"理解model.py"从学术练习变成具体运营优势的地方。

6、悄然改变的小决策

除了四大变更之外，nanoGPT和LLaMA之间还有一系列值得提及的较小变化。每一个单独都不大；但它们加在一起也很可观。

6.1 线性层没有bias

原始GPT-2在几乎所有nn.Linear上都有bias=True。LLaMA在所有地方使用bias=False。两方面的理由：

  1. 内存：每个线性层节省一个向量。在大模型中，这相当于在15T token上节省数分钟的训练时间。
  2. 稳定性：在某些情况下，bias可能增加对异常值的敏感性。没有bias，网络被迫学习通过原点的函数。

这个决定有些争议——一些论文表明在小模型中添加bias略有帮助。但在大规模上，bias=False是正确选择已成为共识。

6.2 分词器：BPE vs. SentencePiece

GPT-2使用BPE（字节对编码），主要在英文文本上训练。LLaMA使用SentencePiece，词汇量大得多（32k → LLaMA 3中的128k），多语言覆盖更好。

这不是架构变更，但影响其他一切——任何"每秒token数"基准测试都取决于"一个token"在你的分词中的含义。LLaMA 3的128k词汇量在相同上下文窗口中比32k的LLaMA 2多装约25%的文本。

6.3 LM头与输入嵌入的权重绑定

在GPT-2中：

# nanoGPT
self.transformer.wte.weight = self.lm_head.weight  # weight tying!

将token投影到向量的矩阵实际上是将向量投影回logits的矩阵的转置。节省了模型三分之一的参数（vocab_size * n_embd很大）。

LLaMA不做权重绑定。embed_tokens.weight和lm_head.weight矩阵是分开的。这个决定增加了参数，但为量化和微调提供了更大的灵活性。这是一种刻意的"退步"反而变得有效的情况。

7、阅读nanochat：nanoGPT之后是什么

2025年10月，Karpathy发布了nanochat，描述为"100美元能买到的最好ChatGPT"，正式作为nanoGPT的继任者（他在README中将nanoGPT标记为已弃用）。

nanochat做了nanoGPT没做的事：

  • 从头训练模型带SFT（监督微调）和类RLHF——不只是语言模型预训练。
  • 使用现代架构（RMSNorm、RoPE、GQA、SwiGLU）而非忠实GPT-2。
  • 单一复杂度旋钮：传入--depth=N，其余超参数以"计算最优"模式自动计算。
  • 在适度的硬件上运行：在8×H100节点上约48美元约2小时获得GPT-2级别的结果（2019年，训练GPT-2花费约43,000美元）。

对于想要阅读完整现代代码库的人来说，nanochat/是当前的推荐。但要从nanoGPT/model.py开始——它仍然是无可超越的教学起点。

8、这对你意味着什么

来自本文的五个运营影响，按通常遇到的顺序排列：

1. 内存带宽主导自回归推理。 不是计算。选择用于部署的GPU时，先看GB/s再看TFLOPS。这就是为什么1 PFLOP的DGX Spark的吞吐量与计算能力只有其一小部分但带宽更多的系统相似——它在大模型上是内存受限的。这也是为什么MI300X（5.3 TB/s）在某些场景下比H100（3.35 TB/s）更适合推理，尽管H100有更多计算能力。

2. KV缓存随序列增长；它不是免费的。 如果你的应用使用32k+上下文，先计算缓存再假设能放得下。低n_kv_heads的模型（LLaMA 3 70B为8）只需要1/8的空间——刻意的选择。

3. RoPE允许无需重新训练的上下文扩展。 想要更大窗口的模型？调整theta（或使用YaRN/Position Interpolation）。在为长上下文微调之前，先试这个——通常就能解决。

4. SwiGLU + RMSNorm是量化效果最好的架构。 AWQ和GPTQ是针对这种架构开发的。LLaMA风格模型量化到INT4/INT8损失极小；GPT-2风格模型根据方法可能损失更多。

5. model.py可以装进你的脑海。 nanoGPT约300行，LLaMA参考实现约500行。一个下午就能读完，接下来一周你做的每个架构决策都基于阅读而非直觉。这是我在应用AI中知道的最高ROI。

原文链接: Reading model.py From Inside: From GPT-2 to LLaMA 3 in 500 Lines

汇智网翻译整理，转载请标明出处