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

1"""Granite MoE Hybrid architecture adapter. 

2 

3Hybrid Mamba2 + Attention with Sparse MoE. Most layers are Mamba SSM blocks; 

4a few are standard attention (determined by config.layer_types). Every layer 

5has a shared MLP and optional sparse MoE. 

6 

7Both attention and Mamba are mapped as optional — each present only on its 

8respective layer type. The Mamba mixer is wired under the canonical ``.mixer`` 

9slot (HF path is ``.mamba``) so SSM analyses (compute_effective_attention) 

10reach it the same way as on NemotronH / Mamba-2. Mamba hooks expose in_proj, 

11conv1d, and inner_norm (the gated two-input MambaRMSNormGated). 

12""" 

13 

14from typing import Any 

15 

16from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter 

17from transformer_lens.model_bridge.generalized_components import ( 

18 EmbeddingBridge, 

19 GatedRMSNormBridge, 

20 JointGateUpMLPBridge, 

21 LinearBridge, 

22 RMSNormalizationBridge, 

23 RotaryEmbeddingBridge, 

24 ScaledResidualBlockBridge, 

25 SSM2MixerBridge, 

26 UnembeddingBridge, 

27) 

28from transformer_lens.model_bridge.generalized_components.depthwise_conv1d import ( 

29 DepthwiseConv1DBridge, 

30) 

31from transformer_lens.model_bridge.supported_architectures.granite import ( 

32 GraniteArchitectureAdapter, 

33) 

34 

35 

36class GraniteMoeHybridArchitectureAdapter(GraniteArchitectureAdapter): 

37 """Hybrid Mamba2 + Attention with Sparse MoE. 

38 

39 Attention is optional (absent on Mamba layers). shared_mlp and MoE are 

40 universal. Inherits Granite config and attention bridge construction. 

41 """ 

42 

43 # Explicit for parity with the Mamba/NemotronH siblings. 

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

45 

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

47 ArchitectureAdapter.__init__(self, cfg) 

48 self._setup_common_config(cfg) 

49 

50 pos_emb_type = getattr(cfg, "position_embedding_type", "rope") 

51 if pos_emb_type != "rope": 

52 self.cfg.positional_embedding_type = "none" 

53 

54 self.supports_fold_ln = False 

55 self.weight_processing_conversions = {} 

56 

57 # Normalize the per-layer mixer-type list as cfg.layers_block_type (HF names 

58 # it `layer_types`) so analysis tools can find the Mamba layers, as on NemotronH. 

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

60 

61 self.component_mapping = self._build_component_mapping() 

62 

63 def _build_mamba_bridge(self) -> SSM2MixerBridge: 

64 """Mamba-2 mixer bridge with in_proj, conv1d, inner_norm hooks. 

65 

66 ``name="mamba"`` is the HF submodule path; the bridge exposes it under the 

67 canonical ``.mixer`` dict key (see ``_build_component_mapping``). inner_norm 

68 is the gated two-input MambaRMSNormGated, so it wraps GatedRMSNormBridge. 

69 """ 

70 return 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 }, 

79 ) 

80 

81 def _build_component_mapping(self) -> dict: 

82 block_submodules: dict = { 

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

84 "ln2": RMSNormalizationBridge(name="post_attention_layernorm", config=self.cfg), 

85 "attn": self._build_attention_bridge(optional=True), 

86 "mixer": self._build_mamba_bridge(), 

87 # input_linear is fused [gate | up]: a plain MLPBridge made 

88 # hook_pre the 2*d_mlp pre-GLU tensor with no hook_pre_linear. The 

89 # bridge registers "gate"/"in" itself, so only "out" is declared. 

90 "shared_mlp": JointGateUpMLPBridge( 

91 name="shared_mlp", 

92 config=self.cfg, 

93 fused_attr="input_linear", 

94 submodules={ 

95 "out": LinearBridge(name="output_linear"), 

96 }, 

97 ), 

98 } 

99 

100 num_experts = getattr(self.cfg, "num_experts", None) or getattr( 

101 self.cfg, "num_local_experts", 0 

102 ) 

103 if num_experts and num_experts > 0: 

104 block_submodules["moe"] = self._build_moe_bridge() 

105 

106 # HF multiplies each sublayer output by residual_multiplier before the 

107 # residual add. hook_attn_out fires on attention layers (mamba layers have 

108 # no attention-position hook). With experts, the MLP branch is 

109 # block_sparse_moe + shared_mlp summed inline — no single module produces 

110 # the contribution, so hook_mlp_out stays absent; without experts it fires 

111 # on the scaled shared_mlp output. 

112 mapping: dict = { 

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

114 "blocks": ScaledResidualBlockBridge( 

115 name="model.layers", 

116 submodules=block_submodules, 

117 residual_contribution_scale=getattr(self.cfg, "residual_multiplier", 1.0), 

118 scaled_mlp_submodule=None if (num_experts and num_experts > 0) else "shared_mlp", 

119 ), 

120 "ln_final": RMSNormalizationBridge(name="model.norm", config=self.cfg), 

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

122 } 

123 

124 if self.cfg.positional_embedding_type == "rotary": 

125 mapping["rotary_emb"] = RotaryEmbeddingBridge(name="model.rotary_emb", config=self.cfg) 

126 

127 return mapping