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

1"""Bamba (``BambaForCausalLM``) adapter: Jamba-lineage hybrid alternating Mamba-2 

2and llama-style GQA attention mixers per ``config.layers_block_type``.""" 

3 

4from typing import Any 

5 

6from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter 

7from transformer_lens.model_bridge.generalized_components import ( 

8 BlockBridge, 

9 EmbeddingBridge, 

10 GatedRMSNormBridge, 

11 LinearBridge, 

12 PositionEmbeddingsAttentionBridge, 

13 RMSNormalizationBridge, 

14 RotaryEmbeddingBridge, 

15 SSM2MixerBridge, 

16 UnembeddingBridge, 

17) 

18from transformer_lens.model_bridge.generalized_components.depthwise_conv1d import ( 

19 DepthwiseConv1DBridge, 

20) 

21 

22 

23class BambaArchitectureAdapter(ArchitectureAdapter): 

24 """Architecture adapter for BambaForCausalLM models. 

25 

26 Both mixers are mapped optional — each present only on its layer type. 

27 The Mamba-2 mixer is wired under the canonical ``.mixer`` slot (HF path 

28 ``.mamba``) so SSM analyses reach it as on GraniteMoeHybrid / NemotronH. 

29 """ 

30 

31 _testing_hybrid = True 

32 _testing_eager = "config" 

33 

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

35 """Initialize the Bamba architecture adapter.""" 

36 super().__init__(cfg) 

37 

38 self._set_rms_rotary_defaults() 

39 # Mamba layers require per-step SSM state; generation is stateful. 

40 self.cfg.is_stateful = True 

41 

42 # Normalize the per-layer mixer-type list as cfg.layers_block_type so 

43 # analysis tools can find the Mamba layers, as on the hybrid siblings. 

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

45 

46 # Mixed mamba/attention layers: keep raw HF weight layout. 

47 self.supports_fold_ln = False 

48 self.weight_processing_conversions = {} 

49 

50 self.component_mapping = { 

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

52 "rotary_emb": RotaryEmbeddingBridge(name="model.rotary_emb", config=self.cfg), 

53 "blocks": BlockBridge( 

54 name="model.layers", 

55 submodules={ 

56 "ln1": RMSNormalizationBridge(name="input_layernorm", config=self.cfg), 

57 "attn": PositionEmbeddingsAttentionBridge( 

58 name="self_attn", 

59 config=self.cfg, 

60 optional=True, 

61 submodules={ 

62 "q": LinearBridge(name="q_proj"), 

63 "k": LinearBridge(name="k_proj"), 

64 "v": LinearBridge(name="v_proj"), 

65 "o": LinearBridge(name="o_proj"), 

66 }, 

67 requires_attention_mask=True, 

68 requires_position_embeddings=True, 

69 ), 

70 "mixer": SSM2MixerBridge( 

71 name="mamba", 

72 config=self.cfg, 

73 optional=True, 

74 submodules={ 

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

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

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

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

79 }, 

80 ), 

81 "ln2": RMSNormalizationBridge(name="pre_ff_layernorm", config=self.cfg), 

82 "mlp": self._gated_mlp(name="feed_forward"), 

83 }, 

84 ), 

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

86 "unembed": UnembeddingBridge(name="lm_head", config=self.cfg), 

87 } 

88 

89 def create_stateful_cache( 

90 self, 

91 hf_model: Any, 

92 batch_size: int, 

93 device: Any, 

94 dtype: Any, 

95 ) -> Any: 

96 """Unified DynamicCache carrying KV entries and SSM conv/recurrent state.""" 

97 from transformers.cache_utils import DynamicCache 

98 

99 return DynamicCache(config=hf_model.config)