Jamba-1.5: Hybrid Transformer-Mamba Models at Scale

📄 arXiv: 2408.12570v1 📥 PDF

作者: Jamba Team, Barak Lenz, Alan Arazi, Amir Bergman, Avshalom Manevich, Barak Peleg, Ben Aviram, Chen Almagor, Clara Fridman, Dan Padnos, Daniel Gissin, Daniel Jannai, Dor Muhlgay, Dor Zimberg, Edden M Gerber, Elad Dolev, Eran Krakovsky, Erez Safahi, Erez Schwartz, Gal Cohen, Gal Shachaf, Haim Rozenblum, Hofit Bata, Ido Blass, Inbal Magar, Itay Dalmedigos, Jhonathan Osin, Julie Fadlon, Maria Rozman, Matan Danos, Michael Gokhman, Mor Zusman, Naama Gidron, Nir Ratner, Noam Gat, Noam Rozen, Oded Fried, Ohad Leshno, Omer Antverg, Omri Abend, Opher Lieber, Or Dagan, Orit Cohavi, Raz Alon, Ro'i Belson, Roi Cohen, Rom Gilad, Roman Glozman, Shahar Lev, Shaked Meirom, Tal Delbari, Tal Ness, Tomer Asida, Tom Ben Gal, Tom Braude, Uriya Pumerantz, Yehoshua Cohen, Yonatan Belinkov, Yuval Globerson, Yuval Peleg Levy, Yoav Shoham

分类: cs.CL, cs.LG

发布日期: 2024-08-22

备注: Webpage: https://www.ai21.com/jamba


💡 一句话要点

Jamba-1.5:基于混合Transformer-Mamba架构的大规模语言模型,实现高吞吐和低内存占用。

🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture) 支柱九:具身大模型 (Embodied Foundation Models)

关键词: 大规模语言模型 混合架构 Transformer Mamba 混合专家模型 量化 长上下文 指令微调

📋 核心要点

  1. 现有Transformer模型在处理长文本时面临计算和内存瓶颈,限制了其在长上下文任务中的应用。
  2. Jamba架构结合Transformer和Mamba的优势,利用混合专家模型,在保证性能的同时,显著降低内存占用和提高吞吐量。
  3. Jamba-1.5模型在长上下文基准测试中优于其他开放权重模型,并提出了ExpertsInt8量化技术,降低了部署成本。

📝 摘要(中文)

本文介绍了Jamba-1.5,这是一种基于Jamba架构的新型指令微调大型语言模型。Jamba是一种混合Transformer-Mamba混合专家架构,可在各种上下文长度上提供高吞吐量和低内存使用率,同时保持与Transformer模型相同或更好的质量。我们发布了两种模型尺寸:Jamba-1.5-Large,具有940亿个激活参数,以及Jamba-1.5-Mini,具有120亿个激活参数。两种模型都针对各种对话和指令遵循能力进行了微调,并具有256K tokens的有效上下文长度,这是开放权重模型中最大的。为了支持经济高效的推理,我们引入了ExpertsInt8,这是一种新颖的量化技术,允许在处理256K-token上下文时,将Jamba-1.5-Large安装在具有8个80GB GPU的机器上,而不会损失质量。在对一系列学术和聊天机器人基准进行评估时,Jamba-1.5模型取得了优异的成绩,同时提供了高吞吐量,并在长上下文基准测试中优于其他开放权重模型。两种尺寸的模型权重均已根据Jamba开放模型许可公开提供,并且我们将ExpertsInt8作为开源发布。

🔬 方法详解

问题定义:现有的大型语言模型,特别是基于Transformer的架构,在处理超长上下文时,计算复杂度和内存需求急剧增加,导致推理速度慢,硬件成本高昂。这限制了它们在需要理解和生成长篇文本的应用中的实用性。

核心思路:Jamba的核心思路是结合Transformer和Mamba架构的优势。Transformer擅长建模全局依赖关系,而Mamba在处理长序列时具有更高的效率和更低的内存占用。通过混合使用这两种架构,并结合混合专家(MoE)模型,Jamba旨在在性能、吞吐量和内存效率之间取得平衡。

技术框架:Jamba的整体架构是一个混合Transformer-Mamba混合专家模型。具体来说,模型可能包含多个层,其中一些层是标准的Transformer层,另一些层是Mamba层。混合专家模型允许模型根据输入的不同部分激活不同的专家网络,从而提高模型的容量和效率。模型还采用了ExpertsInt8量化技术,进一步降低了内存占用。

关键创新:Jamba的关键创新在于其混合Transformer-Mamba架构和ExpertsInt8量化技术。混合架构允许模型利用Transformer的全局建模能力和Mamba的长序列处理效率。ExpertsInt8量化技术则可以在不损失模型性能的前提下,显著降低模型的内存占用,使得大规模模型的部署更加可行。

关键设计:关于关键设计细节,论文提到了256K tokens的有效上下文长度,这是开放权重模型中最大的。此外,ExpertsInt8量化技术允许在8个80GB GPU上运行Jamba-1.5-Large模型,同时处理256K-token上下文,而不会损失质量。具体的Transformer和Mamba层数、混合专家模型的专家数量、以及量化技术的具体实现细节可能在论文中有更详细的描述。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

Jamba-1.5模型在学术和聊天机器人基准测试中取得了优异的成绩,并在长上下文基准测试中优于其他开放权重模型。ExpertsInt8量化技术使得Jamba-1.5-Large模型可以在8个80GB GPU上运行,同时处理256K-token上下文,而不会损失质量,显著降低了部署成本。

🎯 应用场景

Jamba-1.5模型适用于需要处理长文本的各种应用场景,例如长篇文档摘要、代码生成、对话系统、以及需要理解和生成长篇故事或剧本的创意写作。其高吞吐量和低内存占用使其更易于部署在资源受限的环境中,并降低了大规模语言模型的应用门槛。

📄 摘要(原文)

We present Jamba-1.5, new instruction-tuned large language models based on our Jamba architecture. Jamba is a hybrid Transformer-Mamba mixture of experts architecture, providing high throughput and low memory usage across context lengths, while retaining the same or better quality as Transformer models. We release two model sizes: Jamba-1.5-Large, with 94B active parameters, and Jamba-1.5-Mini, with 12B active parameters. Both models are fine-tuned for a variety of conversational and instruction-following capabilties, and have an effective context length of 256K tokens, the largest amongst open-weight models. To support cost-effective inference, we introduce ExpertsInt8, a novel quantization technique that allows fitting Jamba-1.5-Large on a machine with 8 80GB GPUs when processing 256K-token contexts without loss of quality. When evaluated on a battery of academic and chatbot benchmarks, Jamba-1.5 models achieve excellent results while providing high throughput and outperforming other open-weight models on long-context benchmarks. The model weights for both sizes are publicly available under the Jamba Open Model License and we release ExpertsInt8 as open source.