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

5 RMSNormalizationBridge, 

6 RotaryEmbeddingBridge, 

7 ScaledResidualBlockBridge, 

8 UnembeddingBridge, 

9) 

10from transformer_lens.model_bridge.supported_architectures.granite import ( 

11 GraniteArchitectureAdapter, 

12) 

13 

14 

15class GraniteMoeArchitectureAdapter(GraniteArchitectureAdapter): 

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

17 

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

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

20 top-k routing. 

21 """ 

22 

23 def _build_component_mapping(self) -> dict: 

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

25 return { 

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

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

28 # HF multiplies each sublayer output by residual_multiplier before the 

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

30 "blocks": ScaledResidualBlockBridge( 

31 name="model.layers", 

32 submodules={ 

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

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

35 "attn": self._build_attention_bridge(), 

36 "mlp": self._build_moe_bridge(), 

37 }, 

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

39 ), 

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

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

42 }