Skip to main content

自回归模型的两个阶段

  • Prefill 阶段：并行处理输入 prompt 的所有 token，并输出第一个 token；同时生成用于未来解码的 KV Cache。因为输入序列长度很长，所以计算开销大，很小的 batch size （批量处理的序列数）就会让 GPU 打满。增大 batch size，prefill 阶段单个 token 开销几乎是不变的。
  • Decode 阶段：（开启 KV Cache）使用先前的 KV Cache 以自回归方式逐步生成新的 token。Decode 阶段在每个自回归步仅仅会生成一个 token。虽然decode 阶段输入的序列长度一直是 1，但是需要反复读取 KV Cache，故而 IO 开销很大。Decode 时，单个 token 的开销显著大于 prefill 阶段。需要把 decode 阶段的 batch size 配置得非常大才有可能占满 GPU ，但会因为 KV Cache 读写开销太大而变得不现实，所以 Decode 阶段的 GPU 利用率很低。

TTFT（Time To First Token）指的是用户发出请求后，模型输出第一个 token 所需的时间，它主要受预填充（prefill）阶段计算和缓存加载影响，决定了响应的“开口速度”； TPOT（Time Per Output Token）是模型在开始生成后，平均每个 token 的输出耗时，它反映了增量解码阶段（decode）的效率，决定了生成过程的“语速流畅度”。 换句话说，TTFT解决“多久开始说话”，TPOT衡量“说话有多快”。

多头注意力机制

在语言模型（LLM）中，“head” 通常指的是自注意力机制（self-attention mechanism）中的一个注意力头（attention head）。在多头注意力机制（multi-head attention）中，多个注意力头被用来从不同的表示子空间中提取信息，每个注意力头可以被看作是一个独立的注意力机制。通过同时使用多个注意力头，模型能够捕获输入序列中更丰富和多样的特征。

注意力头（Attention Head）是自注意力机制中的一个关键组件，它帮助模型在处理输入序列时识别和关注重要部分。通常，自注意力机制中的注意力头包含以下几个步骤：

  1. 线性变换：

    • 输入序列中的每个元素通常被表示为一个向量。对于每个注意力头，这些输入向量首先通过三个不同的线性变换层，分别生成查询（Query）、键（Key）和值（Value）向量。这些变换通常用参数矩阵 WQW^Q, WKW^K, 和 WVW^V 表示： Q=XWQ,K=XWK,V=XWVQ = XW^Q, \quad K = XW^K, \quad V = XW^V

    其中 XX 是输入序列的表示。

  2. 计算注意力分数：

    • 对于每个查询向量 qiq_i，通过与所有键向量 kjk_j 计算点积以获得注意力分数。这些分数度量了输入序列中各个位置之间的相关性： score(qi,kj)=qi⋅kj\text{score}(q_i, k_j) = q_i \cdot k_j
    • 为了增强数值稳定性，这些分数通常会被除以 dk\sqrt{d_k}，其中 dkd_k 是键向量的维度。

  3. 计算注意力权重：

    • 使用Softmax函数将注意力分数转换为概率分布，即注意力权重。这些权重表示每个位置的重要性： αij=softmax(score(qi,k))j=exp⁡(score(qi,kj))∑lexp⁡(score(qi,kl))\alpha_{ij} = \text{softmax}(\text{score}(q_i, k))_j = \frac{\exp(\text{score}(q_i, k_j))}{\sum_{l} \exp(\text{score}(q_i, k_l))}

  4. 加权求和：

    • 使用计算得到的注意力权重对值向量进行加权求和，得到该注意力头的输出： Attention(Q,K,V)i=∑jαijvj\text{Attention}(Q, K, V)_i = \sum_{j} \alpha_{ij} v_j
    • 每个查询向量 qiq_i 产生一个输出向量，这些输出向量组成了注意力头的输出。

  5. 多头注意力：

    • 在多头注意力机制中，多个注意力头的输出会被连接（concatenate）在一起，然后通过一个线性层进行变换，以整合来自不同头的信息。

通过这些步骤，注意力头能够有效地捕捉输入序列中各位置之间的关系和依赖信息。这种机制使得模型能够灵活地关注序列的不同部分，从而提高对复杂语言模式的理解和生成能力。

  1. head_num (32):

    • 这是模型中多头注意力机制的总注意力头数。多头注意力（Multi-head Attention）是Transformer架构中的一个关键组件，它允许模型在不同的子空间中并行学习和提取特征。32个头意味着该模型可以在32个不同的特征子空间中同时进行注意力计算，增强模型的表现力和对复杂模式的捕捉能力。

  2. kv_head_num (8):

    • 这是与键和值（Key-Value）相关的注意力头数。通常情况下，“kv_head_num” 指的是模型中专门用于处理键和值的注意力头的数量。这可能意味着在某些变种或优化版本中，模型对查询（Query）和键值对（Key-Value）采用不同的头数，以实现特定的性能提升或计算优化。

  3. size_per_head (128):

    • 这是每个注意力头的维度大小。具体来说，“size_per_head” 指的是每个注意力头在计算中生成的向量的维度。128的维度表示每个注意力头输出的向量具有128个特征，这决定了每个头可以捕获和表示的信息量。这个参数与注意力头的数量一起影响模型的总参数量及其计算复杂度。