liveBook · Manning

JAMBA

As explained in LLMs in Production

JAMBA is a hybrid language-model architecture combining state-space and transformer layers with joint attention.

JAMBA targets the long-context and memory limits of attention-only language models by combining MAMBA-style state-space layers, transformer layers, and joint attention. Its hybrid design aims to retain transformer strengths while improving context-length efficiency. section 12.2.2

A practical JAMBA deployment involves several important trade-offs:

  • A 52-billion-parameter mixture-of-experts model can support 140K-token contexts on an 80 GB GPU, outperforming attention-only memory requirements. section 12.2.2
  • Fine-tuning exceeds one 80 GB GPU even in half precision, so Accelerate must parallelize training across several GPUs. section 12.2.2
Listing 12.1 Finetuning and inferencing JAMBA 
from trl import SFTTrainer
from peft import LoraConfig
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    TrainingArguments,
)
from transformers import BitsAndBytesConfig
import torch
from datasets import load_dataset

tokenizer = AutoTokenizer.from_pretrained("ai21labs/Jamba-v0.1")
model = AutoModelForCausalLM.from_pretrained(
    "ai21labs/Jamba-v0.1", device_map="auto"
)

dataset = load_dataset("Abirate/english_quotes", split="train")

The listing provides a minimal starting point for loading ai21labs/Jamba-v0.1 with Transformers, selecting its tokenizer, model, and training dataset. section 12.2.2

How LLMs in Production covers it

How deep it goes
Explains the hybrid architecture and operational workflow, but defers detailed state-space and joint-attention internals.
Examples
Uses quote fine-tuning, model saving, 8-bit loading, and generation prompted with a Super Bowl sentence.
What it assumes
Assumes familiarity with transformer workflows, tokenizers, datasets, GPUs, and parameter-efficient fine-tuning.
What you can do after
Apply it via Hugging Face Transformers, TRL, PEFT, and Accelerate for fine-tuning or long-context inference.
How it connects
Connects state-space models with transformers, joint attention, mixture-of-experts models, LoRA, quantization, and FlashAttention.
Cover of LLMs in Production
liveBookmanning.com
Read LLMs in Production on liveBook Get the book at manning.com

This explainer was generated by Manning's knowledge graph from the text of LLMs in Production. The section links point to the parts of the book it draws on.