[论文解读] Jamba-1.5: Hybrid Transformer-Mamba Models at Scale

Jamba Team, Barak Lenz+59 more|arXiv (Cornell University)|Aug 22, 2024

Fluid Dynamics and Vibration Analysis被引用 4

一句话总结

Jamba-1.5 引入了两种指令微调的大型语言模型——Jamba-1.5-Mini（12B 激活参数）和 Jamba-1.5-Large（94B 激活参数），基于混合 Transformer-Mamba 的专家混合（MoE）架构，在 256K token 上下文长度的长上下文基准测试中实现了最先进性能，具有低内存占用（9GB KV 缓存）和高吞吐量。该模型通过一种新颖的量化技术 ExpertsInt8 实现高效推理，可在不损失质量的前提下部署于 8×80GB GPU 上。

ABSTRACT

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.

研究动机与目标

• 开发可处理长上下文输入的大规模、高效且高性能的指令微调语言模型。
• 在不牺牲质量的前提下，降低大模型的内存占用并提升推理效率。
• 通过新颖的量化技术和架构创新，实现在标准硬件（8×80GB GPU）上部署如 Jamba-1.5-Large 这类大模型。
• 证明混合 Transformer-Mamba-MoE 架构在大规模场景下，尤其在长上下文设置中，优于纯 Transformer 或 Mamba 模型。
• 在 Jamba Open Model License 下发布开放权重模型，确保完全透明并符合 OECD 人工智能原则。

提出的方法

• Jamba-1.5 架构结合了 Transformer 层、Mamba 状态空间层以及包含 16 个专家和每 token 选择 top-2 的专家混合（MoE）模块。
• 每个模块采用 1:7 的注意力层与 Mamba 层比例，MoE 每隔 2 层应用一次，隐藏维度为 8192。
• 提出 ExpertsInt8 量化技术，针对 MoE 层中 85% 的权重，实现无质量损失的 8 位量化。
• 模型在 1000B tokens 数据上进行训练，并在多样化基准上进行指令遵循和对话能力的微调。
• 通过利用 Mamba 的高效序列建模能力和 MoE 的稀疏激活机制，该架构实现了高吞吐量和低 KV 缓存占用（在 256K tokens 时为 9GB）。
• 评估在学术基准（如 MMLU、GSM8K）、聊天机器人基准以及长上下文任务（如 RULER）上进行，与 LLaMA-3、Mistral-Large-2 和 Mixtral 进行对比。

(a) 350M.

实验结果

研究问题

• RQ1 混合 Transformer-Mamba-MoE 架构能否在保持或提升模型质量的前提下扩展至 94B 激活参数？
• RQ2 在大规模场景下，Mamba 与 Transformer 层的混合结构是否优于纯 Mamba-2 或仅使用 Transformer 的模型？
• RQ3 ExpertsInt8 量化能否实现 MoE 层的完整 8 位量化且不造成性能下降，从而支持在 8×80GB GPU 上部署？
• RQ4 与其它开放权重模型相比，该模型在长上下文任务中，特别是在 256K token 上下文长度下的表现如何？
• RQ5 Jamba-1.5 模型在多语言任务中，尤其是在低资源语言上的泛化能力如何？

主要发现

• Jamba-1.5-Large 实现了 94B 激活参数，在使用 ExpertsInt8 量化技术时，可在单台配备 8×80GB GPU 的机器上以 256K 上下文长度运行，且无质量损失。
• 在 256K tokens 时，KV 缓存内存占用降低至 9GB，相比 LLaMA-3.1-70B 和 Mistral-Large-2 减少了 10 倍。
• 在 RULER 基准测试中，Jamba-1.5 是唯一实现有效上下文长度达 256K token 的开放权重模型。
• Jamba-1.5-Mini 在多语言 MMLU 上取得 64.3% 的准确率，优于 LLaMA-3.1-8B 和 Gemma-9B；Jamba-1.5-Large 在同一基准上达到 73.94%。
• 在长上下文评估中，Jamba-1.5-Large 超越 LLaMA-3.1-70B 和 Mistral-Large-2，展现出更优的吞吐量和更低的延迟。
• ExpertsInt8 量化被证明在 MoE 层上有效，可在保持模型质量的同时，实现面向消费级硬件的高效部署。

(b) 1.3B.