Jamba2 3B is an ultra-compact open source model designed to bring enterprise-grade reliability to on-device deployments. At just 3B parameters, it runs efficiently on consumer devices—iPhones, Androids, Macs, and PCs—while maintaining the grounding and instruction-following capabilities required for production use.
Released under Apache 2.0 License with a 256K context window, Jamba2 3B enables developers to build reliable AI applications for edge environments. For more details, read the full release blog post.
Key Advantages
On-device deployment: Runs efficiently on iPhones, Androids, Macs, and PCs
Ultra-compact footprint: 3B parameters enabling edge deployments with minimal resources
Benchmark leadership: Excels on IFBench, IFEval, Collie, and FACTS
256K context window: Processes long documents and knowledge bases
Apache 2.0 License: Fully open source for commercial use
SSM-Transformer architecture: Memory-efficient design for resource-constrained environments
Evaluation Results
Jamba2 3B achieves category-leading performance on instruction following and grounding benchmarks despite its compact size. The model delivers consistent, context-faithful outputs across diverse enterprise tasks including RAG workflows and technical document processing.
Training and Evaluation Details
Jamba2 models were developed using a comprehensive post-training pipeline starting from Jamba 1.5 pre-training. The models underwent mid-training on 500B carefully curated tokens with increased representation of math, code, high-quality web data, and long documents. A state passing phase optimized the Mamba layers for effective context length generalization. Training continued with cold start supervised fine-tuning to establish instruction-following and reasoning capabilities, followed by DPO optimization.
The final training stages involved multiple on-policy reinforcement learning phases, progressively moving from short-context verifiable rewards to longer context training with mixed verifiable and model-based rewards. Evaluation focused on two key enterprise reliability signals: instruction-following benchmarks measuring steerability, and grounding benchmarks testing context faithfulness. Human evaluators assessed performance on real-world enterprise tasks using blind, counterbalanced side-by-side comparisons, rating outputs on factuality, style, constraint-adherence, instruction-following, and helpfulness.
Quickstart
Run with vLLM
Best results require vLLM version 0.10.2 or higher.
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
model = AutoModelForCausalLM.from_pretrained("ai21labs/AI21-Jamba2-3B",
dtype=torch.bfloat16,
attn_implementation="flash_attention_2",
device_map="auto")
tokenizer = AutoTokenizer.from_pretrained("ai21labs/AI21-Jamba-3B")
messages = [
{"role": "system",
"content": "You are an HR Policy Assistant.
Answer employee questions using only the provided policy documents.
If the answer isn't in the documents, say so clearly.
Be concise and cite the specific policy section when possible."
},
{"role": "user",
"content": "Context documents: {retrieved_chunks}.
Employee question: {user_question}.
Answer:"
},
]
prompts = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
outputs = model.generate(**tokenizer(prompts, return_tensors="pt").to(model.device), do_sample=True, temperature=0.6)
generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
print(generated_text)
