Skip to content

[Proposal] Add Zamba2 shared-attention hybrid adapter (Zamba2ForCausalLM) #1463

Description

@jlarson4

Proposal

Add a TransformerBridge adapter for Zamba2ForCausalLM (Zyphra Zamba2), a Mamba2 backbone interleaved with shared global-attention blocks carrying per-layer LoRA adapters.

Motivation

Zamba2's distinguishing feature is that a small number of global-attention blocks are shared across depth (reused at multiple layers, specialized only by per-layer LoRA), running over a Mamba2 backbone. That is a different hybrid than the interleaved Mamba2+attention adapters already in the registry (NemotronH, GraniteMoeHybrid): the shared block lets researchers ask where in-context-learning and retrieval physically live, and what the reused attention contributes on top of the state-space stream. It already has direct interp results on the sub-2B checkpoint (function-vector head identification, AIE scoring, ablation, and steering on Zamba2-1.2B), is native in transformers (≥4.48), and the 1.2B checkpoint runs on modest hardware.

Gap scan (2026-06-25): 9 models, ~1.14M downloads.

Pitch

Implement the Zamba2 block structure: Mamba2 mixer layers plus the shared global-attention block(s) with per-layer LoRA-style adapters (config fields use_shared_attention_adapter / adapter_rank). The single-vs-two shared-block split is encoded in num_mem_blocks (1 for the 1.2B, 2 for the 2.7B/7B). Expose the shared-attention and Mamba2 streams as separate hooks so the two paths can be ablated independently.

  • Claude Code users can scaffold with /add-model-support Zyphra/Zamba2-1.2B-instruct.
  • Register at the four sites listed in contributing.md.
  • Verify smallest-first: Zyphra/Zamba2-1.2B-instruct, then Zyphra/Zamba2-2.7B where hardware allows.

Additional context

Checklist

  • I have checked that there is no similar issue in the repo (required)

Metadata

Metadata

Assignees

Labels

TransformerBridgeBug specific to the new TransformerBridge systemcomplexity-moderateModerately complicated issues for people who have intermediate experience with the codehelp wantedExtra attention is needednew-architectureThis card involves adding a new architecture .

Type

No type

Fields

No fields configured for issues without a type.

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions