transformer_lens.model_bridge.supported_architectures.jamba module

Jamba hybrid attention+Mamba architecture adapter.

Supports JambaForCausalLM (e.g. ai21labs/Jamba-tiny-random, ai21labs/AI21-Jamba-Reasoning-3B, ai21labs/Jamba-v0.1).

Architecture overview: - Heterogeneous layers from config.layers_block_type — each element is

either "mamba" (JambaMambaDecoderLayer) or "attention" (JambaAttentionDecoderLayer). Attention layers recur every attn_layer_period starting at attn_layer_offset (classically 1/8).

  • Every layer has the same residual skeleton: pre-norm (input_layernorm) → mixer (.mamba or .self_attn) → residual → pre-FFN norm (pre_ff_layernorm) → .feed_forward (dense SwiGLU JambaMLP or JambaSparseMoeBlock when layers_num_experts[i] > 1) → residual.

  • The Mamba mixer is Mamba-1 (JambaMambaMixer): in_proj / conv1d / x_proj / dt_proj / out_proj, plus Jamba-specific dt_layernorm / b_layernorm / c_layernorm on the selective params.

  • Attention is GQA without RoPE — absolute-position-free; no model-level rotary module.

  • Generation threads a unified DynamicCache via past_key_values (attention KV + Mamba conv/recurrent states). The Mamba mixer receives that same object as cache_params=past_key_values.

Key adapter decisions: - BlockBridge (not SSMBlockBridge): two norms and a real post-mixer

residual make transformer-shaped hooks (hook_resid_mid via ln2) meaningful — same choice as Falcon-H1 / GraniteMoeHybrid.

  • SSMMixerBridge (PR #1481 Mamba-1 interp surface) wraps .mamba under the canonical .mixer dict key so find_ssm_mixer / compute_effective_attention / eager_scan resolve it. Inner projections match Mamba-1; the three selective-param RMSNorms are mapped so reconstruction and the opt-in eager_scan path apply them (Jamba fork of stock Mamba-1). Default forward still delegates to HF (bit-identical).

  • Attention and mixer are optional=True so setup skips the absent branch per layer type.

  • FFN: dense GatedMLPBridge when num_experts <= 1 (Reasoning-3B); MoEBridge passthrough with optional dense projections and router when num_experts > 1 (tiny-random / 52B), since dense and MoE layers share the HF path .feed_forward.

  • is_stateful = False: generation uses the standard past_key_values path (same rationale as Zamba2 / Falcon-H1). Setting is_stateful would route through the pure-Mamba cache_params loop and diverge.

class transformer_lens.model_bridge.supported_architectures.jamba.JambaArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for JambaForCausalLM.

Interleaved attention + Mamba-1 layers with optional sparse MoE FFN. Attention and Mamba streams are separate optional slots so each can be ablated independently.

applicable_phases: list[int] = [1, 2, 3, 4]
component_mapping: ComponentMapping | None
uses_split_attention: bool
weight_processing_conversions: Dict[str, ParamProcessingConversion | str] | None