LLM智能应用开发

LLM结构的学习路径

• LLM结构解析(开源LlaMA)
• 自定义数据集构造
• 自定义损失函数和模型训练/微调

Transformer经典结构

• Encoder-decoder结构
• 输入部分
  • Input embedding
  • Positional embedding
• Transformer部分
  • Feed forward network
  • Attention module

HF LlaMA模型结构

LlamaForCausalLM(
  (model): LlamaModel(
    ...
    (layers): ModuleList(
      (0-15): 16 x LlamaDecoderLayer(
        (self_attn): LlamaAttention
        (mlp): LlamaMLP
        (input_layernorm): LlamaRMSNorm
        (post_attention_layernorm): LlamaRMSNorm
    )
    ...
  )
  (lm_head): Linear(in_features=2048, out_features=128256, bias=False)
)

LlamaDecoderLayer内部结构

Transformer架构的核心: attention(注意力机制)

(self_attn): LlamaAttention(
  (q_proj): Linear(in_features=2048, out_features=2048, bias=False)
  (k_proj): Linear(in_features=2048, out_features=512, bias=False)
  (v_proj): Linear(in_features=2048, out_features=512, bias=False)
  (o_proj): Linear(in_features=2048, out_features=2048, bias=False)
  (rotary_emb): LlamaRotaryEmbedding()
)

Attention内部结构

• 静: 结构视角(init function...)
  • 4个Linear层
    • q_proj、k_proj、v_proj、o_proj
• 动: 推理视角(Forward，bp靠Autograd自动求导)

Attention模块的输入

问题：QKV是输入吗？

• 非也，输入是上一层的hidden states

class LlamaAttention(nn.Module):
...
def forward(hidden_states)
  ...
  query_states = self.q_proj(hidden_states)
  key_states = self.k_proj(hidden_states)
  value_states = self.v_proj(hidden_states)

• 思考：hidden states的shape是怎样的？

标准Attention的第一步: 获得

• 给定hidden states(后续简写为)，通过前向传播(执行forward)得到
  • 的shape: [batch_size, seq_len, hidden_size]
  • : ,
    • 的shape: [hidden_size, hidden_size]
  • :
    • 的shape: [hidden_size, hidden_size]
  • :
    • 的shape: [hidden_size, hidden_size]

标准Attention的第一步: 获得

• 为方便理解方法，脑补通过tensor.view改变shape
  • [batch_size, seq_len, hidden_size] -> [N, d]
    • N = batch_size * seq_len, d = hidden_size

标准Attention的第二步: 计算

• 给定，计算，此处考虑mask

标准Attention的第三步: 计算Attention

• 给定，计算
  • row-wise softmax:
  • , ,

标准Attention的第四步: 计算输出

• 给定和，计算

标准Attention回顾

• 给定 (来自)

Attention中mask的作用

• 回顾
• <PAD>: 一种表示“padding”的特殊token，用来避免对句子中的某些token的影响
• 为了避免padding对attention的影响，在计算时，我们可以将padding的部分设置为一个很大的数，如

Attention中mask的作用

对应的实现

移步notebook

MuliHeadAttention

• 标准Attention只生成一个输出A，考虑多种角度，期望得到不同的A
  • 靠多个头实现，什么头？？
  • 进行拆分，拆成多个头
  • 拆分为多头：[batch_size, seq_len, num_heads, head_dim]
  • 些许改造Attention计算过程

MultiHeadAttention

• 给定 (shape [bs, seq, hs]),shape简化为
• 多个heads
• shape的变换(tensor.view实现): [N, d] -> [N, num_heads, head_dim]
  • 其中, d = hidden_size = num_heads * head_dim
  • 实现中，[bs, seq, hs] -> [bs, seq, nh, hd]
    • 再transpose为[bs, nh, seq, hd]

Attention variations

• 有了多头，多Query，多Key-value，多Query和Key-value的组合
• 考虑到多Query和Key-value的组合

Grouped Query Attention (GQA)

• 介于多头注意力与多Query方案之间：个Query组共享同一组
• 设注意力头数为，分组数为，则每组包含个Query头
• 仍保持独立的q_proj，投影矩阵缩减为，KV缓存体积缩小倍
• Hugging Face配置：num_heads = h，num_key_value_heads = g

GQA带来的改进与取舍

• 推理显存：每个token的KV缓存从降为，长上下文推理显存压力显著降低
• 吞吐表现：减少k_proj/v_proj的矩阵乘与显存访问，提升批量推理吞吐
• 模型实践：Llama-2/3、Gemma、Mistral等开源模型默认启用GQA (num_key_value_heads=g)
• 表达能力：合理选择（常见或）可兼顾速度与精度，分组过大可能削弱头部多样性

对应的实现

移步notebook

Attention计算开销

• 的计算过程是的复杂度，那么多头的情况下，的计算复杂度是
• 实际上，可依赖GPU并行执行提升速度
  • 分块的并行计算(sm计算单元)
• 如何加速Attention计算？
  • BlockedAttention
  • FlashAttention

BlockedAttention第一步: 获得

• 给定 (shape [batch_size, seq_len, hidden_size]),shape简化为

BlockedAttention第二步: 计算

• 给定,
  • 的shape: , 的shape:

BlockedAttention第二步: 计算

• 给定,
  • 的shape: , 的shape:

BlockedAttention第三步: 计算Attention

给定，计算

BlockedAttention第三步: 计算Attention

给定，计算

BlockedAttention第三步: 计算Attention

给定，计算

BlockedAttention的数值稳定性

• 实现时不能对每个块单独归一化后直接拼接，否则等效于多个softmax
• 需要维护全局的行最大值与行和，新块的贡献使用缩放后累加
• 这是FlashAttention等实现中的通用做法，可避免溢出并保证与标准Attention一致

BlockedAttention第四步: 计算

给定，计算

BlockedAttention回顾