[flyingdoubleG]

Description

Follow-up to #1986 / #1987. On Zamba2 and RecurrentGemma, the residual-decomposition APIs still raise the errors from #1986, because some blocks expose no sublayer at all:

	decompose_resid()	get_full_resid_decomposition(), stack_head_results()
Zamba2 (hybrid layers)	KeyError: 'blocks.1.hook_attn_out'	KeyError: 'blocks.1.attn.hook_z'
RecurrentGemma (every layer)	KeyError: 'blocks.0.hook_attn_out'	KeyError: 'blocks.0.attn.hook_z'

Both blocks declare no variant slot, so the cache falls back to treating them as attention blocks.

Zamba2. Zamba2HybridLayer has no top-level .mamba / .input_layernorm, so the optional mixer and norm were skipped on hybrid layers. The hybrid layer wraps a regular Zamba2MambaDecoderLayer as .mamba_decoder, and the shared transformer's output (after the per-layer .linear) is only added to that layer's Mamba input:

# Zamba2MambaDecoderLayer.forward
residual = hidden_states
hidden_states = hidden_states + transformer_hidden_states  # hybrid layers only
hidden_states = self.input_layernorm(hidden_states)
hidden_states = self.mamba(hidden_states, ...)
hidden_states = residual + hidden_states

So the Mamba-2 output is the only residual write on every layer. A small Zamba2BlockBridge points mixer / norm at mamba_decoder.mamba / mamba_decoder.input_layernorm when the HF layer has a .mamba_decoder. Hybrid layers then get the same mixer hooks as the Mamba layers (hook_mixer_out, mixer.in_proj, …), and the decompositions get a {l}_mixer_out row for them. On a hybrid layer norm.hook_in is the residual plus the shared transformer's output.

RecurrentGemma. The adapter wrapped each decoder layer whole, with only hook_resid_pre / hook_resid_post. Every layer has the same layout, though: temporal_pre_norm → temporal_block → residual add → channel_pre_norm → mlp_block → residual add. Only the temporal block differs (RG-LRU or local attention). This PR bridges the two norms as ln1 / ln2, the temporal block as a passthrough mixer (a plain GeneralizedComponent; cfg.layers_block_type says which kind it is), and the MLP as a GatedMLPBridge. The block gets hook_resid_mid, hook_mixer_in / hook_mixer_out and hook_mlp_out. Decompositions label the temporal block {l}_attn_out on attention layers and {l}_mixer_out on recurrent ones, followed by {l}_mlp_out, and expand_neurons=True works.

Two config fixes were needed for the sums to match:

• HF multiplies the embedding by model.normalizer after embed_tokens, and that buffer is bfloat16 (sqrt(2560) is stored as 50.5). hook_embed was the unscaled embedding, so it did not equal blocks.0.hook_resid_pre. The embedding bridge now gets output_contribution_scale set to the normalizer, as on Granite. It is computed from d_model and read from the loaded model's buffer in prepare_model. This matches Gemma-2/3, whose hook_embed is already the scaled embedding.
• RecurrentGemmaMlp is intermediate_size // 2 wide, but cfg.d_mlp was intermediate_size. That gave get_full_resid_decomposition(expand_neurons=True) twice as many labels as rows. The adapter now sets d_mlp = intermediate_size // 2.

On the tiny random models of the new tests, the max abs error of each sum against the residual stream:

	dev	this PR
Zamba2 decompose_resid()	KeyError	0
Zamba2 get_full_resid_decomposition(expand_neurons=True)	KeyError	1e-7
RecurrentGemma decompose_resid()	KeyError	0
RecurrentGemma decompose_resid(mlp_input=True)	KeyError	0
RecurrentGemma get_full_resid_decomposition(expand_neurons=True)	KeyError	2e-7

Logits still match HF exactly on both (the bridge delegates the layer forwards), and greedy generate matches HF token for token.

User-visible changes:

• Zamba2: hybrid layers now have .mixer / .norm and their hooks, and cache.ssm_layers() includes them.
• RecurrentGemma: blocks have ln1, mixer, ln2 and mlp; hook_embed is the scaled embedding; cfg.d_mlp is the real MLP width.

applicable_phases is unchanged for both.

RWKV-4 also raises KeyError in these APIs. It is not changed here: HF halves the residual stream every rescale_every layers between blocks (so blocks.{l}.hook_resid_post != blocks.{l+1}.hook_resid_pre), and layer 0 applies pre_ln to the embedding, so it needs more than a slot mapping.

Type of change

• Bug fix (non-breaking change which fixes an issue)

Checklist:

• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes
• I have not rewritten tests relating to key interfaces which would affect backward compatibility

Tests:

• tests/integration/model_bridge/test_ssm_hybrid_resid_decomposition.py gets tiny Zamba2 (mamba, hybrid, mamba, hybrid) and RecurrentGemma (recurrent, recurrent, attention, recurrent, hidden_size=80 so the bfloat16 normalizer differs from sqrt(80)) factories. Both run through the existing sum, partition, label, head and LN + projection checks. RecurrentGemma is added to the mlp_input=True check. There are also two direct tests: on Zamba2, resid_post - resid_pre == mixer.hook_out in every layer; on RecurrentGemma, hook_embed == blocks.0.hook_resid_pre and resid_pre + hook_mixer_out == resid_mid.
• The unit tests for both adapters cover the hybrid-layer redirect, the new RecurrentGemma submodules and aliases, the normalizer scale, and d_mlp.
• The slow tests/integration/model_bridge/test_zamba2_adapter.py asserted that hybrid layers have no mixer. Those two tests now check that .mixer / .norm are the wrapped decoder's modules, and that the mixer output is the hybrid layer's residual write. I could not run that file locally because it downloads Zamba2-1.2B.

On dev, 37 of the new and updated tests in test_ssm_hybrid_resid_decomposition.py and test_recurrent_gemma_adapter.py fail, and the Zamba2 unit test file cannot import Zamba2BlockBridge. All 129 tests in these three files pass here.

Local runs (Windows, CPU, torch 2.9.1, --jaxtyping-packages=transformer_lens,beartype.beartype, PYTHONUTF8=1):

• test_ssm_hybrid_resid_decomposition.py and the Zamba2 / RecurrentGemma unit tests: all pass.
• tests/unit -m "not slow": 8180 passed, 14 failed. The same 14 fail on dev in my environment (ten Gemma 3n adapter tests need the optional timm, three test_sparse_probing.py tests check optimizer stop reasons, and test_faithfulness_recovery_is_graded_not_binary is [Bug Report] test_faithfulness_recovery_is_graded_not_binary passes or fails on the sign of float32 rounding #1993).
• test_ssm_effective_attention_dispatch.py: 8 Mamba / Granite / NemotronH SSD tests fail the same way on dev here (no mixer.conv1d.hook_out in the cache), which I take to be my local transformers version.
• No shared code changes. A sweep of tiny random hybrid models still sums to the residual stream within 4e-7 on Jamba, Bamba, OLMo-Hybrid, Qwen3-Next (except expand_neurons=True on its MoE, as before), Mamba and FalconMamba.

black (the locked 23.12.1), isort --check-only and mypy on the changed files pass.

This was prepared with AI assistance (Claude); I reviewed and tested the change.

[pablocs116]

Thanks for this! As the author of the original RecurrentGemma adapter (#1483), this looks right to me. Splitting each layer into ln1/mixer/ln2/mlp is the right granularity now that the passthrough mixer handles the RG-LRU vs. local-attention difference, and the d_mlp fix is correct: RecurrentGemmaMlp sets intermediate_size = config.intermediate_size // 2, so the old value was 2x the real width.

Checked locally at 5feed31 (CPU, fp32):

• The PR's tests (test_recurrent_gemma_adapter.py, test_zamba2_adapter.py, test_ssm_hybrid_resid_decomposition.py, integration test_zamba2_adapter.py): 150 passed, 16 skipped.
• An independent script on a tiny random model (hidden_size 80, block types recurrent/recurrent/attention, 4 layers), against a deep-copied HF reference:
  • cfg.d_mlp == intermediate_size // 2 == mlp_block.gate_proj.out_features
  • logits match HF (max |diff| 7e-7), and resid_post matches HF hidden_states per layer
  • hook_embed == embed_tokens(x) * model.normalizer exactly (8.9375, the bf16 value, not sqrt(80))
  • resid_pre + mixer_out == resid_mid and resid_mid + mlp_out == resid_post on every layer, recurrent and attention
  • decompose_resid and get_full_resid_decomposition sum to the final residual
  • edits are in residual units: zeroing hook_embed matches zeroing HF's embedding, and zeroing blocks.1.hook_mlp_out matches zeroing that layer's MLP
  • greedy generate matches HF

One small, non-blocking note: the _boot helper in test_ssm_hybrid_resid_decomposition.py doesn't call adapter.prepare_model, so the new prepare_model (reading model.normalizer) isn't exercised by the tests. It's redundant with _embedding_normalizer unless the buffer and the config disagree, so this is just a coverage gap. I checked it separately by calling prepare_model the way boot_transformers does.

Also worth a line in the release notes: hook_embed and cfg.d_mlp change values for existing RecurrentGemma users (both for the better).

[SBrouck]

Thanks! @pablocs116 covered RecurrentGemma, so I focused on the Zamba2 side, which builds on the mixer slots from #1987.

The redirect renames norm / mixer in set_original_component. Your test starts with a Mamba layer, so I wanted to make sure a Mamba layer that comes after a hybrid one keeps .mamba. I checked locally at 5feed31 with an independent tiny Zamba2 that starts with a hybrid layer (hybrid, mamba, hybrid, mamba, hybrid, batch 2):

• Each block's mixer is the right HF module (same in_proj weights): mamba_decoder.mamba on hybrid layers, mamba on the others.
• resid_post - resid_pre == mixer.hook_out on every layer (1e-7 or less).
• decompose_resid() and get_full_resid_decomposition() (both expand_neurons) sum to the final residual (8e-8 or less), and ssm_layers() returns all five layers.
• Zeroing blocks.2.mixer.hook_out gives exactly the same logits as zeroing mamba_decoder.mamba's output in an untouched HF copy.
• Logits match HF exactly.

I couldn't download the 1.2B checkpoint here, so I haven't run the slow integration test_zamba2_adapter.py.

[jlarson4]

@flyingdoubleG This bug fix looks good! I was able to run all tests and can confirm everything turns up green on my end. Approved for inclusion in the next release