Coverage for transformer_lens/model_bridge/supported_architectures/zamba2.py: 100%

1"""Zamba2 hybrid Mamba2-Transformer architecture adapter. 

2 

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

4 

5Architecture overview: 

6- Heterogeneous layers defined by ``config.layers_block_type`` — surfaced with 

7 canonical names: ``"linear_attention"`` (pure Mamba-2 SSM, HF ``"mamba"``) or 

8 ``"hybrid"`` (Mamba-2 + shared global-attention block). 

9- Most layers are ``Zamba2MambaDecoderLayer``: a single pre-norm 

10 (``input_layernorm``) followed by a Mamba-2 mixer (``.mamba``). 

11- A recurring subset are ``Zamba2HybridLayer``: each wraps a Mamba-2 decoder 

12 layer plus a SHARED ``Zamba2AttentionDecoderLayer``. The shared attention 

13 block's weights are tied across all hybrid layers, cycling through 

14 ``config.num_mem_blocks`` unique blocks. When 

15 ``config.use_shared_attention_adapter=True``, each hybrid layer carries 

16 an independent per-layer LoRA adapter on top of the shared attention. 

17- No model-level rotary embedding module is wired by the bridge — the 

18 attention block handles RoPE internally via ``position_ids``. 

19- Generation runs on the standard KV-cache path: HF threads a single unified 

20 ``Zamba2HybridDynamicCache`` via ``past_key_values`` (carrying both KV-cache 

21 entries for attention and SSM conv/recurrent states for Mamba-2), so the 

22 bridge does not use the Mamba-specific ``cache_params`` stateful path. 

23 

24Key adapter decisions: 

25- ``SSMBlockBridge`` is used for all layers. Its forward delegates entirely to 

26 the HF layer, giving ``hook_in`` / ``hook_out`` on every layer regardless 

27 of type. 

28- For Mamba layers: ``norm`` (-> ``.input_layernorm``) and ``mixer`` 

29 (-> ``.mamba``) are declared as submodules and expose inner hooks 

30 (in_proj, conv1d, inner_norm, out_proj). 

31- For Hybrid layers: ``norm`` and ``mixer`` are marked ``optional=True`` so 

32 component_setup skips them gracefully (Hybrid layers have no top-level 

33 ``.input_layernorm`` or ``.mamba``). Block-level ``hook_in``/``hook_out`` 

34 still fire on every layer. 

35- ``applicable_phases = [1, 2, 3, 4]``: P1 is exact vs raw HF (pure 

36 passthrough); P2/P3 skip without a HookedTransformer; P4 exercises 

37 ``past_key_values`` cache threading across Mamba-2 and attention layers. 

38""" 

39 

40from typing import Any 

41 

42from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter 

43from transformer_lens.model_bridge.generalized_components import ( 

44 DepthwiseConv1DBridge, 

45 EmbeddingBridge, 

46 GatedRMSNormBridge, 

47 LinearBridge, 

48 RMSNormalizationBridge, 

49 SSM2MixerBridge, 

50 SSMBlockBridge, 

51 UnembeddingBridge, 

52) 

53from transformer_lens.model_bridge.generalized_components.base import ( 

54 GeneralizedComponent, 

55) 

56 

57 

58def _make_optional(component: "GeneralizedComponent") -> "GeneralizedComponent": 

59 """Mark a GeneralizedComponent submodule as optional. 

60 

61 ``component_setup.py`` reads ``getattr(submodule, 'optional', False)`` at 

62 setup time, so setting the attribute directly is safe regardless of whether 

63 the bridge class's ``__init__`` accepts an ``optional`` keyword argument. 

64 """ 

65 component.optional = True 

66 return component 

67 

68 

69class Zamba2ArchitectureAdapter(ArchitectureAdapter): 

70 """Architecture adapter for Zamba2ForCausalLM. 

71 

72 Hybrid Mamba-2 + shared global-attention model. Most layers are pure 

73 Mamba-2 SSM (``"mamba"``); a recurring subset are hybrid layers 

74 (``"hybrid"``) that route through a shared attention block before the 

75 Mamba-2 step. 

76 """ 

77 

78 # P1: exact passthrough vs raw HF; P2/P3: skip without HookedTransformer; 

79 # P4: generation with past_key_values cache threading. 

80 applicable_phases: list[int] = [1, 2, 3, 4] 

81 

82 def __init__(self, cfg: Any) -> None: 

83 super().__init__(cfg) 

84 

85 self.cfg.normalization_type = "RMS" 

86 self.cfg.uses_rms_norm = True 

87 # RoPE is handled internally by the attention block (position_ids are 

88 # threaded through HF layers); no model-level rotary bridge needed. 

89 self.cfg.positional_embedding_type = "none" 

90 # MLP inside the shared attention block uses GELU, not SwiGLU. 

91 self.cfg.gated_mlp = False 

92 self.cfg.attn_only = False 

93 self.cfg.final_rms = True 

94 # NOTE: is_stateful stays False even though Mamba-2 layers carry SSM 

95 # state. In this bridge, ``is_stateful=True`` selects the *Mamba* cache 

96 # path, which threads the cache as ``cache_params=`` and drives a 

97 # conv-kernel ``cache_position``. Zamba2's HF forward instead threads a 

98 # single unified ``Zamba2HybridDynamicCache`` via ``past_key_values=`` 

99 # (both KV entries and SSM conv/recurrent states live in that one 

100 # object). The standard KV-cache generation path already threads 

101 # ``past_key_values`` correctly and matches HF ``generate`` bit-for-bit, 

102 # so we use it rather than the Mamba path (whose ``cache_params`` kwarg 

103 # collides with Zamba2's own ``cache_params=past_key_values`` call). 

104 self.cfg.is_stateful = False 

105 

106 # Expose the per-layer type list so analysis tools can identify which 

107 # layers are Mamba-only vs hybrid: "linear_attention" | "hybrid". 

108 setattr(self.cfg, "layers_block_type", self._canonical_layer_types(cfg)) 

109 

110 # Number of unique shared attention weight blocks (hybrid layers cycle 

111 # through num_mem_blocks independent attention weight sets). 

112 setattr(self.cfg, "num_mem_blocks", getattr(cfg, "num_mem_blocks", 1)) 

113 

114 # Whether per-layer LoRA adapters are active on the shared attention. 

115 setattr( 

116 self.cfg, 

117 "use_shared_attention_adapter", 

118 getattr(cfg, "use_shared_attention_adapter", False), 

119 ) 

120 

121 # Mamba-2 dimensional config (mirrors Mamba2ArchitectureAdapter / 

122 # NemotronHArchitectureAdapter patterns). 

123 mamba_intermediate_size = int(getattr(cfg, "mamba_expand", 2) * self.cfg.d_model) 

124 n_groups = getattr(cfg, "mamba_ngroups", 1) 

125 ssm_state_size = getattr(cfg, "mamba_d_state", 64) 

126 conv_dim = mamba_intermediate_size + 2 * n_groups * ssm_state_size 

127 setattr(self.cfg, "mamba_intermediate_size", mamba_intermediate_size) 

128 setattr(self.cfg, "conv_dim", conv_dim) 

129 

130 self.weight_processing_conversions = {} 

131 

132 self.component_mapping = { 

133 "embed": EmbeddingBridge(name="model.embed_tokens"), 

134 "blocks": SSMBlockBridge( 

135 name="model.layers", 

136 submodules={ 

137 # Pre-norm: present on Zamba2MambaDecoderLayer (.input_layernorm), 

138 # absent on Zamba2HybridLayer (no top-level pre-norm) -> optional. 

139 "norm": _make_optional( 

140 RMSNormalizationBridge(name="input_layernorm", config=self.cfg) 

141 ), 

142 # Mamba-2 mixer: present on Zamba2MambaDecoderLayer (.mamba), 

143 # absent at the top level of Zamba2HybridLayer -> optional. 

144 "mixer": _make_optional( 

145 SSM2MixerBridge( 

146 name="mamba", 

147 config=self.cfg, 

148 submodules={ 

149 # -- Mamba-2 inner submodules (all optional) -- 

150 "in_proj": LinearBridge(name="in_proj", optional=True), 

151 "conv1d": DepthwiseConv1DBridge(name="conv1d", optional=True), 

152 # HF names the gated RMS norm "norm" inside the 

153 # mixer; TL uses "inner_norm" to avoid colliding 

154 # with the block-level norm declared above. 

155 "inner_norm": _make_optional(GatedRMSNormBridge(name="norm")), 

156 "out_proj": LinearBridge(name="out_proj", optional=True), 

157 }, 

158 ) 

159 ), 

160 }, 

161 ), 

162 "ln_final": RMSNormalizationBridge(name="model.final_layernorm", config=self.cfg), 

163 "unembed": UnembeddingBridge(name="lm_head"), 

164 }