AI21 releases Jamba

AI21 has released Jamba, the first production-grade model based on Mamba. It combines the structured state space model (SSM) architecture of Mamba with Transformer to compensate for the limitations of pure SSM, and supports a 256K context window.

The core design is as follows.

• Of the total 52B parameters, only 12B are activated during inference.
• Out of every 8 layers, only 1 is a Transformer layer, with the rest composed of Mamba and MoE (mixture-of-experts).
• It is optimized to handle up to 140K context on a single 80GB GPU.

In terms of performance, it demonstrated 3x throughput compared to Mixtral 8x7B on long-context benchmarks. The model was released as Apache 2.0 open weights on Hugging Face, and Jamba-Instruct is available on the AI21 Platform, as well as being provided as an NVIDIA NIM inference microservice in the NVIDIA API catalog.