模型架构设计

本文档涵盖 LLM 模型架构的核心设计原则、主流架构变体及其工程实现细节。

设计原则

• 可扩展性：支持从 7B 到 700B+ 不同规模的模型训练与推理
• 效率：平衡计算和内存需求，优化吞吐量和延迟
• 灵活性：支持多种任务类型（生成、分类、对话、代码等）
• 稳定性：训练稳定，收敛可靠，不易出现梯度爆炸

典型架构演进

1. Encoder-Decoder（Seq2Seq）

最早的大规模预训练架构，Encoder 和 Decoder 对称堆叠。典型代表：T5、BART。

class Seq2SeqModel(nn.Module):
    def __init__(self, encoder, decoder, d_model=768):
        super().__init__()
        self.encoder = encoder
        self.decoder = decoder
        self.d_model = d_model
    
    def forward(self, input_ids, attention_mask, labels=None):
        encoder_outputs = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
        decoder_outputs = self.decoder(
            input_ids=labels,
            encoder_hidden_states=encoder_outputs.last_hidden_state,
            encoder_attention_mask=attention_mask
        )
        return decoder_outputs

适用场景：翻译、摘要等需要输入输出长度差异较大的任务。

2. Decoder-only（主流）

目前绝大多数 LLM 采用的架构，包括 GPT 系列、LLaMA、PaLM 等。

核心组件：

1. Token Embedding：将离散 token 映射为连续向量空间
2. 位置编码（RoPE 或 ALiBi）：提供序列顺序信息
3. 多层 Transformer Block：自注意力 + FFN 的交替堆叠
4. RMSNorm 归一化：替代 LayerNorm，计算更高效
5. SwiGLU 激活的 FFN：比传统 GELU 提升约 5-10% 效果
6. 输出投影：将隐藏状态映射到 vocab 维度

# 简化的 Transformer Block 实现
class TransformerBlock(nn.Module):
    def __init__(self, dim, n_heads, n_ff, dropout=0.1):
        super().__init__()
        self.norm1 = RMSNorm(dim)
        self.attn = MultiHeadAttention(dim, n_heads)
        self.norm2 = RMSNorm(dim)
        self.ffn = SwiGLUFeedForward(dim, n_ff)
        self.drop = nn.Dropout(dropout)
    
    def forward(self, x, mask=None):
        # Pre-Norm: 先归一化再计算
        x = x + self.drop(self.attn(self.norm1(x), mask))
        x = x + self.drop(self.ffn(self.norm2(x)))
        return x

3. Encoder-only

主要用于理解任务（分类、NER、情感分析）。典型代表：BERT、RoBERTa。

class BertClassifier(nn.Module):
    def __init__(self, bert, n_classes=2):
        super().__init__()
        self.bert = BertModel.from_pretrained('bert-base-chinese')
        self.classifier = nn.Linear(768, n_classes)
    
    def forward(self, input_ids, attention_mask):
        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
        cls_token = outputs.last_hidden_state[:, 0, :]  # [CLS] token
        logits = self.classifier(cls_token)
        return logits

关键配置参数

config = {
    'd_model': 4096,           # 隐藏层维度
    'n_heads': 32,             # 注意力头数
    'n_kv_heads': 8,           # GQA 的 KV 头数
    'n_layers': 32,            # Transformer 层数
    'd_ff': 11008,             # FFN 中间层维度
    'vocab_size': 32000,       # 词表大小
    'max_seq_len': 4096,       # 最大序列长度
    'hidden_dim': 14336,       # 实际 FFN 维度 (d_ff * 8/3)
    'rms_norm_eps': 1e-5,      # RMSNorm epsilon
}

参数规模估算

模型	参数量	计算量	显存需求
LLaMA-7B	7B	14 GFLOPs/token	~14 GB
LLaMA-13B	13B	26 GFLOPs/token	~26 GB
LLaMA-70B	70B	140 GFLOPs/token	~140 GB
LLaMA-405B	405B	810 GFLOPs/token	~810 GB

高效组件详解

1. GQA（Grouped Query Attention）

问题：传统 MHA 中每个查询头都需要自己的 KV 对，KV Cache 随 n_heads 线性增长。

解决方案：将查询头分组，每组共享一个 KV 头。

class GroupedQueryAttention(nn.Module):
    def __init__(self, d_model, n_heads, n_kv_heads):
        super().__init__()
        self.d_model = d_model
        self.n_heads = n_heads
        self.n_kv_heads = n_kv_heads
        self.head_dim = d_model // n_heads
        
        # Q 投影
        self.q_proj = nn.Linear(d_model, n_heads * self.head_dim)
        # KV 投影（共享）
        self.k_proj = nn.Linear(d_model, n_kv_heads * self.head_dim)
        self.v_proj = nn.Linear(d_model, n_kv_heads * self.head_dim)
        self.o_proj = nn.Linear(n_heads * self.head_dim, d_model)
    
    def forward(self, hidden_states, kv_cache=None):
        bsz, seq_len, _ = hidden_states.shape
        
        # 投影
        q = self.q_proj(hidden_states).view(bsz, seq_len, self.n_heads, self.head_dim)
        k = self.k_proj(hidden_states).view(bsz, seq_len, self.n_kv_heads, self.head_dim)
        v = self.v_proj(hidden_states).view(bsz, seq_len, self.n_kv_heads, self.head_dim)
        
        # 扩展 KV 以匹配 Q 的头数
        k = k.repeat_interleave(self.n_heads // self.n_kv_heads, dim=2)
        v = v.repeat_interleave(self.n_heads // self.n_kv_heads, dim=2)
        
        # 注意力计算
        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
        attn_weights = F.softmax(scores, dim=-1)
        output = torch.matmul(attn_weights, v)
        
        return self.o_proj(output.view(bsz, seq_len, -1))

优势：

• KV Cache 减少为原来的 n_kv_heads / n_heads
• 推理速度提升 2-3 倍（LLaMA-70B 使用 8 组 KV）

2. SwiGLU 激活

传统 FFN：FFN(x) = W2 * GELU(W1 * x + b1) + b2

SwiGLU：FFN(x) = W3 * SiLU(W1 * x + b1) ⊗ (W2 * x + b2)

class SwiGLUFeedForward(nn.Module):
    def __init__(self, dim, hidden_dim, activation=F.silu):
        super().__init__()
        self.w1 = nn.Linear(dim, hidden_dim, bias=False)
        self.w2 = nn.Linear(dim, hidden_dim, bias=False)
        self.w3 = nn.Linear(hidden_dim, dim, bias=False)
        self.activation = activation
    
    def forward(self, x):
        # SwiGLU: gate * (x projected)
        gate = self.w1(x)
        proj = self.w2(x)
        return self.w3(self.activation(gate) * proj)

优势：

• 相比 GELU 提升约 5-10% 效果
• 参数量相同（不需要额外投影矩阵）

3. RMSNorm

LayerNorm：(x - mean) / sqrt(var + eps) * gamma + beta

RMSNorm：x / sqrt(mean(x^2) + eps) * gamma

class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-5):
        super().__init__()
        self.eps = eps
        self.gamma = nn.Parameter(torch.ones(dim))
    
    def forward(self, x):
        # 计算均方根
        norm = x.norm(2, dim=-1, keepdim=True)
        x_normed = x / (norm + self.eps)
        return x_normed * self.gamma

优势：

• 去除均值计算，加速 10-20%
• 效果与 LayerNorm 相当

4. RoPE（Rotary Position Embedding）

传统位置编码：直接加到 embedding 上，不支持外推。

RoPE：通过旋转矩阵注入位置信息，天然支持外推。

def apply_rope(x, cos, sin):
    """应用 RoPE 位置编码"""
    bsz, seq_len, n_heads, head_dim = x.shape
    x = x.view(bsz, seq_len, n_heads, head_dim // 2, 2)
    
    # 旋转操作
    x1 = x[..., 0]
    x2 = x[..., 1]
    o1 = cos * x1 - sin * x2
    o2 = cos * x2 + sin * x1
    
    return o1.view(bsz, seq_len, n_heads, head_dim // 2).flatten(-2)

架构变体对比

特性	LLaMA	Mistral	Mixtral	Falcon
注意力	MHA	GQA	GQA	MQA
FFN	SwiGLU	SwiGLU	SwiGLU	GELU
Norm	RMSNorm	RMSNorm	RMSNorm	RMSNorm
位置编码	RoPE	RoPE	RoPE	ALiBi
MoE	否	否	是	否

总结

模型架构设计需要在性能、效率和灵活性之间找到平衡。当前的主流趋势是：

• Decoder-only 架构成为 LLM 的标准选择
• GQA 和 RMSNorm 成为高效推理的标配
• SwiGLU 替代 GELU 获得更好的效果
• RoPE 提供灵活的位置编码方案
• MoE 架构进一步扩展模型规模