Jamba-1.5: Hybrid Transformer-Mamba Models at Scale
目录
0. 摘要
我们推出了 Jamba-1.5,这是一种基于我们 Jamba 架构的新指令微调大型语言模型。Jamba 是一种混合 Transformer-Mamba 的 MoE 架构,能够在不同上下文长度中提供高吞吐量和低内存使用,同时保持与 Transformer 模型相同或更好的质量。我们发布了两种模型尺寸:Jamba-1.5-Large,具有 94B 有效参数,和 Jamba-1.5-Mini,具有 12B 有效参数。这两个模型都经过微调,以支持各种对话和指令跟随能力,并且具有 256K tokens 的有效上下文长度,是开放权重模型中最长的。为了支持成本效益的推理,我们引入了 ExpertsInt8,这是一种新颖的量化技术,能够在处理 256K tokens 上下文时,将 Jamba-1.5-Large 适配于一台配备 8 个 80GB GPU 的机器上运行,并且不会损失质量。在一系列学术和聊天机器人基准测试中,Jamba 模型表现优异,提供高吞吐量并在长上下文基准测试中优于其他开放权重模型。
2. 模型架构
(2024,Attention-Mamba,MoE 替换 MLP)Jamba:混合 Transformer-Mamba 语言模型
(2023,SSM,门控 MLP,选择性输入,上下文压缩)Mamba:具有选择性状态空间的线性时间序列建模
Jamba-1.5-Large 基于我们开发的 Jamba [24] 混合解码器架构,该架构结合了 Transformer 层 [36] 与 Mamba 层 [13],一种状态空间模型 (SSM) [14, 15],以及专家混合 (MoE) 模块 [8, 34]。有关此架构的详细描述,请参见 [24]。
在 Jamba [24] 的工作中,我们发现 Transformer、Mamba 和 MoE 元素的组合有助于在吞吐量、内存使用和质量之间实现平衡。Jamba-1.5-Large 在更大规模上展示了这种灵活性。
Jamba-1.5-Large 遵循相同的 Jamba 结构,但具有更大的容量。它拥有 94B 有效参数,总参数量为 398B。该模型有 9 个模块,每个模块具有以下规格:
- 每个模块有 l = 8 层。
- a : m = 1 : 7 的注意力层与 Mamba 层的比例。在我们关于 Jamba 的研究中发现这个比例是最优的 [24],后续的工作也证实了类似的比例成功 [6, 37]。
- 每 e = 2 层使用 MoE 替代单一的 MLP。共有 n = 16 个专家,每个 token 选择最优的 K = 2 个。
- 隐藏状态的维度为 8192。
- 注意力查询头的数量为 64,KV 头的数量为 8。
表 1 将 Jamba-1.5 模型与相似规模的公开模型进行了比较。Jamba-1.5-Mini 的有效参数数量与 Mixtral 8x7B 相当,而 Jamba-1.5-Large 的有效参数数量介于 LLaMA-3.1-70B 和 Mistral-Large-2 之间。同时,Jamba 模型在 KV 缓存内存使用量(在 256K tokens 上下文下)方面远小于所有其他模型,相比同类模型减少了大约一个数量级。
在这些设置下,并结合我们的专门量化技术(第 3.1 节),Jamba-1.5-Large 可以在一台配备 8 个80GB GPU 的机器上运行,支持长达 256K tokens 的上下文长度。
(2024|ICML,Mamba2,SSD,SSM,SMA,矩阵变换,张量收缩,张量并行)Transformer 是 SSM
对于此次发布,我们还尝试了 Mamba-2 [6],这是 Mamba 的一个更快且改进的版本,据报道它在性能上超越了单独使用 Mamba 和 Transformers 的模型。然而,正如图 1 所示,我们发现,在混合架构中,Mamba-1-Attention 组合的效果优于 Mamba-2-Attention,因此我们在 Jamba-1.5-Large 中使用了 Mamba-1。我们还发现混合架构的性能优于纯 Mamba-2。我们推测这可能是因为 Mamba-2 相比于 Mamba-1 的一些优势(特别是使用更大状态尺寸的能力)在 Mamba 层之间交错全注意力层时不那么显著,因为全注意力层可以从整个上下文中汇聚信息。
3. 服务考量与改进
我们分享了一些见解和改进,旨在实现大规模高效服务 Jamba 模型。
3.1 ExpertsInt8 量化
为了支持 Jamba-1.5-Large 的高效服务,我们开发了一种新型量化技术,称为 ExpertsInt8。我们观察到,超过 85% 的模型权重位于 MoE 层,超过 90% 位于 MoE 或 MLP 层。我们希望在保持快速 BF16 内核的好处的同时,对这些权重进行量化。为此,我们将 MoE 和 MLP 权重量化为 INT8,存储为 INT8,并在实际计算之前将它们解量化回 BF16。重要的是,解量化步骤直接发生在 vLLM [18] 中的 fused_moe 内核内部。这样,解量化过程增加的开销微乎其微,甚至在延迟上比 BF16 更有优势。【我们将这归因于内核在相对较小的权重和激活块上操作,这些块在执行计算之前从 GPU HBM 移动到 SRAM。在我们的实现中,当权重量化为 int8 时,它们从 HBM 移动到 SRAM,因此由于内存占用减少了一半,所需时间也减少了。】
我们已经将修改后的 fused_moe 内核贡献给了 vLLM。【拉取请求在此处:https://github.com/vllm-project/vllm/pull/7415】
我们的 ExpertsInt8 方法具有几个优势。
- 首先,它非常快速;量化过程仅需在模型加载时几秒钟。
- 其次,与 vLLM 中大多数其他技术不同,它不依赖于需要数小时或数天且可能不稳定的校准过程。
- 第三,我们仍然可以使用 BF16 来处理大规模激活。
- 第四,它可以在 A100 GPU 上使用,而 FP8 仅在 H100 上可用。
- 最后,我们的量化在延迟上与 FP8 相当,同时超越了其他量化技术,并且不会导致质量损失。
图 2 比较了使用不同量化技术的延迟,包括 Jamba-1.5-Mini、Jamba-1.5-Large 和两个 Mixtral 模型(8x78B 和 8x22B)。在 H100 GPU 上,ExpertsInt8 的延迟与 FP8 相匹配。在 A100 上,由于 FP8 不可用,ExpertsInt8 是一种有吸引力的技术,显著超越了 GPTQ [9]。结合上述 ExpertsInt8 的优点,这使得它成为服务大型 MoE 模型的一个有吸引力的量化技术。
3.2 激活损失
在预训练过程中,我们发现某些激活值,特别是特定专家的输出以及最后的 Mamba 层的输出,在处理特定输入 token 时,逐渐增大,最终达到高达 4 × 10^6 的值。尽管我们发现这并未对使用 BF16 精度进行的预训练造成损害,但这些激活值的幅度可能会在推理过程中引发数值问题,因为一些量化库仅支持 FP16 精度,而 FP16 的最大范围为 64K。
为了解决这些问题,我们添加了一个“激活损失”(Activation Loss)项,其值与前向传播中激活值的均方值成正比,并设有可配置的 α 因子,以惩罚较大的激活值。通过实验,我们发现这种辅助损失对训练没有影响,即使 α 值达到至少 10^{−3}。对于 Jamba-1.5-Large,我们使用了 α = 10^{−5},这足以将激活值减少到一个可接受的范围(最大 2K-3K)。此外,添加这一辅助损失几乎瞬间降低了激活值,使得它仅在训练结束时添加也不会影响训练速度和质量。
为了验证这种方法,我们在模型上使用 FP16 激活值运行了完整的评估套件,结果与使用 BF16 的评估结果相同,没有出现 NaN/溢出。
4. 吞吐量和延迟分析
得益于混合 Jamba 架构,我们的 Jamba-1.5 模型提供了出色的吞吐量和延迟性能。图 3 和图 4 分别展示了 Jamba-1.5-Mini 和 Jamba-1.5-Large 的表现。如图所示,我们的模型在延迟和吞吐量方面均显著优于相同规模的模型。它们在处理长上下文时展现出显著优势,存在较大的性能差距。重要的是,Jamba-1.5-Large 在处理长上下文时依然高效,而大型的 LLaMA3-405B 不能在相同硬件上运行。【注:Large 比 Mini 有更高的时延和更低的吞吐量,直观理解就是速度换性能】
5. 训练
5.1 训练基础设施和数据
Jamba-1.5-Large 在 NVIDIA H100 GPU 上训练,使用我们内部开发的专有框架,包括 FSDP、张量并行、序列并行和专家并行。对于专家并行,我们适配了 MegaBlocks [10]。
5.2 训练阶段
该模型的训练分为三个阶段。
- 在预训练阶段,模型首先在我们内部的数据集上进行训练,该数据集最后更新于 2024 年 3 月。我们的预训练数据集是公开的网页文档、代码、书籍和科学文章的混合体。我们的预处理流程包括解析、质量过滤和去重。为了最大化利用公开数据,我们开发了自己的解析器,并使用它来提取文本和格式。数据混合的具体组成通过各种消融实验确定。该阶段包括多语言数据,重点关注以下语言:英语、西班牙语、法语、葡萄牙语、意大利语、荷兰语、德语、阿拉伯语和希伯来语。
- 然后,模型进行了一个短的中期训练,重点训练长文档,以强调其长距离能力。
- 最后,模型经过了后训练(post-training),如下一节所述。
5.3 后训练
我们后训练的方法旨在同时实现两个目标:(i)为模型提供各种技能和对话能力;(ii)保留预训练中的能力,特别是中期训练中的长上下文能力。这两个目标部分存在冲突,因为大多数现有的后训练数据集包含相对较短的示例。
考虑到这些因素,我们的后训练过程包括在高质量对话数据、特定技能数据和长上下文数据上进行监督微调 [32, 39]。混合这些不同类型的数据旨在保留长上下文能力并获得所需的技能。如下面的评估所示,我们发现我们的模型在长上下文评估中表现非常好。
在进行监督微调时,我们大量使用合成数据,这在最近的基础模型中很常见 [7],并反映了我们构建复合 AI 系统 [20] 的结构化数据的方法。我们开发了多种不同的数据合成流程,针对不同的模型能力。所有流程都应用以下模式:(i)在目标分布中采样或生成提示;(ii)从语言模型中生成响应;(iii)根据自动验证和评分过滤或排名响应;(iv)后编辑以去除伪影并适应所需格式。我们使用不同的模型、提示、采样、过滤和编辑来处理不同的数据管道,从而组成最终的数据混合体。
我们基于大量主要内部的自动化指标选择了最终的训练配方(数据混合和超参数)。两个 Jamba-1.5 模型使用相同的控制 token 和格式模板进行微调,我们将这些作为我们发布的一部分,提供 HF 兼容的标记器和聊天模板;有关详细信息,请参见模型卡。
以下是一些合成数据生成的显著示例:
-
基于表格的 QA:我们生成表格数据及其相应的问题-答案对,如我们在表格理解 [20] 中所示。然后,我们使用语言模型将表格转换为自然语言段落。我们生成的训练示例包括提取、聚合和归属任务,涉及给定表格中特定行或列的文本。
-
文档 QA:给定一个文档,我们提示语言模型生成问题-答案对,适用于单个或多个段落。我们有时通过添加类似的文本将这些示例嵌入更长的上下文中,以鼓励长上下文理解和归属。
-
工具使用:我们使用开源的 Glaive 函数调用数据集 [1] 作为起点,通过各种启发式方法和对输出模式的验证进行过滤。为了支持并行函数调用,我们首先为 Glaive 中的每个函数生成多个有效的参数分配。接下来,我们从这些有效的参数分配中采样子集,针对相同函数和不同函数,生成与函数调用集对应的用户请求。最后,我们提示函数调用语言模型对这些生成的用户请求做出响应,并仅保留与原始参数分配匹配的函数调用的响应。
-
可控性:我们定义了一组可以轻松验证的指令,并合成了包括一个或多个约束的通用文档草拟任务的提示。我们从语言模型中生成这些提示的完成,并基于对我们细粒度指令的验证以及通用奖励模型进行拒绝采样。为了支持系统消息中的指令,我们选择了多种这种类型的提示,这些提示共享细粒度指令实例,并将这些提示重新格式化为多轮对话,将指令移到系统消息中。
5.4 一些观察
我们分享了一些关于 Jamba-1.5 开发过程中的观察。这些观察虽然尚未完全探索,但希望能够激发社区进一步研究这些问题。首先,虽然我们仅在后训练阶段包含了非常小比例的非英语数据,仅针对几个语言和特定技能,我们的 Jamba-1.5 模型在多语言环境下表现相当出色。我们确实在预训练阶段包括了多语言数据,如上所述。因此,我们推测模型能够在主要以英语进行后训练时,利用预训练阶段学到的知识。
其次,我们高效的 Jamba 架构降低了在长上下文上的微调成本,使我们能够在给定预算下进行更多实验。因此,我们能够在后训练阶段实验多种不同的训练配方。
最后,尽管偏好微调算法如 PPO [33] 或 DPO [29] 可以改善模型输出与人类意图之间的对齐,我们发现细致的合成数据生成、数据过滤和监督微调的结合对获得强大的后训练模型至关重要。
6. 评估
虽然我们相信基准测试仅与实际应用的成功和用户满意度部分相关,但我们仍报告了一些关键公共基准的结果。首先,我们报告标准学术基准的结果。然后,我们评估模型在聊天机器人基准上的表现。最后,我们在多个长上下文评估和多语言评估中评估 Jamba-1.5-Large。
我们将其与最近的相同规模范围的开放权重模型进行比较:在比较 Jamba-1.5-Large 时,与 LLaMA-3.1 70B 和 Mistral-Large-2-123B;在比较 Jamba-1.5-Mini 时,与 LLaMA-3.1-8B 和 Gemma-2-9B 进行比较。
6.1 学术基准
6.2 聊天机器人评估
在本节中,我们对 Jamba-1.5 模型在两个聊天机器人(chatbot)场景中的表现进行了吞吐量评估:
- Arena-Hard [22],这是一个包含 500 个具有挑战性的用户查询的数据集,使用 GPT4-Turbo 作为评判;
- WildBench [25],该数据集也使用 GPT4-Turbo 作为评判,但进行了长度偏差的缓解处理。
6.3 长上下文评估
我们在 RULER 基准上进行评估,RULER 是一组 13 个合成任务,旨在评估语言模型的长上下文能力。RULER 包括 8 种变体的针在大 haystack 检索任务 [17, 21, 27, 28],其中包含多个“针” [2]。此外,它还有一个变量跟踪任务,需要返回一系列变量绑定,两个聚合任务,需要返回最常见的词汇,以及两个问答任务,其中包含来自自然数据集 [30, 41] 的答案的段落被插入到随机段落中,以模拟长上下文。
接下来,我们在 ∞BENCH 数据集上进行评估,该数据集旨在评估语言模型的长上下文能力,平均长度为 100K 词汇。我们专注于两个英文任务,理解长篇小说:问答(EN.QA)和多项选择题问答(EN.MC)。如表 5 所示,Jamba-1.5 模型在这方面表现非常出色,超越了同样规模的 LLaMA-3.1 和 Mistral-Large-2 模型。(由于 Gemma-2 9B 的上下文窗口较短(8K),我们未报告其结果。)
6.4 多语言能力
我们对 Jamba-1.5 在非英语语言中的能力进行了基本评估。具体来说,我们报告了在多语言 MMLU 数据集 [19] 上的结果,该数据集通过 LM Evaluation Harness [11] 分发。如表 6 所示,Jamba-1.5-Mini 的表现与对比模型相当或更好。Jamba-1.5-Large 略微落后于其可比模型,但仍展现出良好的多语言能力。
8. 结论
我们介绍了 Jamba-1.5-Large 和 Jamba-1.5-Mini,这两个基于 Jamba 混合 Transformer-Mamba 架构的大规模模型。两个模型在学术基准、聊天机器人评估和长上下文评估中均表现出色,同时提供了改进的延迟和吞吐量,特别是在处理长上下文时。我们发布了模型权重,希望社区能够使用这些模型并在此技术基础上进行进一步开发。
论文地址:https://arxiv.org/abs/2408.12570
Jamba 开源模型许可:https://www.ai21.com/licenses/jamba-open-model-license。
项目页面:https://huggingface.co/ai21labs
Jamba-1.5-Mini:https://huggingface.co/ai21labs/AI21-Jamba-1.5-Mini
Jamba-1.5-Large:https://huggingface.co/ai21labs/AI21-Jamba-1.5-Large
公和众与号:EDPJ(进 Q 交流群:922230617 或加 VX:CV_EDPJ 进 V 交流群)
加 VX 群请备注学校 / 单位 + 研究方向
CV 进计算机视觉群
KAN 进 KAN 群