LLaMA架构深度解析：从开源大模型设计到高效推理的工程实现全流程

一、引言：LLaMA如何改变开源大模型格局

2023年Meta发布LLaMA系列模型，标志着开源大模型时代的到来。LLaMA-1证明了使用公开数据训练可以接近GPT-3级别的性能；LLaMA-2引入了GQA和更长的上下文；LLaMA-3在15T token上训练，性能全面超越GPT-3.5。LLaMA的架构设计影响了几乎所有后续开源模型——Alpaca、Vicuna、CodeLlama等都是基于LLaMA微调的。LLaMA的工程价值在于其简洁高效的架构：Decoder-Only、RoPE、RMSNorm、SwiGLU、GQA的组合成为现代大模型的事实标准。本文将深入解析LLaMA的架构设计并提供完整实现。

二、LLaMA架构组件

import torch
import torch.nn as nn
import torch.nn.functional as F
import math
from typing import List, Dict, Tuple, Optional, Any
from dataclasses import dataclass

@dataclass
class LLaMAConfig:
    """LLaMA配置"""
    vocab_size: int = 32000
    hidden_size: int = 4096
    intermediate_size: int = 11008
    num_hidden_layers: int = 32
    num_attention_heads: int = 32
    num_kv_heads: int = 32  # LLaMA-1: MHA, LLaMA-2 70B: GQA(8)
    rms_norm_eps: float = 1e-6
    rope_theta: float = 10000.0
    max_position_embeddings: int = 4096
    tie_word_embeddings: bool = False

class LLaMARMSNorm(nn.Module):
    """RMSNorm"""
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))
    
    def forward(self, x):
        variance = x.pow(2).mean(-1, keepdim=True)
        x = x * torch.rsqrt(variance + self.eps)
        return self.weight * x

class LLaMARoPE(nn.Module):
    """旋转位置编码"""
    def __init__(self, dim: int, max_len: int = 4096, theta: float = 10000.0):
        super().__init__()
        inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer('inv_freq', inv_freq)
        self._build_cache(max_len)
    
    def _build_cache(self, max_len):
        positions = torch.arange(max_len).float()
        freqs = torch.outer(positions, self.inv_freq)
        emb = torch.cat([freqs, freqs], dim=-1)
        self.register_buffer('cos_cached', emb.cos())
        self.register_buffer('sin_cached', emb.sin())
    
    def forward(self, q, k, positions=None):
        seq_len = q.shape[2]
        if positions is not None:
            cos = self.cos_cached[positions].unsqueeze(1)
            sin = self.sin_cached[positions].unsqueeze(1)
        else:
            cos = self.cos_cached[:seq_len].unsqueeze(0).unsqueeze(0)
            sin = self.sin_cached[:seq_len].unsqueeze(0).unsqueeze(0)
        
        def rotate(x):
            x1 = x[..., :x.shape[-1] // 2]
            x2 = x[..., x.shape[-1] // 2:]
            return torch.cat([-x2, x1], dim=-1)
        
        return q * cos + rotate(q) * sin, k * cos + rotate(k) * sin

class LLaMAAttention(nn.Module):
    """LLaMA注意力层"""
    def __init__(self, config: LLaMAConfig):
        super().__init__()
        self.n_heads = config.num_attention_heads
        self.n_kv_heads = config.num_kv_heads
        self.n_rep = self.n_heads // self.n_kv_heads
        self.head_dim = config.hidden_size // self.n_heads
        
        self.q_proj = nn.Linear(config.hidden_size, self.n_heads * self.head_dim, bias=False)
        self.k_proj = nn.Linear(config.hidden_size, self.n_kv_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(config.hidden_size, self.n_kv_heads * self.head_dim, bias=False)
        self.o_proj = nn.Linear(self.n_heads * self.head_dim, config.hidden_size, bias=False)
        
        self.rope = LLaMARoPE(self.head_dim, config.max_position_embeddings, config.rope_theta)
    
    def forward(self, x, mask=None):
        B, L, _ = x.shape
        q = self.q_proj(x).view(B, L, self.n_heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(B, L, self.n_kv_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(B, L, self.n_kv_heads, self.head_dim).transpose(1, 2)
        
        q, k = self.rope(q, k)
        
        if self.n_rep > 1:
            k = k.unsqueeze(2).expand(-1, -1, self.n_rep, -1, -1).reshape(B, self.n_heads, L, self.head_dim)
            v = v.unsqueeze(2).expand(-1, -1, self.n_rep, -1, -1).reshape(B, self.n_heads, L, self.head_dim)
        
        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))
        attn = F.softmax(scores.float(), dim=-1).to(q.dtype)
        out = torch.matmul(attn, v).transpose(1, 2).contiguous().view(B, L, -1)
        return self.o_proj(out)

class LLaMAMLP(nn.Module):
    """SwiGLU前馈网络"""
    def __init__(self, config: LLaMAConfig):
        super().__init__()
        self.gate = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.up = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.down = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
    
    def forward(self, x):
        return self.down(F.silu(self.gate(x)) * self.up(x))

class LLaMABlock(nn.Module):
    """LLaMA Transformer层"""
    def __init__(self, config: LLaMAConfig):
        super().__init__()
        self.norm1 = LLaMARMSNorm(config.hidden_size, config.rms_norm_eps)
        self.attn = LLaMAAttention(config)
        self.norm2 = LLaMARMSNorm(config.hidden_size, config.rms_norm_eps)
        self.mlp = LLaMAMLP(config)
    
    def forward(self, x, mask=None):
        x = x + self.attn(self.norm1(x), mask)
        x = x + self.mlp(self.norm2(x))
        return x

class LLaMA(nn.Module):
    """完整LLaMA模型"""
    def __init__(self, config: LLaMAConfig):
        super().__init__()
        self.config = config
        self.embed = nn.Embedding(config.vocab_size, config.hidden_size)
        self.layers = nn.ModuleList([LLaMABlock(config) for _ in range(config.num_hidden_layers)])
        self.norm = LLaMARMSNorm(config.hidden_size, config.rms_norm_eps)
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        if config.tie_word_embeddings:
            self.lm_head.weight = self.embed.weight
        self._init_weights()
    
    def _init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, mean=0, std=0.02)
            elif isinstance(m, nn.Embedding):
                nn.init.normal_(m.weight, mean=0, std=0.02)
        for name, p in self.named_parameters():
            if 'down' in name or 'o_proj' in name:
                p.data.normal_(mean=0, std=0.02 / math.sqrt(2 * self.config.num_hidden_layers))
    
    def forward(self, input_ids, labels=None):
        x = self.embed(input_ids)
        L = input_ids.shape[1]
        mask = torch.ones(L, L, device=x.device).tril().unsqueeze(0).unsqueeze(0)
        for layer in self.layers:
            x = layer(x, mask)
        x = self.norm(x)
        logits = self.lm_head(x)
        loss = None
        if labels is not None:
            loss = F.cross_entropy(logits[..., :-1, :].contiguous().view(-1, logits.size(-1)),
                                   labels[..., 1:].contiguous().view(-1), ignore_index=-100)
        return {'logits': logits, 'loss': loss}
    
    @torch.no_grad()
    def generate(self, input_ids, max_new=100, temperature=0.7, top_p=0.9):
        for _ in range(max_new):
            logits = self.forward(input_ids[:, -self.config.max_position_embeddings:])['logits'][:, -1, :] / temperature
            sorted_logits, sorted_idx = torch.sort(logits, descending=True)
            cum_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
            remove = cum_probs > top_p
            remove[..., 1:] = remove[..., :-1].clone()
            remove[..., 0] = False
            sorted_logits[remove] = float('-inf')
            probs = F.softmax(sorted_logits, dim=-1)
            next_token = sorted_idx.gather(-1, torch.multinomial(probs, 1))
            input_ids = torch.cat([input_ids, next_token], dim=1)
        return input_ids

def test_llama():
    config = LLaMAConfig(vocab_size=1000, hidden_size=128, intermediate_size=256,
                         num_hidden_layers=2, num_attention_heads=4, num_kv_heads=2,
                         max_position_embeddings=128)
    model = LLaMA(config)
    ids = torch.randint(0, 1000, (2, 32))
    out = model(ids, labels=ids)
    print(f"Logits: {out['logits'].shape}, Loss: {out['loss'].item():.4f}")
    print(f"Params: {sum(p.numel() for p in model.parameters()):,}")

if __name__ == "__main__":
    test_llama()

三、LLaMA系列对比

def llama_comparison():
    models = [
        ("LLaMA-1 7B", "1T", 2048, "MHA", "公开数据，接近GPT-3"),
        ("LLaMA-1 13B", "1T", 2048, "MHA", "性能更优"),
        ("LLaMA-1 65B", "1.4T", 2048, "MHA", "旗舰，Chinchilla最优"),
        ("LLaMA-2 7B", "2T", 4096, "MHA", "商用许可"),
        ("LLaMA-2 70B", "2T", 4096, "GQA(8)", "GQA降低推理开销"),
        ("LLaMA-3 8B", "15T", 8192, "GQA(8)", "大幅扩展数据"),
        ("LLaMA-3 70B", "15T", 8192, "GQA(8)", "全面超越GPT-3.5"),
        ("LLaMA-3.1 405B", "15T", 131072, "GQA(8)", "最大开源模型"),
    ]
    print(f"{'模型':<20} {'训练token':<10} {'上下文':<10} {'注意力':<10} {'特点'}")
    for name, tokens, ctx, attn, desc in models:
        print(f"{name:<20} {tokens:<10} {ctx:<10} {attn:<10} {desc}")
    
    print("\nLLaMA核心设计选择:")
    choices = [
        ("RMSNorm替代LayerNorm", "更高效，去均值计算"),
        ("SwiGLU替代ReLU", "门控机制提升表达能力"),
        ("RoPE位置编码", "相对位置，支持长度外推"),
        ("Pre-Norm架构", "训练更稳定"),
        ("GQA(LLaMA-2起)", "减少KV Cache 75%"),
        ("无bias", "减少参数，更简洁"),
        ("大词表(128K LLaMA-3)", "多语言支持"),
    ]
    for name, desc in choices:
        print(f"  - {name}: {desc}")

if __name__ == "__main__":
    llama_comparison()

四、总结

LLaMA系列通过RMSNorm、SwiGLU、RoPE、GQA等简洁高效的架构选择，定义了现代大模型的设计范式。其开源策略催生了整个开源大模型生态，Alpaca、Vicuna、CodeLlama等衍生模型证明了LLaMA作为基础模型的泛化能力。LLaMA-3在15T token上的训练展示了数据规模对模型性能的关键作用。理解LLaMA的架构设计，是理解现代大模型工程实践的基础。