feat: add Zamba2ForCausalLM architecture adapter#1486
Merged
Conversation
Closes TransformerLensOrg#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.
mukund1985
force-pushed
the
feat/zamba2-adapter
branch
from
98a7560 to
73d1576
Compare
Contributor
Author
|
@jlarson4 ready for review when you get a chance. One thing worth flagging: the PR touches sources/transformers.py and sources/_bridge_builder.py to add five Mamba-family config attributes to _HF_PASSTHROUGH_ATTRS — called out in the PR body under "Framework change". Happy to discuss the approach if you'd prefer a different way to handle those. |
…RES + CANONICAL_AUTHORS_BY_ARCH)
1 task
Collaborator
|
Thanks @mukund1985! I'll review once the CI checks finish |
Collaborator
|
@mukund1985 this looks great! Merging |
This was referenced Jul 6, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Closes #1463
What this adds
Zamba2ArchitectureAdapterforZamba2ForCausalLM(Zyphra/Zamba2-1.2B, Zamba2-7B, Zamba2-7B-Instruct).Zamba2 is a hybrid Mamba-2 / shared global-attention model. Most layers are pure Mamba-2 SSM (
"mamba"inlayers_block_type); a recurring subset are hybrid layers ("hybrid") that run through a sharedZamba2AttentionDecoderLayerwhose weights are tied acrossnum_mem_blocksslots.Design
SSMBlockBridge for all layers.
hook_in/hook_outfire on every layer regardless of type. For Mamba layers,normandmixersubmodules are wired normally (exposingin_proj,conv1d,inner_norm,out_projhooks). For hybrid layers, both are markedoptional=Truesocomponent_setupskips them gracefully — hybrid layers have no top-level.input_layernormor.mamba.is_stateful=False. Zamba2 threads its unified cache viapast_key_values(the standard KV path). Usingis_stateful=Truerouted generation into the Mambacache_params=path, which collided with Zamba2's own internal cache call and produced wrong tokens. Verified empirically that the standard KV path produces token-identical output tohf_model.generate().positional_embedding_type="none". RoPE is handled inside the HF attention block viaposition_ids; no model-level rotary bridge needed.Framework change (shared files)
Five config attributes (
mamba_expand,mamba_ngroups,num_mem_blocks,layers_block_type,use_shared_attention_adapter) were not in_HF_PASSTHROUGH_ATTRS, so they silently fell back to adapter defaults. Added them to bothsources/transformers.pyandsources/_bridge_builder.py. These are Mamba-family attributes that future SSM adapters will also need.Tests
21 integration tests against
Zyphra/Zamba2-1.2B(38 layers: 32 mamba + 6 hybrid):max_diff == 0.0)hf_model.generate()bit-for-bit