transformer_lens.model_bridge.supported_architectures.zamba2 module

Zamba2 hybrid Mamba2-Transformer architecture adapter.

Supports Zamba2ForCausalLM (e.g. Zyphra/Zamba2-1.2B, Zamba2-7B, Zamba2-7B-Instruct).

Architecture overview: - Heterogeneous layers defined by config.layers_block_type — surfaced with

canonical names: "linear_attention" (pure Mamba-2 SSM, HF "mamba") or "hybrid" (Mamba-2 + shared global-attention block).

  • Most layers are Zamba2MambaDecoderLayer: a single pre-norm (input_layernorm) followed by a Mamba-2 mixer (.mamba).

  • A recurring subset are Zamba2HybridLayer: each wraps a Mamba-2 decoder layer plus a SHARED Zamba2AttentionDecoderLayer. The shared attention block’s weights are tied across all hybrid layers, cycling through config.num_mem_blocks unique blocks. When config.use_shared_attention_adapter=True, each hybrid layer carries an independent per-layer LoRA adapter on top of the shared attention.

  • No model-level rotary embedding module is wired by the bridge — the attention block handles RoPE internally via position_ids.

  • Generation runs on the standard KV-cache path: HF threads a single unified Zamba2HybridDynamicCache via past_key_values (carrying both KV-cache entries for attention and SSM conv/recurrent states for Mamba-2), so the bridge does not use the Mamba-specific cache_params stateful path.

Key adapter decisions: - SSMBlockBridge is used for all layers. Its forward delegates entirely to

the HF layer, giving hook_in / hook_out on every layer regardless of type.

  • For Mamba layers: norm (-> .input_layernorm) and mixer (-> .mamba) are declared as submodules and expose inner hooks (in_proj, conv1d, inner_norm, out_proj).

  • For Hybrid layers: norm and mixer are marked optional=True so component_setup skips them gracefully (Hybrid layers have no top-level .input_layernorm or .mamba). Block-level hook_in/hook_out still fire on every layer.

  • applicable_phases = [1, 2, 3, 4]: P1 is exact vs raw HF (pure passthrough); P2/P3 skip without a HookedTransformer; P4 exercises past_key_values cache threading across Mamba-2 and attention layers.

class transformer_lens.model_bridge.supported_architectures.zamba2.Zamba2ArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for Zamba2ForCausalLM.

Hybrid Mamba-2 + shared global-attention model. Most layers are pure Mamba-2 SSM ("mamba"); a recurring subset are hybrid layers ("hybrid") that route through a shared attention block before the Mamba-2 step.

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