[Paper Review] 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 Analysis4 citations

TL;DR

Jamba-1.5 introduces two instruction-tuned large language models—Jamba-1.5-Mini (12B active parameters) and Jamba-1.5-Large (94B active parameters)—based on a hybrid Transformer-Mamba mixture-of-experts architecture, achieving state-of-the-art performance on long-context benchmarks with 256K token context length, low memory usage (9GB KV cache), and high throughput. The model is optimized for efficient inference via a novel quantization technique, ExpertsInt8, enabling deployment on 8x80GB GPUs without quality loss.

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.

Motivation & Objective

• To develop large-scale, efficient, and high-performing instruction-tuned language models capable of handling long-context inputs.
• To reduce memory usage and improve inference efficiency in large models without sacrificing quality.
• To enable deployment of large models like Jamba-1.5-Large on standard hardware (8x80GB GPUs) via novel quantization and architectural innovations.
• To demonstrate that hybrid Transformer-Mamba-MoE architectures outperform pure Transformer or Mamba models at scale, especially in long-context settings.
• To release open-weight models under the Jamba Open Model License with full transparency and alignment to OECD AI principles.

Proposed method

• The Jamba-1.5 architecture combines Transformer layers, Mamba state-space layers, and a mixture-of-experts (MoE) module with 16 experts and top-2 gating per token.
• Each block uses an 1:7 ratio of attention to Mamba layers, with MoE applied every 2 layers, and a hidden dimension of 8192.
• ExpertsInt8 quantization is introduced, targeting the 85% of weights in MoE layers, enabling 8-bit quantization with no quality degradation.
• The model is trained on 100B tokens and fine-tuned for instruction-following and conversational capabilities across diverse benchmarks.
• The architecture achieves high throughput and low KV cache usage (9GB at 256K tokens) by leveraging Mamba’s efficient sequence modeling and MoE’s sparse activation.
• Evaluation is conducted on academic benchmarks (e.g., MMLU, GSM8K), chatbot benchmarks, and long-context tasks like RULER, with comparisons to LLaMA-3, Mistral-Large-2, and Mixtral.

(a) 350M.

Experimental results

Research questions

• RQ1 Can a hybrid Transformer-Mamba-MoE architecture scale to 94B active parameters while maintaining or improving model quality?
• RQ2 Does the combination of Mamba and Transformer layers outperform pure Mamba-2 or Transformer-only models in a hybrid setup at scale?
• RQ3 Can ExpertsInt8 quantization enable full 8-bit quantization of the MoE layers without performance degradation, allowing deployment on 8x80GB GPUs?
• RQ4 How does the model perform on long-context tasks, particularly at 256K token context length, compared to other open-weight models?
• RQ5 To what extent do Jamba-1.5 models generalize to multilingual tasks, especially in low-resource languages?

Key findings

• Jamba-1.5-Large achieves 94B active parameters and fits on a single machine with 8x80GB GPUs at 256K context length using ExpertsInt8 quantization, with no quality loss.
• The model reduces KV cache memory usage to 9GB at 256K tokens, representing a 10x reduction compared to LLaMA-3.1-70B and Mistral-Large-2.
• On the RULER benchmark, Jamba-1.5 models are the only open-weight models with an effective context length of 256K tokens.
• Jamba-1.5-Mini achieves 64.3% on multilingual MMLU, outperforming LLaMA-3.1-8B and Gemma-9B, and Jamba-1.5-Large reaches 73.94% on the same benchmark.
• Jamba-1.5-Large outperforms LLaMA-3.1-70B and Mistral-Large-2 in long-context evaluations, demonstrating superior throughput and latency.
• ExpertsInt8 quantization is shown to be effective for MoE layers, preserving model quality while enabling efficient deployment on consumer-grade hardware.

(b) 1.3B.