项目技术路线指导2-week1

项目技术路线指导-2

Flash Attention和ROPE

Transformer 的计算复杂度随着输入序列长度的增加而呈平方增长，这使得其在处理长序列时面临巨大的挑战。为了解决这个问题，我们引入了 FlashAttention，一种新型的优化技术，通过优化存储器使用，提高了 Transformer 的执行效率。

FlashAttention是一种新型的注意力算法，旨在减少对高带宽内存（High Bandwidth Memory，HBM）的读写操作，从而提高Transformer模型的效率。该算法的主要思想是使注意力机制具备IO意识，考虑到GPU内存级别之间的读写操作。通过减少对HBM的读写操作，FlashAttention可以有效降低模型的内存访问瓶颈，提升训练速度和推理效率。另一项关键技术

另一方面，如何更好地表示位置信息，是 NLP 中的一个重要问题。为了解决这个问题，我们引入了 Rotary Position Embedding (RoPe)，一种新型的位置编码方案。RoPe 的主要思想是：用绝对位置编码来表征相对位置编码,主要就是对attention中的q, k向量注入了绝对位置信息，然后用更新的q,k向量做attention中的内积就会引入相对位置信息了。RoPE通过将位置向量进行旋转操作，将位置信息与输入元素的表示进行融合，使模型能够更好地捕捉序列中的位置关系和上下文信息。

FlashAttention

FlashAttention 的主要思路是尽量使用高速的静态随机存取存储器(SRAM)，减少对存取速度慢的 HBM（High Band Memory）的使用。具体来说，FlashAttention 的计算过程如下：

计算 S = QK^T ∈ R^(N×N)

计算内积 P = softmax(S) ∈ R^(N×N)

计算归一化的注意力 O = PV ∈ R^(N×d)

得到输出 Output

其中，Q, K, V ∈ R^(N×d)，N 是序列长度，d 是 head 的维度。FlashAttention 的关键在于尽量对矩阵乘法进行分块，每个小块都在 on-chip 的（192K）SRAM 进行。这主要通过两个相关的技术实现：Tiling 和 Recomputation。

Tiling技术

Tiling技术，也称为块划分或数据分块，是一种内存优化技术，旨在减少缓存失效并提高数据访问的局部性。在深度学习和其他数组密集型计算任务中，Tiling技术通过将大的数据结构（如矩阵、张量）分解成更小的“块”或“tiles”，使得每个小块可以独立地装入CPU或GPU的快速访问缓存中进行处理。这样，计算设备可以更快地访问和处理数据，因为它减少了内存访问时间和数据在内存层级中的传输量。

工作原理:

1. 数据分割：原始的大矩阵或数组被划分为较小的子矩阵或子数组。
2. 局部性优化：每个tile足够小，以至于可以完全装入CPU/GPU的L1/L2缓存中，减少访问主存储器的需求。
3. 并行处理：每个tile可以在不同的处理单元上并行处理，提高处理效率。

Recomputation

Recomputation技术，又称为重计算或检查点技术，是一种减少神经网络训练过程中内存使用的技术。在神经网络训练过程中，为了进行反向传播，需要存储前向传播的所有中间激活值。这些值在计算梯度时被用来应用链式法则，但它们占用大量内存。Recomputation技术的核心思想是在反向传播过程中不存储所有的中间值，而是在需要这些值的时候重新计算它们，从而节省内存。

工作原理:

1. 选择存储点：在前向传播过程中，只选择某些关键层（通常是深度网络中的几个层）存储中间结果，其余层的中间结果不存储。
2. 重计算：在执行反向传播时，对于未存储中间结果的层，重新执行前向传播的计算来获得所需的中间值。
3. 内存与计算权衡：这种方法减少了内存的使用，但增加了计算量，因为某些计算需要执行两次。

Rotary Position Embedding (RoPe)

RoPe 的工作原理主要基于复数运算。首先，将 Q 矩阵的第 m 个字的表示向量 qm 和 K 矩阵的第 n 个字的表示向量 kn 看作是复数。然后，对 qm 和 kn 进行变换：

f(q, m) = qm * e^(jmθ)

f(k, n) = kn * e^(jnθ)

其中，e^(jmθ) 和 e^(jnθ) 是位置 m 和 n 的绝对位置编码信息。然后，将变换后的数据做内积：

<f(q, m), f(k, n)> = f(q, m) * f(k, n)^*

其中，<> 表示内积计算，^* 表示复数的共轭。通过这样的计算，我们可以得到 (m - n) 的相对位置信息，而初始的变换也只是绝对位置信息的编码。因此，RoPe 实现了用绝对位置编码来表征相对位置编码。

应用

在标准的Transformer模型（包括UniLM）中，attention机制是通过以下公式计算的：

计算Query，Key，Value：Q = W_q * X，K = W_k * X，V = W_v * X

计算attention score：S = softmax(Q * K^T / sqrt(d_k))

计算output：O = S * V

其中，X是输入，W_q，W_k，W_v是权重矩阵，d_k是Key的维度。

在FlashAttention中，softmax计算被分解为分块计算，每个块的元素最大值和分子、分母的计算，然后再合并。具体来说，假设输入矩阵X被分为两块X^(1)和X(2)，那么分块softmax的计算过程如下：

计算每个块的元素最大值：m(X) = max([m(X^(1)), m(X^(2))])

计算分块softmax的分子：f(X) = [e^(m(X(1)) - m(X)) * f(X^(1)), e^(m(X(2)) - m(X)) * f(X^(2))]

计算分块softmax的分母：l(X) = [e^(X(1) - m(X)) * l(X^(1)), e^(X(2) - m(X)) * l(X^(2))]

计算softmax：softmax(X) = f(X) / l(X)

Multi Query Attention是Multi-Head Attention的变体形式，其中，Key和Value只有一个头，Query是多头。通过共享Key得到attention score时，MQA具有以下优势：更加省内存，计算量更低。在实现上，MQA和MHA主要是在计算Key和Value的过程中有计算量的差异，由于训练阶段由于数据是并行的，这种差异整体不明显，而在推理阶段，在memory cache的基础上，MQA中每个step的V的计算量为 dk，而MHA中计算量则为 dkh。

在MQA中，attention score的计算过程如下：

计算Query，Key，Value：Q = W_q * X，K = W_k * M，V = W_v * M

计算attention score：logits = Q * K^T

计算weights：weights = softmax(logits + mask)

计算output：O = weights * V

其中，X和M是输入，W_q，W_k，W_v是权重矩阵，mask是掩码矩阵。

在标准的Transformer模型中，位置编码通常是通过将一个位置编码向量添加到输入序列的每个元素上来实现的。这种方法可以有效地为模型提供序列中元素的位置信息，但是它不能直接处理相对位置信息。

RoPE的提出是为了解决这个问题。具体来说，RoPE将位置信息编码为一个旋转矩阵，然后将这个旋转矩阵应用到输入序列的每个元素上。这样，模型就可以直接处理相对位置信息，而不仅仅是绝对位置信息。

RoPE的公式如下：

计算位置编码：PE = W * pos

计算旋转矩阵：R = exp(i * PE)

计算旋转后的输入：X’ = X * R

其中，pos是位置向量，W是权重矩阵，i是虚数单位，X是输入，'*'表示矩阵乘法，'exp’表示指数函数。

在实际使用中，我们发现不仅要求LLM能够识别单词和短语，更需要深入理解长文本序列，以有效捕获上下文中的复杂关系和语义。然而，传统的位置编码方法在处理不同长度序列时可能受到限制，导致模型难以准确捕捉长距离的依赖关系。为了克服这一挑战，我们引入了一项创新技术，即“自适应正弦余弦位置编码”，以提升模型在长序列建模方面的能力。

我们的方法采用了一种欺骗模型的思路，即根据输入序列的长度动态地调整位置编码的频率，从而更好地适应序列的变化，让模型认为其仍然在有效的长度内。

具体来说，我们通过计算自适应基数（base）来实现频率的调整，即

alpha=base /(1024−1)

base=base ×alpha^(d/(d−2))

其中，L 代表输入序列的长度，d 表示模型的维度，base 是初始基数。然后，我们根据自适应基数计算位置编码的频率，

freq=1/(base ^(i/d))

这样一来，每个位置的编码频率都会因序列长度的变化而不同，从而使模型能够更好地适应不同长度的序列。

在实际应用中，当输入序列的长度超过之前计算的最大长度时，我们会重新计算位置编码的频率，并生成新的位置编码。然而，对于较短的序列，我们仍然使用之前计算好的位置编码，以保持计算效率。这种策略保证了模型在处理不同长度的序列时能够充分应用自适应位置编码的优势。