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

1"""Granite MoE architecture adapter.""" 

2 

3from transformer_lens.model_bridge.generalized_components import ( 

4 RMSNormalizationBridge, 

5 RotaryEmbeddingBridge, 

6 ScaledResidualBlockBridge, 

7 UnembeddingBridge, 

8) 

9from transformer_lens.model_bridge.supported_architectures.granite import ( 

10 GraniteArchitectureAdapter, 

11) 

12 

13 

14class GraniteMoeArchitectureAdapter(GraniteArchitectureAdapter): 

15 """Architecture adapter for IBM Granite MoE models. 

16 

17 Identical to dense Granite but replaces the gated MLP with a Sparse Mixture 

18 of Experts block (block_sparse_moe) using batched expert parameters and 

19 top-k routing. 

20 """ 

21 

22 def _build_component_mapping(self) -> dict: 

23 """Build component mapping with MoE instead of dense MLP.""" 

24 return { 

25 "embed": self._build_embedding_bridge(), 

26 "rotary_emb": RotaryEmbeddingBridge(name="model.rotary_emb"), 

27 # HF multiplies each sublayer output by residual_multiplier before the 

28 # residual add; hook_attn_out / hook_mlp_out expose the scaled contribution. 

29 "blocks": ScaledResidualBlockBridge( 

30 name="model.layers", 

31 submodules={ 

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

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

34 "attn": self._build_attention_bridge(), 

35 "mlp": self._build_moe_bridge(), 

36 }, 

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

38 ), 

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

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

41 }