From 73d1576db57ac9508eecbdc1efde187bfeff0924 Mon Sep 17 00:00:00 2001 From: Mukund Pandey Date: Mon, 6 Jul 2026 15:37:46 +0100 Subject: [PATCH 1/2] feat: add Zamba2ForCausalLM architecture adapter Closes #1463 Supports Zyphra/Zamba2-1.2B, Zamba2-7B, Zamba2-7B-Instruct via TransformerBridge. Zamba2 is a hybrid Mamba-2 / shared-attention model where layers_block_type alternates between "mamba" (pure SSM) and "hybrid" (SSM + shared global-attention block, weight-tied across num_mem_blocks slots). Key design decisions: - SSMBlockBridge for all layers: hook_in/hook_out fire regardless of layer type. Mamba layers expose mixer submodule hooks (in_proj, conv1d, inner_norm, out_proj); hybrid layers have norm and mixer marked optional=True so component_setup skips them gracefully. - is_stateful=False: Zamba2 threads its unified cache via past_key_values (standard KV path), not the Mamba cache_params path. Setting is_stateful=True collided with Zamba2's own cache call. - positional_embedding_type="none": RoPE is handled inside the HF attention block; no model-level rotary bridge needed. - Added mamba_expand, mamba_ngroups, num_mem_blocks, layers_block_type, use_shared_attention_adapter to _HF_PASSTHROUGH_ATTRS in both sources/transformers.py and sources/_bridge_builder.py so these config fields reach the adapter reliably. 21 integration tests pass against Zyphra/Zamba2-1.2B (38 layers: 32 mamba + 6 hybrid). Forward parity is exact (max_diff == 0.0); greedy generation matches HF bit-for-bit across both layer types. --- .../model_bridge/test_zamba2_adapter.py | 302 ++++++++++++++++++ .../factories/architecture_adapter_factory.py | 2 + .../model_bridge/sources/_bridge_builder.py | 6 + .../model_bridge/sources/transformers.py | 6 + .../supported_architectures/__init__.py | 16 +- .../supported_architectures/zamba2.py | 166 ++++++++++ 6 files changed, 492 insertions(+), 6 deletions(-) create mode 100644 tests/integration/model_bridge/test_zamba2_adapter.py create mode 100644 transformer_lens/model_bridge/supported_architectures/zamba2.py diff --git a/tests/integration/model_bridge/test_zamba2_adapter.py b/tests/integration/model_bridge/test_zamba2_adapter.py new file mode 100644 index 000000000..b2ed82451 --- /dev/null +++ b/tests/integration/model_bridge/test_zamba2_adapter.py @@ -0,0 +1,302 @@ +"""Integration tests for the Zamba2 architecture adapter. + +Verifies forward-pass and generation parity against Zyphra/Zamba2-1.2B: +- Forward-pass logits match HF exactly (bridge delegates the full forward to HF) +- Greedy multi-token generation matches HF bit-for-bit (exercises the unified + Zamba2HybridDynamicCache threaded via past_key_values across Mamba-2 and + shared global-attention layers) +- Sanity checks: config flags, block count, hook coverage on both layer types + +Zamba2 has two layer types in ``config.layers_block_type``: + - ``"mamba"`` — pure Mamba-2 SSM (Zamba2MambaDecoderLayer) + - ``"hybrid"`` — Mamba-2 + shared global-attention (Zamba2HybridLayer) + +Run with GPU acceleration: + CUDA_VISIBLE_DEVICES=0 pytest tests/integration/model_bridge/test_zamba2_adapter.py -v -s + +On a CPU-only machine: + pytest tests/integration/model_bridge/test_zamba2_adapter.py -v -s +""" + +import gc + +import pytest +import torch + +from transformer_lens.model_bridge.bridge import TransformerBridge +from transformer_lens.model_bridge.generalized_components import ( + RMSNormalizationBridge, + SSM2MixerBridge, + SSMBlockBridge, +) + +pytestmark = pytest.mark.slow + +MODEL = "Zyphra/Zamba2-1.2B" + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _device() -> str: + return "cuda" if torch.cuda.is_available() else "cpu" + + +def _dtype() -> torch.dtype: + # bfloat16 on GPU to match HF defaults; float32 on CPU for numerical safety + return torch.bfloat16 if torch.cuda.is_available() else torch.float32 + + +# --------------------------------------------------------------------------- +# Session fixture — load once, share across all test classes +# --------------------------------------------------------------------------- + + +@pytest.fixture(scope="module") +def zamba2_bridge(): + device = _device() + dtype = _dtype() + bridge = TransformerBridge.boot_transformers(MODEL, device=device, dtype=dtype) + yield bridge + del bridge + if torch.cuda.is_available(): + torch.cuda.empty_cache() + for _ in range(3): + gc.collect() + + +# --------------------------------------------------------------------------- +# Config and bridge structure +# --------------------------------------------------------------------------- + + +class TestZamba2BridgeCreation: + """Smoke-test that the bridge loads with the right config flags.""" + + def test_norm_bridges_are_rms(self, zamba2_bridge: TransformerBridge) -> None: + """normalization_type='RMS' must wire RMSNormalizationBridge, not LayerNorm. + + The block pre-norm (on Mamba layers) and the final norm should both be + RMS. Asserting the concrete bridge type catches a regression where the + adapter wires the wrong normalization component. + """ + lbt = getattr(zamba2_bridge.cfg, "layers_block_type", []) + first_mamba = next(i for i, t in enumerate(lbt) if t == "mamba") + assert isinstance(zamba2_bridge.blocks[first_mamba].norm, RMSNormalizationBridge) + assert isinstance(zamba2_bridge.ln_final, RMSNormalizationBridge) + + def test_block_count(self, zamba2_bridge: TransformerBridge) -> None: + # Zamba2-1.2B has 38 layers (32 "mamba" + 6 "hybrid"); the block count + # must match the per-layer type list rather than a hard-coded magic + # number so the test tracks the checkpoint's actual config. + lbt = getattr(zamba2_bridge.cfg, "layers_block_type", []) + assert len(zamba2_bridge.blocks) == len(lbt) == 38 + + def test_blocks_are_ssm_block_bridge(self, zamba2_bridge: TransformerBridge) -> None: + assert isinstance(zamba2_bridge.blocks[0], SSMBlockBridge) + + def test_mamba_layers_have_mixer(self, zamba2_bridge: TransformerBridge) -> None: + """Mamba (mamba) layers should expose SSM2MixerBridge as .mixer.""" + lbt = getattr(zamba2_bridge.cfg, "layers_block_type", []) + first_mamba = next(i for i, t in enumerate(lbt) if t == "mamba") + assert isinstance(zamba2_bridge.blocks[first_mamba].mixer, SSM2MixerBridge) + + def test_hybrid_layers_have_no_mixer(self, zamba2_bridge: TransformerBridge) -> None: + """Hybrid layers have no top-level .mamba; .mixer should be absent or None.""" + lbt = getattr(zamba2_bridge.cfg, "layers_block_type", []) + first_hybrid = next(i for i, t in enumerate(lbt) if t == "hybrid") + # The optional=True submodule should not be set up by component_setup + mixer = getattr(zamba2_bridge.blocks[first_hybrid], "mixer", None) + assert mixer is None or getattr( + mixer, "optional", False + ), "Hybrid layer should not have a wired .mixer (no top-level .mamba attribute)" + + def test_layers_block_type_populated(self, zamba2_bridge: TransformerBridge) -> None: + lbt = getattr(zamba2_bridge.cfg, "layers_block_type", []) + assert len(lbt) == len(zamba2_bridge.blocks) + # Both layer types must appear + assert "mamba" in lbt, "No mamba (Mamba) layers found" + assert "hybrid" in lbt, "No hybrid (Mamba + attention) layers found" + + def test_mamba_intermediate_size_positive(self, zamba2_bridge: TransformerBridge) -> None: + assert getattr(zamba2_bridge.cfg, "mamba_intermediate_size", 0) > 0 + + def test_conv_dim_positive(self, zamba2_bridge: TransformerBridge) -> None: + assert getattr(zamba2_bridge.cfg, "conv_dim", 0) > 0 + + def test_uses_standard_kv_cache_path(self, zamba2_bridge: TransformerBridge) -> None: + # Zamba2 threads its unified cache via past_key_values (standard KV + # path), NOT the Mamba cache_params path. is_stateful=True would select + # the Mamba path, whose cache_params kwarg collides with Zamba2's own + # forward. Keeping it False routes generation through the past_key_values + # path, which matches HF generate() bit-for-bit (see TestZamba2Generation). + assert zamba2_bridge.cfg.is_stateful is False + + def test_positional_embedding_none(self, zamba2_bridge: TransformerBridge) -> None: + # RoPE is handled inside the HF attention block; no model-level embedding + assert zamba2_bridge.cfg.positional_embedding_type == "none" + + +# --------------------------------------------------------------------------- +# Forward-pass parity +# --------------------------------------------------------------------------- + + +class TestZamba2ForwardPass: + """Bridge logits must match HF logits exactly. + + Zamba2ArchitectureAdapter uses SSMBlockBridge with a pure passthrough + forward, so the bridge never reimplements any computation. Parity with + HF should be exact (diff == 0), not just close. + """ + + @pytest.fixture(scope="class") + def tokens(self) -> torch.Tensor: + return torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8]]) + + def test_forward_returns_logits( + self, zamba2_bridge: TransformerBridge, tokens: torch.Tensor + ) -> None: + tokens = tokens.to(_device()) + with torch.no_grad(): + out = zamba2_bridge(tokens) + assert out.shape == (1, 8, zamba2_bridge.cfg.d_vocab) + assert not torch.isnan(out).any(), "NaN in bridge logits" + assert not torch.isinf(out).any(), "Inf in bridge logits" + + def test_forward_matches_hf_exactly( + self, zamba2_bridge: TransformerBridge, tokens: torch.Tensor + ) -> None: + tokens = tokens.to(_device()) + hf_model = zamba2_bridge.original_model + with torch.no_grad(): + bridge_out = zamba2_bridge(tokens) + hf_out = hf_model(tokens).logits + max_diff = (bridge_out.float() - hf_out.float()).abs().max().item() + assert max_diff == 0.0, ( + f"Bridge vs HF forward max diff = {max_diff:.2e}. " + "Expected 0 because SSMBlockBridge.forward() is a pure passthrough." + ) + + def test_forward_no_nan_on_longer_sequence(self, zamba2_bridge: TransformerBridge) -> None: + # 32 tokens exercises multiple Mamba SSM steps and shared attention passes + tokens = torch.arange(1, 33).unsqueeze(0).to(_device()) + with torch.no_grad(): + out = zamba2_bridge(tokens) + assert not torch.isnan(out).any(), "NaN in logits for 32-token sequence" + + +# --------------------------------------------------------------------------- +# Multi-token generation parity (exercises DynamicCache state handling) +# --------------------------------------------------------------------------- + + +class TestZamba2Generation: + """Bridge greedy generation must match HF native generate() exactly. + + DynamicCache carries both KV-cache entries (from the shared attention + blocks in hybrid layers) and Mamba-2 conv/recurrent states. Token-level + equality with HF confirms the state threading is correct across both + layer types. + """ + + @pytest.fixture(scope="class") + def prompt(self) -> torch.Tensor: + return torch.tensor([[1, 2, 3, 4]]) + + def test_generation_produces_tokens( + self, zamba2_bridge: TransformerBridge, prompt: torch.Tensor + ) -> None: + prompt = prompt.to(_device()) + with torch.no_grad(): + result = zamba2_bridge.generate(prompt, max_new_tokens=5, do_sample=False) + assert isinstance(result, torch.Tensor) + assert result.shape == (1, 9) # 4 prompt + 5 new + + def test_greedy_matches_hf_exactly( + self, zamba2_bridge: TransformerBridge, prompt: torch.Tensor + ) -> None: + """Bit-for-bit equality with HF generate() over 8 new tokens.""" + prompt = prompt.to(_device()) + hf_model = zamba2_bridge.original_model + with torch.no_grad(): + bridge_out = zamba2_bridge.generate(prompt, max_new_tokens=8, do_sample=False) + hf_out = hf_model.generate(prompt, max_new_tokens=8, do_sample=False, pad_token_id=0) + assert torch.equal(bridge_out, hf_out), ( + f"Token mismatch between bridge and HF.\n" + f" bridge : {bridge_out.tolist()}\n" + f" hf : {hf_out.tolist()}\n" + "DynamicCache state threading across Mamba-2 and hybrid attention layers " + "is likely wrong." + ) + + def test_generation_is_deterministic( + self, zamba2_bridge: TransformerBridge, prompt: torch.Tensor + ) -> None: + prompt = prompt.to(_device()) + with torch.no_grad(): + out1 = zamba2_bridge.generate(prompt, max_new_tokens=4, do_sample=False) + out2 = zamba2_bridge.generate(prompt, max_new_tokens=4, do_sample=False) + assert torch.equal(out1, out2), "Greedy generation is not deterministic" + + +# --------------------------------------------------------------------------- +# Hook coverage: bridge hooks fire for both Mamba and hybrid layers +# --------------------------------------------------------------------------- + + +class TestZamba2HookCoverage: + """run_with_cache captures residual stream and mixer hooks.""" + + @pytest.fixture(scope="class") + def cache(self, zamba2_bridge: TransformerBridge): + tokens = torch.tensor([[1, 2, 3, 4, 5]]).to(_device()) + with torch.no_grad(): + _, cache = zamba2_bridge.run_with_cache(tokens) + return cache + + def test_block_hooks_fire_on_all_layers(self, cache, zamba2_bridge: TransformerBridge) -> None: + """hook_in and hook_out must fire on every layer regardless of type.""" + n_blocks = len(zamba2_bridge.blocks) + # Sample first, an early, the middle, and the last layer (indices derived + # from the actual block count, not a hard-coded magic number). + for i in sorted({0, 1, n_blocks // 2, n_blocks - 1}): + assert f"blocks.{i}.hook_in" in cache, f"Missing hook_in for block {i}" + assert f"blocks.{i}.hook_out" in cache, f"Missing hook_out for block {i}" + + def test_mamba_mixer_submodule_hooks_fire( + self, cache, zamba2_bridge: TransformerBridge + ) -> None: + """Mamba (mamba) layers must expose in_proj / conv1d / out_proj hooks.""" + lbt = getattr(zamba2_bridge.cfg, "layers_block_type", []) + mamba_indices = [i for i, t in enumerate(lbt) if t == "mamba"] + assert mamba_indices, "No mamba layers found in layers_block_type" + for i in mamba_indices[:3]: + for submod in ("in_proj", "conv1d", "out_proj"): + key_in = f"blocks.{i}.mixer.{submod}.hook_in" + key_out = f"blocks.{i}.mixer.{submod}.hook_out" + assert key_in in cache, f"Missing {key_in}" + assert key_out in cache, f"Missing {key_out}" + + def test_hybrid_layers_no_mixer_hooks(self, cache, zamba2_bridge: TransformerBridge) -> None: + """Hybrid layers have no top-level .mamba, so no mixer submodule hooks.""" + lbt = getattr(zamba2_bridge.cfg, "layers_block_type", []) + hybrid_indices = [i for i, t in enumerate(lbt) if t == "hybrid"] + assert hybrid_indices, "No hybrid layers found in layers_block_type" + for i in hybrid_indices[:3]: + # mixer hooks must not appear for hybrid layers + assert ( + f"blocks.{i}.mixer.in_proj.hook_in" not in cache + ), f"Unexpected mixer hook on hybrid layer {i}" + + def test_no_transformer_specific_hooks(self, cache) -> None: + """SSMBlockBridge must not inject transformer-shaped hook names.""" + forbidden = ("hook_resid_mid", "hook_attn_out", "hook_mlp_out") + bad = [k for k in cache if any(f in k for f in forbidden)] + assert bad == [], f"Unexpected transformer-shaped hooks: {bad[:5]}" + + def test_no_nan_in_cache(self, cache) -> None: + for key, val in cache.items(): + if isinstance(val, torch.Tensor) and val.is_floating_point(): + assert not torch.isnan(val).any(), f"NaN in cache['{key}']" diff --git a/transformer_lens/factories/architecture_adapter_factory.py b/transformer_lens/factories/architecture_adapter_factory.py index 5c73e1cf2..78725edf1 100644 --- a/transformer_lens/factories/architecture_adapter_factory.py +++ b/transformer_lens/factories/architecture_adapter_factory.py @@ -82,6 +82,7 @@ T5ArchitectureAdapter, T5GemmaArchitectureAdapter, XGLMArchitectureAdapter, + Zamba2ArchitectureAdapter, ) # Export supported architectures @@ -163,6 +164,7 @@ "MT5ForConditionalGeneration": T5ArchitectureAdapter, "T5GemmaForConditionalGeneration": T5GemmaArchitectureAdapter, "XGLMForCausalLM": XGLMArchitectureAdapter, + "Zamba2ForCausalLM": Zamba2ArchitectureAdapter, "NanoGPTForCausalLM": NanogptArchitectureAdapter, "TransformerLensNative": NativeArchitectureAdapter, "MinGPTForCausalLM": MingptArchitectureAdapter, diff --git a/transformer_lens/model_bridge/sources/_bridge_builder.py b/transformer_lens/model_bridge/sources/_bridge_builder.py index a74d2ef9d..e94de5141 100644 --- a/transformer_lens/model_bridge/sources/_bridge_builder.py +++ b/transformer_lens/model_bridge/sources/_bridge_builder.py @@ -73,6 +73,12 @@ "cond_dim", "adaln", "cross_attn", + # Zamba2 (Mamba-2 + shared-attention hybrid) + "mamba_expand", + "mamba_ngroups", + "num_mem_blocks", + "layers_block_type", + "use_shared_attention_adapter", ] diff --git a/transformer_lens/model_bridge/sources/transformers.py b/transformer_lens/model_bridge/sources/transformers.py index 7929ee7bd..f22167589 100644 --- a/transformer_lens/model_bridge/sources/transformers.py +++ b/transformer_lens/model_bridge/sources/transformers.py @@ -572,6 +572,12 @@ def boot( "cond_dim", "adaln", "cross_attn", + # Zamba2 (Mamba-2 + shared-attention hybrid) + "mamba_expand", + "mamba_ngroups", + "num_mem_blocks", + "layers_block_type", + "use_shared_attention_adapter", ] for attr in _HF_PASSTHROUGH_ATTRS: val = getattr(hf_config, attr, None) diff --git a/transformer_lens/model_bridge/supported_architectures/__init__.py b/transformer_lens/model_bridge/supported_architectures/__init__.py index ee7bc5b6a..20a3ed6a1 100644 --- a/transformer_lens/model_bridge/supported_architectures/__init__.py +++ b/transformer_lens/model_bridge/supported_architectures/__init__.py @@ -12,12 +12,12 @@ from transformer_lens.model_bridge.supported_architectures.baichuan import ( BaichuanArchitectureAdapter, ) -from transformer_lens.model_bridge.supported_architectures.bd3lm import ( - BD3LMArchitectureAdapter, -) from transformer_lens.model_bridge.supported_architectures.bart import ( BartArchitectureAdapter, ) +from transformer_lens.model_bridge.supported_architectures.bd3lm import ( + BD3LMArchitectureAdapter, +) from transformer_lens.model_bridge.supported_architectures.bert import ( BertArchitectureAdapter, ) @@ -63,6 +63,9 @@ from transformer_lens.model_bridge.supported_architectures.glm4_moe import ( Glm4MoeArchitectureAdapter, ) +from transformer_lens.model_bridge.supported_architectures.glm_moe_dsa import ( + GlmMoeDsaArchitectureAdapter, +) from transformer_lens.model_bridge.supported_architectures.gpt2 import ( GPT2ArchitectureAdapter, ) @@ -78,9 +81,6 @@ from transformer_lens.model_bridge.supported_architectures.gptj import ( GptjArchitectureAdapter, ) -from transformer_lens.model_bridge.supported_architectures.glm_moe_dsa import ( - GlmMoeDsaArchitectureAdapter, -) from transformer_lens.model_bridge.supported_architectures.granite import ( GraniteArchitectureAdapter, ) @@ -219,6 +219,9 @@ from transformer_lens.model_bridge.supported_architectures.xglm import ( XGLMArchitectureAdapter, ) +from transformer_lens.model_bridge.supported_architectures.zamba2 import ( + Zamba2ArchitectureAdapter, +) __all__ = [ "ApertusArchitectureAdapter", @@ -294,4 +297,5 @@ "T5ArchitectureAdapter", "T5GemmaArchitectureAdapter", "XGLMArchitectureAdapter", + "Zamba2ArchitectureAdapter", ] diff --git a/transformer_lens/model_bridge/supported_architectures/zamba2.py b/transformer_lens/model_bridge/supported_architectures/zamba2.py new file mode 100644 index 000000000..3d7f61f50 --- /dev/null +++ b/transformer_lens/model_bridge/supported_architectures/zamba2.py @@ -0,0 +1,166 @@ +"""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`` — each element + is either ``"mamba"`` (pure Mamba-2 SSM) 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. +""" + +from typing import Any + +from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter +from transformer_lens.model_bridge.generalized_components import ( + DepthwiseConv1DBridge, + EmbeddingBridge, + GatedRMSNormBridge, + LinearBridge, + RMSNormalizationBridge, + SSM2MixerBridge, + SSMBlockBridge, + UnembeddingBridge, +) +from transformer_lens.model_bridge.generalized_components.base import ( + GeneralizedComponent, +) + + +def _make_optional(component: "GeneralizedComponent") -> "GeneralizedComponent": + """Mark a GeneralizedComponent submodule as optional. + + ``component_setup.py`` reads ``getattr(submodule, 'optional', False)`` at + setup time, so setting the attribute directly is safe regardless of whether + the bridge class's ``__init__`` accepts an ``optional`` keyword argument. + """ + component.optional = True + return component + + +class Zamba2ArchitectureAdapter(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. + """ + + # P1: exact passthrough vs raw HF; P2/P3: skip without HookedTransformer; + # P4: generation with past_key_values cache threading. + applicable_phases: list[int] = [1, 2, 3, 4] + + def __init__(self, cfg: Any) -> None: + super().__init__(cfg) + + self.cfg.normalization_type = "RMS" + self.cfg.uses_rms_norm = True + # RoPE is handled internally by the attention block (position_ids are + # threaded through HF layers); no model-level rotary bridge needed. + self.cfg.positional_embedding_type = "none" + # MLP inside the shared attention block uses GELU, not SwiGLU. + self.cfg.gated_mlp = False + self.cfg.attn_only = False + self.cfg.final_rms = True + # NOTE: is_stateful stays False even though Mamba-2 layers carry SSM + # state. In this bridge, ``is_stateful=True`` selects the *Mamba* cache + # path, which threads the cache as ``cache_params=`` and drives a + # conv-kernel ``cache_position``. Zamba2's HF forward instead threads a + # single unified ``Zamba2HybridDynamicCache`` via ``past_key_values=`` + # (both KV entries and SSM conv/recurrent states live in that one + # object). The standard KV-cache generation path already threads + # ``past_key_values`` correctly and matches HF ``generate`` bit-for-bit, + # so we use it rather than the Mamba path (whose ``cache_params`` kwarg + # collides with Zamba2's own ``cache_params=past_key_values`` call). + self.cfg.is_stateful = False + + # Expose the per-layer type list so analysis tools can identify which + # layers are Mamba-only vs hybrid. + # Values: "mamba" (Mamba-2) | "hybrid" (Mamba-2 + attention). + layers_block_type = list(getattr(cfg, "layers_block_type", [])) + setattr(self.cfg, "layers_block_type", layers_block_type) + + # Number of unique shared attention weight blocks (hybrid layers cycle + # through num_mem_blocks independent attention weight sets). + setattr(self.cfg, "num_mem_blocks", getattr(cfg, "num_mem_blocks", 1)) + + # Whether per-layer LoRA adapters are active on the shared attention. + setattr( + self.cfg, + "use_shared_attention_adapter", + getattr(cfg, "use_shared_attention_adapter", False), + ) + + # Mamba-2 dimensional config (mirrors Mamba2ArchitectureAdapter / + # NemotronHArchitectureAdapter patterns). + mamba_intermediate_size = int(getattr(cfg, "mamba_expand", 2) * self.cfg.d_model) + n_groups = getattr(cfg, "mamba_ngroups", 1) + ssm_state_size = getattr(cfg, "mamba_d_state", 64) + conv_dim = mamba_intermediate_size + 2 * n_groups * ssm_state_size + setattr(self.cfg, "mamba_intermediate_size", mamba_intermediate_size) + setattr(self.cfg, "conv_dim", conv_dim) + + self.weight_processing_conversions = {} + + self.component_mapping = { + "embed": EmbeddingBridge(name="model.embed_tokens"), + "blocks": SSMBlockBridge( + name="model.layers", + submodules={ + # Pre-norm: present on Zamba2MambaDecoderLayer (.input_layernorm), + # absent on Zamba2HybridLayer (no top-level pre-norm) -> optional. + "norm": _make_optional( + RMSNormalizationBridge(name="input_layernorm", config=self.cfg) + ), + # Mamba-2 mixer: present on Zamba2MambaDecoderLayer (.mamba), + # absent at the top level of Zamba2HybridLayer -> optional. + "mixer": _make_optional( + SSM2MixerBridge( + name="mamba", + config=self.cfg, + submodules={ + # -- Mamba-2 inner submodules (all optional) -- + "in_proj": LinearBridge(name="in_proj", optional=True), + "conv1d": DepthwiseConv1DBridge(name="conv1d", optional=True), + # HF names the gated RMS norm "norm" inside the + # mixer; TL uses "inner_norm" to avoid colliding + # with the block-level norm declared above. + "inner_norm": _make_optional(GatedRMSNormBridge(name="norm")), + "out_proj": LinearBridge(name="out_proj", optional=True), + }, + ) + ), + }, + ), + "ln_final": RMSNormalizationBridge(name="model.final_layernorm", config=self.cfg), + "unembed": UnembeddingBridge(name="lm_head"), + } From d62f2f2900195cebf9fe1da72ded381a95014f31 Mon Sep 17 00:00:00 2001 From: Mukund Pandey Date: Mon, 6 Jul 2026 16:45:01 +0100 Subject: [PATCH 2/2] fix: add Zamba2ForCausalLM to model registry (HF_SUPPORTED_ARCHITECTURES + CANONICAL_AUTHORS_BY_ARCH) --- transformer_lens/tools/model_registry/__init__.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/transformer_lens/tools/model_registry/__init__.py b/transformer_lens/tools/model_registry/__init__.py index b01c8fe85..9c4cd03fa 100644 --- a/transformer_lens/tools/model_registry/__init__.py +++ b/transformer_lens/tools/model_registry/__init__.py @@ -116,6 +116,7 @@ "MT5ForConditionalGeneration", "T5GemmaForConditionalGeneration", "XGLMForCausalLM", + "Zamba2ForCausalLM", } # Foundation-trained orgs per architecture. Source of truth for the scraper's @@ -192,6 +193,7 @@ "T5ForConditionalGeneration": ["google-t5", "google", "Salesforce", "MBZUAI"], "T5GemmaForConditionalGeneration": ["google"], "XGLMForCausalLM": ["facebook"], + "Zamba2ForCausalLM": ["Zyphra"], } __all__ = [