Coverage for transformer_lens/model_bridge/supported_architectures/mamba2.py: 96%

1"""Architecture adapter for HF's Mamba2ForCausalLM, plus the effective attention helper.""" 

2import warnings 

3from typing import Any, Dict, Optional, Union 

4 

5import torch 

6 

7from transformer_lens.ActivationCache import ActivationCache 

8from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter 

9from transformer_lens.model_bridge.bridge import TransformerBridge 

10from transformer_lens.model_bridge.generalized_components import ( 

11 DepthwiseConv1DBridge, 

12 EmbeddingBridge, 

13 GatedRMSNormBridge, 

14 LinearBridge, 

15 RMSNormalizationBridge, 

16 SSM2MixerBridge, 

17 SSMBlockBridge, 

18 UnembeddingBridge, 

19) 

20 

21 

22class Mamba2ArchitectureAdapter(ArchitectureAdapter): 

23 """Wraps HF's Mamba2ForCausalLM. 

24 

25 Differs from Mamba-1 at the mixer level: fused in_proj (no x_proj/dt_proj), 

26 two-input inner norm, multi-head structure with ``num_heads``/``head_dim``/ 

27 ``n_groups``, and an ``[num_heads]``-shaped ``dt_bias``. Shares 

28 ``SSMBlockBridge``, ``DepthwiseConv1DBridge``, and the stateful generation 

29 loop with Mamba-1. 

30 """ 

31 

32 # White-box forward: P1 is exact vs raw HF (mixer delegates to HF); P2/P3 skip 

33 # without a HookedTransformer; P4 is generation. 

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

35 

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

37 super().__init__(cfg) 

38 

39 self.cfg.normalization_type = "RMS" 

40 self.cfg.uses_rms_norm = True 

41 self.cfg.positional_embedding_type = "none" 

42 self.cfg.gated_mlp = False 

43 self.cfg.attn_only = False 

44 self.cfg.final_rms = True 

45 self.cfg.is_stateful = True 

46 

47 # Most SSM config fields come from _HF_PASSTHROUGH_ATTRS. Mamba2Config 

48 # has no `intermediate_size` field, so we compute it from expand and 

49 # derive conv_dim from that. setattr() avoids mypy attr-defined errors 

50 # since cfg is duck-typed for architecture-specific extensions. 

51 expand = getattr(self.cfg, "expand", 2) 

52 hidden_size = self.cfg.d_model 

53 intermediate_size = expand * hidden_size 

54 setattr(self.cfg, "intermediate_size", intermediate_size) 

55 

56 num_heads = self.cfg.n_heads 

57 state_size = getattr(self.cfg, "state_size", 128) 

58 n_groups = getattr(self.cfg, "n_groups", 1) 

59 conv_dim = intermediate_size + 2 * n_groups * state_size 

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

61 

62 # HF splits in_proj 5 ways but two d_mlp slots are always size 0. 

63 # Stored so the integration test can catch a future HF change that 

64 # introduces non-zero d_mlp. 

65 in_proj_out_features = 2 * intermediate_size + conv_dim + num_heads 

66 setattr(self.cfg, "expected_in_proj_out_features", in_proj_out_features) 

67 

68 self.weight_processing_conversions = {} 

69 

70 self.component_mapping = { 

71 "embed": EmbeddingBridge(name="backbone.embeddings"), 

72 "blocks": SSMBlockBridge( 

73 name="backbone.layers", 

74 submodules={ 

75 "norm": RMSNormalizationBridge(name="norm", config=self.cfg), 

76 "mixer": SSM2MixerBridge( 

77 name="mixer", 

78 config=self.cfg, 

79 submodules={ 

80 "in_proj": LinearBridge(name="in_proj"), 

81 "conv1d": DepthwiseConv1DBridge(name="conv1d"), 

82 # TL calls this "inner_norm" to disambiguate from 

83 # the block-level norm; name="norm" is the HF path. 

84 "inner_norm": GatedRMSNormBridge(name="norm"), 

85 "out_proj": LinearBridge(name="out_proj"), 

86 }, 

87 ), 

88 }, 

89 ), 

90 "ln_final": RMSNormalizationBridge(name="backbone.norm_f", config=self.cfg), 

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

92 } 

93 

94 def create_stateful_cache( 

95 self, 

96 hf_model: Any, 

97 batch_size: int, 

98 device: Any, 

99 dtype: torch.dtype, 

100 ) -> Any: 

101 """Build a cache for the stateful generation loop.""" 

102 from transformers.cache_utils import DynamicCache 

103 from transformers.models.mamba2 import modeling_mamba2 

104 

105 cache_cls = getattr(modeling_mamba2, "Mamba2Cache", None) 

106 if cache_cls is not None: 106 ↛ 107line 106 didn't jump to line 107 because the condition on line 106 was never true

107 return cache_cls(hf_model.config, batch_size, device=device, dtype=dtype) 

108 

109 return DynamicCache(config=hf_model.config) 

110 

111 

112def compute_effective_attention( 

113 bridge: TransformerBridge, 

114 cache: ActivationCache, 

115 layer: Optional[int] = None, 

116 include_dt_scaling: bool = False, 

117) -> Union[torch.Tensor, Dict[int, torch.Tensor]]: 

118 """Mamba-2 effective attention for one or all layers. 

119 

120 .. deprecated:: 

121 Use the family-agnostic ``cache.compute_ssm_effective_attention(layer=...)`` 

122 instead. This thin wrapper delegates to it and ignores ``bridge`` (the 

123 cache already knows its model). 

124 """ 

125 warnings.warn( 

126 "mamba2.compute_effective_attention is deprecated; use " 

127 "cache.compute_ssm_effective_attention(layer=..., include_dt_scaling=...).", 

128 DeprecationWarning, 

129 stacklevel=2, 

130 ) 

131 return cache.compute_ssm_effective_attention(layer=layer, include_dt_scaling=include_dt_scaling) 

132 

133 

134def compute_ssm_state( 

135 bridge: TransformerBridge, 

136 cache: ActivationCache, 

137 layer: Optional[int] = None, 

138 time_step: Optional[int] = None, 

139) -> Union[torch.Tensor, Dict[int, torch.Tensor]]: 

140 """Reconstruct the recurrent SSM state ``S`` for one or all Mamba-2 layers. 

141 

142 .. deprecated:: 

143 Use ``cache.compute_ssm_state(layer=..., time_step=...)`` instead. This 

144 thin wrapper delegates to it and ignores ``bridge``. 

145 """ 

146 warnings.warn( 

147 "mamba2.compute_ssm_state is deprecated; use " 

148 "cache.compute_ssm_state(layer=..., time_step=...).", 

149 DeprecationWarning, 

150 stacklevel=2, 

151 ) 

152 return cache.compute_ssm_state(layer=layer, time_step=time_step)