ViperEkura's Blog
关于llama架构的分析
ViperEkura Lv1
  2026-01-07 18:57:23   2026-01-07 20:18:11  

1. 整体架构概览

Llama 3采用了经典的纯解码器Transformer架构，整体设计围绕自回归语言生成任务进行深度优化。模型核心由词嵌入层、多个堆叠的Transformer块以及输出层构成，每个Transformer块内部包含多头注意力机制和前馈网络两部分。特别值得注意的是其位置编码方案——使用了改进版的旋转位置编码，基础频率参数设定为50万，这比原始RoPE的1万要高出许多，使得模型能够更好地处理长序列文本并具备更强的位置外推能力。

在注意力机制方面，Llama 3引入了分组查询注意力设计，允许键值头的数量少于查询头，通过重复利用键值头来匹配查询头的数量，这种设计在保持模型性能的同时显著降低了内存占用。前馈网络采用了SwiGLU激活函数，其独特的门控机制提供了比传统MLP更强的表达能力。整个模型还采用了RMSNorm进行层归一化，计算效率高于标准的LayerNorm，同时配合模型并行策略和KV缓存技术，在训练和推理过程中都实现了优异的速度表现和内存效率。

1.1 模型初始化

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
class Transformer(nn.Module):
    def __init__(self, params: ModelArgs):
        super().__init__()
        self.params = params
        self.vocab_size = params.vocab_size
        self.n_layers = params.n_layers
        
        # 1. 词嵌入层（并行化）
        self.tok_embeddings = VocabParallelEmbedding(
            params.vocab_size, params.dim, init_method=lambda x: x
        )
        
        # 2. Transformer层堆叠
        self.layers = torch.nn.ModuleList()
        for layer_id in range(params.n_layers):
            self.layers.append(TransformerBlock(layer_id, params))
        
        # 3. 输出层
        self.norm = RMSNorm(params.dim, eps=params.norm_eps)
        self.output = ColumnParallelLinear(
            params.dim, params.vocab_size, bias=False, init_method=lambda x: x
        )
        
        # 4. 预计算旋转位置编码
        self.freqs_cis = precompute_freqs_cis(
            params.dim // params.n_heads,  # 每个头的维度
            params.max_seq_len * 2,        # 两倍序列长度，用于缓存
            params.rope_theta,             # RoPE基础频率
        )

关键点：

  • 使用 VocabParallelEmbedding 将词表分割到多个GPU
  • 每层都是独立的 TransformerBlock
  • 输出前进行 RMSNorm 归一化
  • 预先计算旋转位置编码，减少运行时计算

2. 注意力机制详解

2.1 注意力头配置

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
class Attention(nn.Module):
    def __init__(self, args: ModelArgs):
        super().__init__()
        # GQA（分组查询注意力）配置
        self.n_kv_heads = args.n_heads if args.n_kv_heads is None else args.n_kv_heads
        
        # 模型并行配置
        model_parallel_size = fs_init.get_model_parallel_world_size()
        self.n_local_heads = args.n_heads // model_parallel_size
        self.n_local_kv_heads = self.n_kv_heads // model_parallel_size
        self.n_rep = self.n_local_heads // self.n_local_kv_heads  # 重复因子
        
        # 投影层（模型并行）
        self.wq = ColumnParallelLinear(args.dim, args.n_heads * self.head_dim, bias=False)
        self.wk = ColumnParallelLinear(args.dim, self.n_kv_heads * self.head_dim, bias=False)
        self.wv = ColumnParallelLinear(args.dim, self.n_kv_heads * self.head_dim, bias=False)
        self.wo = RowParallelLinear(args.n_heads * self.head_dim, args.dim, bias=False)
        
        # KV缓存（用于推理加速）
        self.cache_k = torch.zeros((args.max_batch_size, args.max_seq_len, 
                                   self.n_local_kv_heads, self.head_dim)).cuda()
        self.cache_v = torch.zeros((args.max_batch_size, args.max_seq_len, 
                                   self.n_local_kv_heads, self.head_dim)).cuda()

关键点：

  • 支持 GQA：当 n_kv_heads < n_heads 时，KV头会被复用
  • 每个GPU只计算部分注意力头（n_local_heads）
  • KV缓存预分配内存，加速自回归生成

2.2 注意力前向传播

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
def forward(self, x: torch.Tensor, start_pos: int, freqs_cis: torch.Tensor, mask: Optional[torch.Tensor]):
    bsz, seqlen, _ = x.shape
    
    # 1. 线性投影得到 Q, K, V
    xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)
    
    # 2. 重塑为多头格式
    xq = xq.view(bsz, seqlen, self.n_local_heads, self.head_dim)
    xk = xk.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)
    xv = xv.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)
    
    # 3. 应用旋转位置编码（RoPE）
    xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs_cis)
    
    # 4. 更新KV缓存
    self.cache_k[:bsz, start_pos : start_pos + seqlen] = xk
    self.cache_v[:bsz, start_pos : start_pos + seqlen] = xv
    
    # 5. 获取缓存的KV（包含历史所有token）
    keys = self.cache_k[:bsz, : start_pos + seqlen]
    values = self.cache_v[:bsz, : start_pos + seqlen]
    
    # 6. 如果使用GQA，复制KV头以匹配Q头数量
    keys = repeat_kv(keys, self.n_rep)      # (bs, seqlen, n_local_heads, head_dim)
    values = repeat_kv(values, self.n_rep)  # (bs, seqlen, n_local_heads, head_dim)
    
    # 7. 注意力计算
    xq = xq.transpose(1, 2)                 # (bs, n_local_heads, seqlen, head_dim)
    keys = keys.transpose(1, 2)             # (bs, n_local_heads, cache_len, head_dim)
    values = values.transpose(1, 2)
    
    # 8. 缩放点积注意力
    scores = torch.matmul(xq, keys.transpose(2, 3)) / math.sqrt(self.head_dim)
    
    # 9. 应用因果掩码
    if mask is not None:
        scores = scores + mask  # mask是下三角矩阵，包含 -inf
    
    # 10. Softmax和输出
    scores = F.softmax(scores.float(), dim=-1).type_as(xq)
    output = torch.matmul(scores, values)  # (bs, n_local_heads, seqlen, head_dim)
    
    # 11. 合并多头输出
    output = output.transpose(1, 2).contiguous().view(bsz, seqlen, -1)
    return self.wo(output)  # 输出投影

3. RoPE（旋转位置编码）实现

3.1 预计算旋转矩阵

1
2
3
4
5
6
7
8
9
10
11
12
13
14
def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0):
    # 1. 计算基础频率（公式：theta^{-2i/dim}）
    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
    
    # 2. 创建位置序列
    t = torch.arange(end, device=freqs.device, dtype=torch.float32)
    
    # 3. 计算外积：每个位置 * 每个频率
    freqs = torch.outer(t, freqs)  # (end, dim//2)
    
    # 4. 转换为复数形式：e^{i*θ} = cosθ + i*sinθ
    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)  # complex64
    
    return freqs_cis

数学原理：

  • 对于位置 m 和维度 i，旋转角度为：m * theta^{-2i/d}
  • 复数形式：(cos(mθ), sin(mθ)) 应用于 Q 和 K 的相邻维度对

3.2 应用旋转位置编码

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
def apply_rotary_emb(xq: torch.Tensor, xk: torch.Tensor, freqs_cis: torch.Tensor):
    # 1. 将最后两个维度重塑为复数形式
    # 输入形状: (batch, seqlen, n_heads, head_dim)
    # 重塑为: (batch, seqlen, n_heads, head_dim//2, 2)
    # 然后转换为复数: (batch, seqlen, n_heads, head_dim//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))
    
    # 2. 广播频率到合适的形状
    freqs_cis = reshape_for_broadcast(freqs_cis, xq_)
    
    # 3. 复数乘法实现旋转
    # 对于复数 a+bi 和 cosθ+isinθ，乘法实现旋转
    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)

4. 前馈网络（SwiGLU）

4.1 FeedForward 实现

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
class FeedForward(nn.Module):
    def __init__(self, dim: int, hidden_dim: int, multiple_of: int, ffn_dim_multiplier: Optional[float]):
        super().__init__()
        # 1. 计算隐藏层维度（默认是 4/3 * dim）
        hidden_dim = int(2 * hidden_dim / 3)
        
        # 2. 可选的隐藏层缩放因子
        if ffn_dim_multiplier is not None:
            hidden_dim = int(ffn_dim_multiplier * hidden_dim)
        
        # 3. 对齐到 multiple_of 的倍数（内存优化）
        hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)
        
        # 4. SwiGLU 的三个投影层
        self.w1 = ColumnParallelLinear(dim, hidden_dim, bias=False, gather_output=False)
        self.w2 = RowParallelLinear(hidden_dim, dim, bias=False, input_is_parallel=True)
        self.w3 = ColumnParallelLinear(dim, hidden_dim, bias=False, gather_output=False)

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

SwiGLU 公式：

1
FFN(x) = W₂ · (silu(W₁·x) ⊙ W₃·x)

其中 silu 是 Swish 激活函数：x * sigmoid(x)

5. Transformer Block 结构

5.1 单层 Transformer

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
class TransformerBlock(nn.Module):
    def __init__(self, layer_id: int, args: ModelArgs):
        super().__init__()
        self.n_heads = args.n_heads
        self.dim = args.dim
        self.head_dim = args.dim // args.n_heads
        self.attention = Attention(args)
        self.feed_forward = FeedForward(
            dim=args.dim,
            hidden_dim=4 * args.dim,  # 默认 4倍扩展
            multiple_of=args.multiple_of,
            ffn_dim_multiplier=args.ffn_dim_multiplier,
        )
        self.layer_id = layer_id
        self.attention_norm = RMSNorm(args.dim, eps=args.norm_eps)
        self.ffn_norm = RMSNorm(args.dim, eps=args.norm_eps)

    def forward(self, x: torch.Tensor, start_pos: int, freqs_cis: torch.Tensor, mask: Optional[torch.Tensor]):
        # 1. 残差连接 + 注意力（Pre-Norm）
        h = x + self.attention(self.attention_norm(x), start_pos, freqs_cis, mask)
        
        # 2. 残差连接 + FFN（Pre-Norm）
        out = h + self.feed_forward(self.ffn_norm(h))
        return out

架构特点：

  • Pre-LayerNorm：归一化在子层之前
  • 残差连接：每个子层都有残差连接
  • 并行计算：注意力头和FFN都支持模型并行
 REWARD AUTHOR
  1. 1. 整体架构概览
    1. 1.1 模型初始化
  2. 2. 注意力机制详解
    1. 2.1 注意力头配置
    2. 2.2 注意力前向传播
  3. 3. RoPE（旋转位置编码）实现
    1. 3.1 预计算旋转矩阵
    2. 3.2 应用旋转位置编码
  4. 4. 前馈网络（SwiGLU）
    1. 4.1 FeedForward 实现
  5. 5. Transformer Block 结构
    1. 5.1 单层 Transformer
  1. 1. 整体架构概览
    1. 1.1 模型初始化
  2. 2. 注意力机制详解
    1. 2.1 注意力头配置
    2. 2.2 注意力前向传播
  3. 3. RoPE（旋转位置编码）实现
    1. 3.1 预计算旋转矩阵
    2. 3.2 应用旋转位置编码
  4. 4. 前馈网络（SwiGLU）
    1. 4.1 FeedForward 实现
  5. 5. Transformer Block 结构
    1. 5.1 单层 Transformer