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 GatedRMSNormBridge, 

19 JointGateUpMLPBridge, 

20 LinearBridge, 

21 RMSNormalizationBridge, 

22 RotaryEmbeddingBridge, 

23 ScaledResidualBlockBridge, 

24 SSM2MixerBridge, 

25 UnembeddingBridge, 

26) 

27from transformer_lens.model_bridge.generalized_components.depthwise_conv1d import ( 

28 DepthwiseConv1DBridge, 

29) 

30from transformer_lens.model_bridge.supported_architectures.granite import ( 

31 GraniteArchitectureAdapter, 

32) 

33 

34 

35class GraniteMoeHybridArchitectureAdapter(GraniteArchitectureAdapter): 

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

37 

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

39 universal. Inherits Granite config and attention bridge construction. 

40 """ 

41 

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

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

44 

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

46 ArchitectureAdapter.__init__(self, cfg) 

47 self._setup_common_config(cfg) 

48 

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

50 if pos_emb_type != "rope": 

51 self.cfg.positional_embedding_type = "none" 

52 

53 self.supports_fold_ln = False 

54 self.weight_processing_conversions = {} 

55 

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

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

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

59 

60 self.component_mapping = self._build_component_mapping() 

61 

62 def _build_mamba_bridge(self) -> SSM2MixerBridge: 

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

64 

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

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

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

68 """ 

69 return SSM2MixerBridge( 

70 name="mamba", 

71 config=self.cfg, 

72 optional=True, 

73 submodules={ 

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

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

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

77 }, 

78 ) 

79 

80 def _build_component_mapping(self) -> dict: 

81 block_submodules: dict = { 

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

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

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

85 "mixer": self._build_mamba_bridge(), 

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

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

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

89 "shared_mlp": JointGateUpMLPBridge( 

90 name="shared_mlp", 

91 config=self.cfg, 

92 fused_attr="input_linear", 

93 submodules={ 

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

95 }, 

96 ), 

97 } 

98 

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

100 self.cfg, "num_local_experts", 0 

101 ) 

102 if num_experts and num_experts > 0: 

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

104 

105 # HF multiplies each sublayer output by residual_multiplier before the 

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

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

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

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

110 # on the scaled shared_mlp output. 

111 mapping: dict = { 

112 "embed": self._build_embedding_bridge(), 

113 "blocks": ScaledResidualBlockBridge( 

114 name="model.layers", 

115 submodules=block_submodules, 

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

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

118 ), 

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

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

121 } 

122 

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

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

125 

126 return mapping