llm-architectures-explained

/models

Jamba v0.1

AI21 Labs · Jamba · open weights

Facts and where they come from

Released2024-03config.jsonconfig.jsonHugging Face repository creation date (api.createdAt)
Licenceapache-2.0config.jsonconfig.jsonREADME metadata: license
Total parameters52Blabmodel cardREADME: 12B active parameters and a total of 52B parameters
Active parameters12Blabmodel cardREADME: 12B active
Context length256K tokensconfig.jsonconfig.jsonmax_position_embeddings
Norm placementprecodemodelling codetransformers 5.18.0 / repo modelling code for jamba: input_layernorm before attention, post_attention_layernorm before the MLP
Norm typeRMSNormcodemodelling codetransformers 5.18.0 / repo modelling code for jamba: input_layernorm before attention, post_attention_layernorm before the MLP
QK-normnocodemodelling codeno q/k normalisation in the attention block
Positional encodingnone; attention layers have no positional encoding; the Mamba layers carry orderpaperarXiv 2403.19887Jamba paper (arXiv 2403.19887) §2: no explicit positional information
Parallel attention and MLPnocodemodelling codetransformers 5.18.0 / repo modelling code for jamba: input_layernorm before attention, post_attention_layernorm before the MLP

Architecture, drawn from the data

28× Mamba + 4× GQA 32q/8kv. Each column is one layer: its token mixer above, its feed-forward block below. Paler columns reuse another layer’s keys and values.

Jamba v0.1: layer stack and blockslayers (32)mixer / FFNlayer 0: Mamba: inner 8192 · state 16layer 0: Gated MLP: 14336layer 1: Mamba: inner 8192 · state 16layer 1: MoE: 16 experts, 2 active · expert 14336layer 2: Mamba: inner 8192 · state 16layer 2: Gated MLP: 14336layer 3: Mamba: inner 8192 · state 16layer 3: MoE: 16 experts, 2 active · expert 14336layer 4: GQA: 32 query / 8 KV heads · head 128layer 4: Gated MLP: 14336layer 5: Mamba: inner 8192 · state 16layer 5: MoE: 16 experts, 2 active · expert 14336layer 6: Mamba: inner 8192 · state 16layer 6: Gated MLP: 14336layer 7: Mamba: inner 8192 · state 16layer 7: MoE: 16 experts, 2 active · expert 14336layer 8: Mamba: inner 8192 · state 16layer 8: Gated MLP: 14336layer 9: Mamba: inner 8192 · state 16layer 9: MoE: 16 experts, 2 active · expert 14336layer 10: Mamba: inner 8192 · state 16layer 10: Gated MLP: 14336layer 11: Mamba: inner 8192 · state 16layer 11: MoE: 16 experts, 2 active · expert 14336layer 12: GQA: 32 query / 8 KV heads · head 128layer 12: Gated MLP: 14336layer 13: Mamba: inner 8192 · state 16layer 13: MoE: 16 experts, 2 active · expert 14336layer 14: Mamba: inner 8192 · state 16layer 14: Gated MLP: 14336layer 15: Mamba: inner 8192 · state 16layer 15: MoE: 16 experts, 2 active · expert 14336layer 16: Mamba: inner 8192 · state 16layer 16: Gated MLP: 14336layer 17: Mamba: inner 8192 · state 16layer 17: MoE: 16 experts, 2 active · expert 14336layer 18: Mamba: inner 8192 · state 16layer 18: Gated MLP: 14336layer 19: Mamba: inner 8192 · state 16layer 19: MoE: 16 experts, 2 active · expert 14336layer 20: GQA: 32 query / 8 KV heads · head 128layer 20: Gated MLP: 14336layer 21: Mamba: inner 8192 · state 16layer 21: MoE: 16 experts, 2 active · expert 14336layer 22: Mamba: inner 8192 · state 16layer 22: Gated MLP: 14336layer 23: Mamba: inner 8192 · state 16layer 23: MoE: 16 experts, 2 active · expert 14336layer 24: Mamba: inner 8192 · state 16layer 24: Gated MLP: 14336layer 25: Mamba: inner 8192 · state 16layer 25: MoE: 16 experts, 2 active · expert 14336layer 26: Mamba: inner 8192 · state 16layer 26: Gated MLP: 14336layer 27: Mamba: inner 8192 · state 16layer 27: MoE: 16 experts, 2 active · expert 14336layer 28: GQA: 32 query / 8 KV heads · head 128layer 28: Gated MLP: 14336layer 29: Mamba: inner 8192 · state 16layer 29: MoE: 16 experts, 2 active · expert 14336layer 30: Mamba: inner 8192 · state 16layer 30: Gated MLP: 14336layer 31: Mamba: inner 8192 · state 16layer 31: MoE: 16 experts, 2 active · expert 1433601631× 16normMamba: inner 8192 · state 16+normMoE: 16 experts, 2 active · expert 14336+× 12normMamba: inner 8192 · state 16+normGated MLP: 14336+× 4normGQA: 32 query / 8 KV heads · head 128+normGated MLP: 14336+Mambafull attentiondense FFNMoE FFN

Modelled costs

From the cost model, batch size 1. Totals the lab states are in the table above; differences come from rounding, from what a lab counts, or from parts the model does not describe (listed on the about page).

Parameters (modelled)51.6B
Active per token (modelled)12.1B
Without embeddings and output head51B total, 11.6B active
Published weights (Hugging Face count)51.6B
KV cache per token, BF16 (layers that grow with context)16 KiB
KV cache + state at 256K tokens, BF164.01 GiB
Decode FLOPs per token at 4K context24 GFLOP
Prefill FLOPs for a 4K prompt95.4 TFLOP

KV cache against context

Jamba v0.1: KV cache bytes against context length101001,00010,000100,0009.5 MiB95 MiB950 MiBcontext (tokens)KV cache + state (BF16)Jamba v0.1

Compare with other models →

Every architecture field

FieldValueSource
d_model4,096config.jsonconfig.jsonhidden_size
vocab65,536config.jsonconfig.jsonvocab_size
tied_embeddingsfalseconfig.jsonconfig.jsontie_word_embeddings
mixers.full.typeattncodemodelling codeattention block
mixers.full.heads32config.jsonconfig.jsonnum_attention_heads
mixers.full.kv_heads8config.jsonconfig.jsonnum_key_value_heads
mixers.full.head_dim128codemodelling codetransformers 5.18.0: head_dim = hidden_size / num_attention_heads
mixers.mamba.typemamba1codemodelling codeMamba (selective SSM)
mixers.mamba.d_inner8,192codemodelling codemamba_expand x hidden_size
mixers.mamba.state16config.jsonconfig.jsonmamba_d_state
mixers.mamba.conv_kernel4config.jsonconfig.jsonmamba_d_conv
mixers.mamba.dt_rank256codemodelling codemamba_dt_rank: auto = ceil(hidden_size / 16)
ffns.dense.typedensecodemodelling codeMLP block
ffns.dense.d_ff14,336config.jsonconfig.jsonintermediate_size
ffns.dense.gatedtruecodemodelling codeMLP: gated (SwiGLU/GeGLU)
ffns.moe.typemoecodemodelling codeMoE block
ffns.moe.experts16config.jsonconfig.jsonnum_experts
ffns.moe.active2config.jsonconfig.jsonnum_experts_per_tok
ffns.moe.d_expert14,336config.jsonconfig.jsonintermediate_size
ffns.moe.gatedtruecodemodelling codeexperts are gated MLPs
layout1× mamba/dense · 1× mamba/moe · 1× mamba/dense · 1× mamba/moe · 1× full/dense · 1× mamba/moe · 1× mamba/dense · 1× mamba/moe · 1× mamba/dense · 1× mamba/moe · 1× mamba/dense · 1× mamba/moe · 1× full/dense · 1× mamba/moe · 1× mamba/dense · 1× mamba/moe · 1× mamba/dense · 1× mamba/moe · 1× mamba/dense · 1× mamba/moe · 1× full/dense · 1× mamba/moe · 1× mamba/dense · 1× mamba/moe · 1× mamba/dense · 1× mamba/moe · 1× mamba/dense · 1× mamba/moe · 1× full/dense · 1× mamba/moe · 1× mamba/dense · 1× mamba/moeconfig.jsonconfig.jsonattn_layer_period/offset, expert_layer_period/offset

Sources