Llama 3架构深度解析:从源码看GQA和RMSNorm

Llama 3发布后,我花了整整一周时间读源码、做实验、测性能。这篇文章不是论文解读,而是从工程实践角度, 分析Llama 3相比Llama 2的关键改进,以及这些改进在实际部署中的影响。

先说结论

如果你没时间看完全文,这是我的核心发现:

1. GQA是最大亮点:推理吞吐量提升2.3倍,显存占用降75%,几乎无精度损失
2. 128K词表不是越大越好:中文效率提升明显,但embedding层吃掉0.5B参数
3. RMSNorm vs LayerNorm:速度快15%,训练稳定性相当
4. 实际部署:8B模型在A100上batch=8时,比Llama 2 7B快40%

背景:为什么要研究Llama 3

我们公司在用Llama 2 7B做对话服务,部署了4张A100。峰值QPS能到120,但显存占用很高, 只能跑batch_size=4。听说Llama 3推理更快,就想看看能不能升级。

读完Llama 3的论文和代码后,发现改动不大,主要就三个:

• Attention机制从MHA改成GQA
• 词表从32K扩到128K
• 训练数据从2T扩到15T

数据和词表是Meta的优势,我们学不来。但GQA这个架构改动,值得深入研究。

GQA:降低推理成本的核心技术

什么是GQA?

先看传统的Multi-Head Attention(MHA):

# Llama 2的MHA
n_heads = 32  # Q、K、V都有32个头
head_dim = 128
d_model = 4096

# 每个头都有独立的K、V投影
W_q = nn.Linear(d_model, n_heads * head_dim)  # (4096, 4096)
W_k = nn.Linear(d_model, n_heads * head_dim)  # (4096, 4096)
W_v = nn.Linear(d_model, n_heads * head_dim)  # (4096, 4096)

# KV Cache大小(推理时)
# [batch, seq_len, n_heads, head_dim]
# batch=1, seq=2048: 2048 * 32 * 128 * 2(K+V) * 2(bytes) ≈ 32MB per layer
# 32 layers: 1GB

Llama 3的GQA(Grouped-Query Attention):

# Llama 3的GQA  
n_heads = 32        # Q还是32个头
n_kv_heads = 8      # K、V只有8个头
head_dim = 128

W_q = nn.Linear(d_model, n_heads * head_dim)      # (4096, 4096)
W_k = nn.Linear(d_model, n_kv_heads * head_dim)  # (4096, 1024) ⬅️ 小了4倍!
W_v = nn.Linear(d_model, n_kv_heads * head_dim)  # (4096, 1024)

# KV Cache大小
# [batch, seq_len, n_kv_heads, head_dim]  
# batch=1, seq=2048: 2048 * 8 * 128 * 2 * 2 ≈ 8MB per layer
# 32 layers: 256MB  ⬅️ 降低了75%!

GQA的实现细节

关键问题:32个Q头怎么匹配8个KV头?

class GroupedQueryAttention(nn.Module):
    def __init__(self, d_model=4096, n_heads=32, n_kv_heads=8):
        super().__init__()
        self.n_heads = n_heads
        self.n_kv_heads = n_kv_heads
        self.n_rep = n_heads // n_kv_heads  # 32 / 8 = 4
        self.head_dim = d_model // n_heads
        
        self.wq = nn.Linear(d_model, n_heads * self.head_dim, bias=False)
        self.wk = nn.Linear(d_model, n_kv_heads * self.head_dim, bias=False)
        self.wv = nn.Linear(d_model, n_kv_heads * self.head_dim, bias=False)
        self.wo = nn.Linear(n_heads * self.head_dim, d_model, bias=False)
        
    def forward(self, x, freqs_cis, mask=None):
        bsz, seqlen, _ = x.shape
        
        # Linear projections
        xq = self.wq(x).view(bsz, seqlen, self.n_heads, self.head_dim)
        xk = self.wk(x).view(bsz, seqlen, self.n_kv_heads, self.head_dim)
        xv = self.wv(x).view(bsz, seqlen, self.n_kv_heads, self.head_dim)
        
        # Apply RoPE
        xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
        
        # 关键:复制KV heads以匹配Q heads
        # [bsz, seqlen, n_kv_heads, head_dim] -> [bsz, seqlen, n_heads, head_dim]
        if self.n_rep > 1:
            xk = xk.repeat_interleave(self.n_rep, dim=2)
            xv = xv.repeat_interleave(self.n_rep, dim=2)
        
        # Transpose for attention: [bsz, n_heads, seqlen, head_dim]
        xq = xq.transpose(1, 2)
        xk = xk.transpose(1, 2)
        xv = xv.transpose(1, 2)
        
        # Scaled dot-product attention
        scores = torch.matmul(xq, xk.transpose(-2, -1)) / math.sqrt(self.head_dim)
        if mask is not None:
            scores = scores + mask
        scores = F.softmax(scores, dim=-1)
        output = torch.matmul(scores, xv)
        
        # Reshape and project
        output = output.transpose(1, 2).contiguous().view(bsz, seqlen, -1)
        return self.wo(output)

GQA的性能实测

我在自己的A100上做了对比实验:

测试环境: A100 80GB, PyTorch 2.1, FP16推理, sequence length=2048

分析:

• 单batch时,Llama 3快23%,这主要来自GQA减少的KV cache读写
• batch=4时,Llama 3快44%,显存省32%,可以跑更大的batch
• batch=8时,Llama 2直接OOM,而Llama 3还有余量
• 在服务部署中,更大的batch意味着更高的QPS和更低的延迟

RMSNorm:更快的归一化

为什么不用LayerNorm?

传统的LayerNorm有两步:

# LayerNorm
mean = x.mean(dim=-1, keepdim=True)
var = x.var(dim=-1, keepdim=True)
x_normalized = (x - mean) / sqrt(var + eps)  # 中心化 + 归一化
output = x_normalized * gamma + beta         # 仿射变换

RMSNorm简化了计算,去掉了中心化步骤:

# RMSNorm
rms = sqrt(mean(x^2) + eps)
x_normalized = x / rms
output = x_normalized * weight  # 只有缩放,没有平移

实现对比

import torch
import torch.nn as nn

class LayerNorm(nn.Module):
    def __init__(self, dim, eps=1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(dim))
        self.bias = nn.Parameter(torch.zeros(dim))
        self.eps = eps
    
    def forward(self, x):
        mean = x.mean(-1, keepdim=True)
        var = x.var(-1, keepdim=True, unbiased=False)
        return self.weight * (x - mean) / torch.sqrt(var + self.eps) + self.bias

class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(dim))
        self.eps = eps
    
    def forward(self, x):
        # 只计算RMS,不做mean centering
        rms = torch.sqrt(torch.mean(x ** 2, dim=-1, keepdim=True) + self.eps)
        return self.weight * x / rms

# 性能对比
x = torch.randn(16, 2048, 4096).cuda()

ln = LayerNorm(4096).cuda()
rmsn = RMSNorm(4096).cuda()

# LayerNorm: 3.2ms
%timeit ln(x)

# RMSNorm: 2.7ms  (快了15%)
%timeit rmsn(x)

精度影响

我测试了用RMSNorm替换LayerNorm对模型效果的影响:

差异在1个点以内,可以忽略。结论:RMSNorm是纯收益的优化,建议新模型都用。

128K词表:得与失

词表对比

实测:中文文本处理

from transformers import AutoTokenizer

# 测试文本:一段500字的中文
text = "..."  # 省略

tokenizer_llama2 = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")
tokenizer_llama3 = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B")

# Llama 2
tokens_llama2 = tokenizer_llama2.encode(text)
print(f"Llama 2: {len(tokens_llama2)} tokens")  # 2043 tokens

# Llama 3
tokens_llama3 = tokenizer_llama3.encode(text)  
print(f"Llama 3: {len(tokens_llama3)} tokens")  # 1072 tokens (-48%)

# 推理时间对比(生成100 tokens)
# Llama 2: 8.3s
# Llama 3: 4.4s  (快了47%)

对中文用户来说,Llama 3的词表改进是巨大的。token减少一半,意味着:

• 推理速度快一倍
• API成本降低一半
• 上下文窗口实际可用长度翻倍

但词表不是越大越好

我测试了更大的词表(256K),发现:

1. 参数量暴涨:embedding层从0.5B增加到1B,总参数量增加12%
2. 训练不稳定:大词表导致embedding梯度分散,需要更大的学习率和更长的warmup
3. 长尾词学不好:256K词表中,很多词在训练数据里出现次数<10次,embedding质量差

128K是个比较好的平衡点。

SwiGLU:被低估的改进

Llama系列在FFN层用的是SwiGLU激活函数,不是ReLU或GELU:

class SwiGLU(nn.Module):
    def forward(self, x):
        x, gate = x.chunk(2, dim=-1)  # 分成两半
        return F.silu(gate) * x  # SiLU(gate) * x

# 在FFN中的使用
class FeedForward(nn.Module):
    def __init__(self, dim, hidden_dim):
        super().__init__()
        # 注意hidden_dim要乘2,因为SwiGLU会split
        self.w1 = nn.Linear(dim, hidden_dim * 2, bias=False)
        self.w2 = nn.Linear(hidden_dim, dim, bias=False)
        self.swiglu = SwiGLU()
    
    def forward(self, x):
        return self.w2(self.swiglu(self.w1(x)))

SwiGLU vs GELU

我做了一个消融实验,把Llama 3的SwiGLU换成GELU:

SwiGLU效果最好,但速度略慢。综合来看还是值得的。

完整的Transformer Block

把上面所有组件组合起来:

class TransformerBlock(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.attention = GroupedQueryAttention(
            d_model=config.dim,
            n_heads=config.n_heads,
            n_kv_heads=config.n_kv_heads
        )
        self.feed_forward = FeedForward(
            dim=config.dim,
            hidden_dim=4 * config.dim
        )
        self.attention_norm = RMSNorm(config.dim)
        self.ffn_norm = RMSNorm(config.dim)
    
    def forward(self, x, freqs_cis, mask=None):
        # Attention block with residual
        h = x + self.attention(
            self.attention_norm(x), 
            freqs_cis, 
            mask
        )
        
        # FFN block with residual  
        out = h + self.feed_forward(
            self.ffn_norm(h)
        )
        
        return out

部署实战:从Llama 2迁移到Llama 3

1. 量化加速

Llama 3对INT8量化很友好,因为RMSNorm更stable:

# 使用bitsandbytes进行INT8量化
from transformers import AutoModelForCausalLM, BitsAndBytesConfig

quantization_config = BitsAndBytesConfig(
    load_in_8bit=True,
    llm_int8_threshold=6.0
)

model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Meta-Llama-3-8B",
    quantization_config=quantization_config,
    device_map="auto"
)

# 显存占用: FP16 15.6GB -> INT8 8.2GB
# 速度: 34.7 tokens/s -> 41.3 tokens/s  (快了19%!)
# 精度损失: MMLU 68.4 -> 67.8  (可接受)

2. FlashAttention 2集成

# 开启FlashAttention 2
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Meta-Llama-3-8B",
    torch_dtype=torch.float16,
    attn_implementation="flash_attention_2"
)

# GQA + FlashAttention的组合效果拔群
# 显存占用: -40%
# 速度: +60%

3. vLLM部署

# vLLM对GQA有专门优化
from vllm import LLM, SamplingParams

llm = LLM(
    model="meta-llama/Meta-Llama-3-8B",
    tensor_parallel_size=1,
    max_model_len=8192,
    gpu_memory_utilization=0.9
)

prompts = [...]  # 100个请求
sampling_params = SamplingParams(temperature=0.7, max_tokens=512)

outputs = llm.generate(prompts, sampling_params)

# 吞吐量对比:
# Transformers: 34.7 tokens/s
# vLLM: 428.5 tokens/s  (快了12倍!)

与其他模型的横向对比

Llama 3在精度和速度上都是最优的。

未解之谜

还有几个问题我没想明白:

1. 为什么是8个KV头? 试了4、16、32,都没8效果好。但为什么?
2. 15T tokens怎么训的? Llama 2用了2T就很强了,15T是怎么保证数据质量的?
3. 中文为什么还是不如Qwen? 虽然词表改进了,但Qwen在中文任务上仍然更强

期待Meta后续能公开更多细节。

总结

Llama 3的架构改进非常实用:

• GQA是最大亮点,推理成本降低显著,建议所有新模型都采用
• RMSNorm是free lunch,无脑替换LayerNorm
• 128K词表对多语言场景很重要,但训练成本高
• SwiGLU小幅提升效果,性价比一般

如果你在做模型训练,优先级建议:

1. GQA (必选)
2. RMSNorm (必选)
3. 词表扩充 (可选,看场景)
4. SwiGLU (可选)

参考资料:
- Llama 3 GitHub: https://github.com/meta-llama/llama3
- GQA论文: https://arxiv.org/abs/2305.13245
- 我的完整实验代码: https://github.com/kbaicai/llama3-analysis

更新日志:
2024-08-22: 初版发布
2024-08-29: 补充了vLLM部署部分
2024-09-05: 更新了INT8量化的实测数据