该存储库包含 LLaMA 2(大型语言模型元 AI)模型的实现,这是一种生成预训练 Transformer (GPT) 变体。实现重点在于模型架构和推理过程。代码经过重组并添加了大量注释,以便于轻松理解架构的关键部分。
RMS-Normalization: RMSNorm是原始层归一化(LayerNorm)的简化。 LayerNorm 是一种正则化技术,可以处理内部协变量偏移问题,从而稳定层激活并提高模型收敛性。它在 LLaMA 2 中被证明是相当成功的。
激活函数: LLaMA 2 使用 SwiGLU 激活函数而不是 ReLU,从而提高了训练性能。
旋转位置嵌入 (RoPE):受 GPT-Neo-X 项目的启发,LLaMA 2 在每一层都加入了旋转位置嵌入,增强了模型的位置理解。
增加上下文长度和分组查询注意力 (GQA): LLaMA 2 模型具有双倍上下文窗口(从 2048 到 4096 个标记)并采用分组查询注意力。这可以更好地处理长文档、聊天历史和摘要任务。
KV 缓存是此实现中采用的一项关键优化技术,用于加速语言模型 (LM) 解码的推理过程。在自回归解码过程中,每个标记都是根据先前的标记进行预测的,模型中的自注意力是因果关系。这意味着令牌的表示仅基于其自身和先前的令牌而不是未来的令牌来计算。
在自注意力中,使用键、值和查询投影来投影输入序列。 KV 缓存有效地存储键和值投影的结果,消除了未来解码迭代中冗余计算的需要。因此,可以从缓存中检索在自回归解码期间保持固定的令牌表示,从而显着提高推理速度。
这种 KV 缓存技术是一项关键的架构功能,可提高 LLaMA 模型在解码过程中的效率和速度。
LLaMA 2 模型融合了 Shazeer (2019) 提出的多查询注意力 (MQA) 概念的变体,这是多头注意力 (MHA) 算法的改进。 MQA 提高了注意力机制的效率,同时保持了最小的准确性下降。
在传统的多头注意力中,整个注意力计算被复制 h 次,其中 h 是注意力头的数量。然而,GQA 通过从 K 和 V 值中删除或显着减少头部尺寸 (h) 来减少计算冗余。在 MQA 中,查询值 (Q) 的每个“头”都会经历相同的 K 和 V 变换,从而优化注意力计算。
这种改进带来了与 MHA 类似的计算性能,但显着减少了从内存读取/写入的数据量。因此,GQA 提高了性能(通过增加算术强度)和内存空间效率(通过减少存储的 KV 缓存数据量),使其成为 LLaMA 架构的一个有价值的补充。
在 LLaMA 2 模型中,旋转位置嵌入 (RoPE) 通过将位置信息合并到令牌表示中,在增强注意力机制方面发挥着至关重要的作用。 “注意力”的概念很强大,但为了确保计算出的注意力有意义,代币需要有一个位置的概念。
位置嵌入有两种主要类型:绝对和相对。绝对位置嵌入对输入短语中单词的绝对位置进行编码,而相对位置嵌入对两个单词之间的相对位置进行编码。这些嵌入提供了重要的位置信息,帮助标记理解它们在序列中的上下文。
旋转位置嵌入采用独特的方法,利用旋转矩阵来嵌入位置信息。目标是确保向量 q 和 k 在位置 m 和 n 处的内积仅取决于 q、k 及其相对距离 (m — n)。旋转矩阵(其中角度是向量的位置)通过矩阵乘法嵌入到原始向量中,与此标准对齐。
这种整合位置信息的创新方法增强了模型理解令牌关系和上下文的能力,有助于改进注意力机制。
model.py
:包含 LLaMA 变压器模型的实现,以及解释每个组件和功能的详细注释。
inference.py
:演示如何使用经过训练的 LLaMA 模型进行推理,提供对输入和输出处理的深入了解。
请随意探索代码,纠正任何错误,并尝试 LLaMA 2 模型!