Source-linked AI summary
Jamba: A Hybrid Transformer-Mamba Language Model
Opher Lieber, Barak Lenz, Hofit Bata, Gal Cohen, Jhonathan Osin, Itay Dalmedigos, Erez Safahi, Shaked Meirom, Yonatan Belinkov, Shai Shalev-Shwartz, Omri Abend, Raz Alon, Tomer Asida, Amir Bergman, Roman Glozman, Michael Gokhman, Avashalom Manevich, Nir Ratner, Noam Rozen, Erez Shwartz, Mor Zusman, Yoav Shoham
TL;DR
Jamba addresses the memory, compute, and long-context limitations of attention-only language models while retaining their performance advantages. It combines Transformer and Mamba layers with MoE, and the resulting model performs comparably to Mixtral-8x7B and Llama-2 70B, supports 256K-token contexts, and offers substantially higher long-context throughput. The release is a pretrained base model intended for further adaptation and study.
Problem
Transformer models have high memory and compute requirements for long contexts, while Mamba models can lag similarly sized Transformers in performance.
Method
Jamba interleaves Transformer and Mamba layers and applies MoE to selected MLP layers, increasing capacity while keeping active-parameter compute manageable.
Results
Jamba performs comparably to Mixtral-8x7B and Llama-2 70B, supports 256K-token contexts, and reaches up to 3x Mixtral throughput for long contexts.
Takeaways & Limitations
The hybrid architecture provides a configurable way to trade off memory usage, throughput, and quality while enabling long-context language modeling.
Takeaways & Limitations
The released Jamba checkpoint is a pretrained base model without alignment, instruction tuning, or moderation mechanisms and requires additional adaptation before production or end-user use.
Abstract
from arXiv · showhide
We present Jamba, a new base large language model based on a novel hybrid Transformer-Mamba mixture-of-experts (MoE) architecture. Specifically, Jamba interleaves blocks of Transformer and Mamba layers, enjoying the benefits of both model families. MoE is added in some of these layers to increase model capacity while keeping active parameter usage manageable. This flexible architecture allows resource- and objective-specific configurations. In the particular configuration we have implemented, we end up with a powerful model that fits in a single 80GB GPU. Built at large scale, Jamba provides high throughput and small memory footprint compared to vanilla Transformers, and at the same time state-of-the-art performance on standard language model benchmarks and long-context evaluations. Remarkably, the model presents strong results for up to 256K tokens context length. We study various architectural decisions, such as how to combine Transformer and Mamba layers, and how to mix experts, and show that some of them are crucial in large scale modeling. We also describe several interesting properties of these architectures which the training and evaluation of Jamba have revealed, and plan to release checkpoints from various ablation runs, to encourage further exploration of this novel architecture. We make the weights of our implementation of Jamba publicly available under a permissive license.
1 Introduction
Jamba combines Transformer and Mamba layers with sparse MoE to balance language-model quality, throughput, memory use, and long-context support. Its released configuration fits on one 80GB GPU, performs comparably to strong larger models, and supports contexts up to 256K tokens.
- Architecture: Jamba combines Transformer and Mamba layers with MoE modules, forming a hybrid architecture designed to balance performance, throughput, and memory footprint.The architecture supports alternative configurations for different hardware and performance requirements.
- Motivation: Transformer models face KV-cache and throughput limitations on long contexts, while Mamba improves efficiency but can lag similarly sized Transformers.Jamba is motivated by combining the complementary strengths of both model families.
- Architecture: MoE increases total model capacity without proportionally increasing compute, because only selected experts contribute active parameters at each token.The implementation uses 16 experts and selects the top 2 experts per token in every other layer.
- Results: Jamba performs comparably to Mixtral-8x7B and Llama-2 70B across a wide range of benchmarks, while supporting 256K-token contexts and outperforming Mixtral on most long-context evaluations.The model provides 3x Mixtral throughput for long contexts and fits on one GPU with 8-bit weights beyond 128K tokens.
- Release: The authors release Jamba under Apache 2.0 and plan to publish checkpoints from ablations covering architectural trade-offs and design choices.The ablations reached up to 7B parameters and 250B training tokens.
- Scope: The released model is a pretrained base model without alignment, instruction tuning, or moderation mechanisms, so additional adaptation is required before production or end-user use.This notice defines an explicit deployment boundary for the released checkpoint.
2 Model Architecture
Jamba combines Transformer and Mamba layers with MoE modules in a flexible hybrid decoder architecture. These choices target the competing goals of quality, throughput, and memory efficiency, especially for long contexts.
- Design objectives: Jamba targets low memory usage, high throughput, and high quality by trading some attention layers for Mamba layers and controlling active MoE parameters.The architecture is designed around both active-parameter usage and KV-cache size rather than total parameter count alone.
- Core architecture: Jamba blocks combine Transformer or Mamba layers with MLPs, while some MLPs are replaced by MoE layers.The architecture mixes attention-to-Mamba layers at a configurable ratio and routes each token to top experts.
- Mixture of experts: MoE increases total model capacity while keeping active parameters and compute comparatively small.The router selects K experts from n available experts, and MoE can be applied every e layers.
- Architectural trade-offs: Increasing the Mamba-to-attention ratio reduces KV-cache memory and improves throughput on long sequences, but may lower model capabilities.The trade-off arises because Mamba layers avoid the attention cache and are more compute-efficient for long contexts.
3 Reaping the Benefits
The released configuration was selected to fit on one 80GB GPU while balancing quality and throughput. It supports long contexts and delivers substantial context-length and throughput advantages over comparable open models.
- 3.1 Jamba Implementation for a Single 80GB GPU: Jamba uses four blocks, each with eight layers and a 1:7 attention-to-Mamba ratio.The configuration also applies MoE every other layer, with 16 experts and two selected per token.
- 3.1 Jamba Implementation for a Single 80GB GPU: The 1:7 ratio was chosen as the most compute-efficient variant among the best-performing quality variants.The selection was based on preliminary ablation experiments.
- 3.1 Jamba Implementation for a Single 80GB GPU: Jamba’s expert configuration was balanced to fit on one 80GB GPU while retaining input memory and limiting compute and communication requirements.The design averages approximately eight experts per layer and uses two experts per token.
- 3.1 Jamba Implementation for a Single 80GB GPU: 2x the context length of Mixtral and 7x that of Llama-2-70B fit on a single A100 80GB GPU with Jamba.The comparison concerns maximum context length under the same single-GPU constraint.
- 3.1 Jamba Implementation for a Single 80GB GPU: Jamba was trained to 1M tokens, while the released model supports context lengths up to 256K tokens.The released context limit is lower than the maximum training context length.
- 3.2 Throughput Analysis: 3x the throughput of Mixtral is reported for Jamba with 128K-token contexts.In the separate batch-size experiment, Jamba also achieved a 3x throughput increase over Mixtral on one A100 GPU at 8K context.
4 Training Infrastructure and Dataset
Jamba was trained with distributed infrastructure on NVIDIA H100 GPUs and an in-house dataset combining web, book, and code data. The dataset pipeline applied quality filtering and deduplication.
- Training infrastructure: Training used NVIDIA H100 GPUs and an in-house framework supporting FSDP, tensor parallelism, sequence parallelism, and expert parallelism.These components enabled efficient large-scale training.
- Dataset: The in-house dataset contains web, book, and code text, with a last update in March 2024.Data processing included quality filters and deduplication.
5 Evaluation
Jamba performs comparably to leading publicly available models of similar or larger size while offering stronger throughput and long-context capabilities. Evaluations show competitive academic benchmark performance, favorable few-shot classification results, and advantages on long-context question answering.
- Academic benchmarks: In most tasks, Jamba performs comparably to leading publicly available models of similar or larger size, including Llama-2 70B and Mixtral.Jamba uses 12B active parameters and 52B total available parameters, while Mixtral uses 12.9B active parameters and Llama-2 70B has 70B total parameters.
- Academic benchmarks: Jamba achieves comparable performance with up to 3x better throughput than Llama-2 70B and Mixtral.The paper reports this comparison across the academic benchmark evaluation.
- Long-context evaluation: Jamba supports context lengths up to 256K tokens and has strong needle-in-a-haystack recall for statements placed in the middle of those contexts.The released model handles up to 256K-token contexts, while training experiments reached up to 1M tokens.
- Long-context evaluation: On few-shot classification with many examples, Jamba outperforms Mixtral on Trec-Fine and Banking77, while the models are on par on NLU Intent and CLINC150.The results use exact match with greedy decoding and extend examples up to 128K tokens.
- Long-context evaluation: On long-context question answering, Jamba outperforms Mixtral on most datasets and on average using F1.The evaluation repurposes five long-context datasets in a 3-shot format, with average input lengths ranging from 6K to 62K tokens before few-shot expansion.
6 Ablations and Insights
The ablations show that Attention–Mamba hybrids outperform pure Attention and pure Mamba models across multiple scales, while MoE improves the hybrid and RMSNorm stabilizes large-model training. They also identify a potential in-context-learning limitation of pure Mamba that hybrid attention can address.
- Benefits of combining Attention and Mamba: Attention–Mamba hybrids outperform pure Attention and pure Mamba models at 1.3B parameters, with virtually no performance difference between 1:3 and 1:7 layer ratios.The models were trained for 250B tokens; the 1:7 ratio was selected for larger experiments because it is more compute-efficient.
- Benefits of combining Attention and Mamba: At 7B parameters, the hybrid outperforms pure Attention and pure Mamba while obtaining better throughput than the vanilla Transformer.The comparison used models trained on 50B tokens.
- Why does the Combination Work?: Pure Mamba performs substantially worse than pure Attention on IMDB, QuAC, and NarrativeQA, whereas the hybrid performs similarly to Attention.The authors associate these failures with difficulty following answer formats and a potential limitation in emergent in-context learning.
- Why does the Combination Work?: Attention in the hybrid focuses on labels from few-shot examples, and the authors found 12 such induction heads across all three attention layers.Figure 8 visualizes an example head in a 1.3B hybrid model on an IMDB example where pure Mamba failed and the hybrid succeeded.
- Mixture-of-experts: MoE improves the Attention–Mamba hybrid at large scale while increasing available capacity through experts and controlling compute via the number of active experts.The reported configuration uses 16 total experts, 2 experts per token, and MoE every 2 layers.
- Training stability: RMSNorm added to internal Mamba activations prevents loss spikes when scaling to the largest released model.The spikes were traced to large activation values inside Mamba layers.
7 Conclusion
Jamba combines Attention, Mamba, and MoE modules in an open implementation designed to balance performance, memory, and throughput. The released model supports 256K-token contexts and fits in a single 80GB GPU even when processing 140K-token texts.
- Conclusion: Jamba combines Attention and Mamba layers with MoE modules, providing flexibility to balance performance and memory requirements while maintaining high throughput.The conclusion describes the implementation as open and reports state-of-the-art performance and long-context support.
- Conclusion: 256K-token context support and single-80GB-GPU deployment are achieved by the largest released model, which has 12B active and 52B total parameters.The model fits in one 80GB GPU even when processing 140K-token texts.