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

1"""Granite architecture adapter. 

2 

3Base adapter for the IBM Granite model family. Provides shared config setup and 

4helper methods used by GraniteMoe and GraniteMoeHybrid variants. 

5""" 

6 

7from typing import Any 

8 

9import torch 

10 

11from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter 

12from transformer_lens.model_bridge.generalized_components import ( 

13 AttentionBridge, 

14 EmbeddingBridge, 

15 GatedMLPBridge, 

16 MoEBridge, 

17 MoERouterBridge, 

18 RMSNormalizationBridge, 

19 RotaryEmbeddingBridge, 

20 ScaledResidualBlockBridge, 

21 UnembeddingBridge, 

22) 

23 

24 

25class GraniteArchitectureAdapter(ArchitectureAdapter): 

26 """Architecture adapter for IBM Granite models (dense). 

27 

28 Granite is a Llama-like architecture with RMSNorm, rotary position embeddings 

29 (RoPE), GQA, and a gated MLP (SiLU activation). Granite-specific scaling 

30 multipliers are applied by the HF model's native forward pass; 

31 ScaledResidualBlockBridge accounts for residual_multiplier so 

32 hook_attn_out / hook_mlp_out expose the scaled residual contributions. 

33 

34 Optional Parameters (may not exist in state_dict): 

35 ------------------------------------------------- 

36 Granite models do NOT have biases on attention and MLP projections: 

37 

38 - blocks.{i}.attn.b_Q/b_K/b_V/b_O - No bias on attention projections 

39 - blocks.{i}.mlp.b_in/b_gate/b_out - No bias on MLP projections 

40 - blocks.{i}.ln1.b, blocks.{i}.ln2.b, ln_final.b - RMSNorm has no bias 

41 """ 

42 

43 _testing_hybrid = True 

44 _testing_eager = None 

45 

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

47 """Initialize the Granite architecture adapter.""" 

48 super().__init__(cfg) 

49 

50 self._setup_common_config(cfg) 

51 self.weight_processing_conversions = {**self._qkvo_weight_conversions()} 

52 self.component_mapping = self._build_component_mapping() 

53 

54 def _setup_common_config(self, cfg: Any) -> None: 

55 """Set up config variables shared across all Granite variants.""" 

56 self._set_rms_rotary_defaults() 

57 self.cfg.default_prepend_bos = False 

58 

59 def _build_attention_bridge(self, optional: bool = False) -> AttentionBridge: 

60 """Build the standard Granite attention bridge (GraniteMoeHybrid passes optional).""" 

61 return self._qkvo_attention_bridge(optional=optional) 

62 

63 def _build_mlp_bridge(self) -> GatedMLPBridge: 

64 """Build the dense gated MLP bridge.""" 

65 return self._gated_mlp() 

66 

67 def _build_moe_bridge(self) -> MoEBridge: 

68 """Sparse MoE block with a hookable router. 

69 

70 logits_index=-1: GraniteMoe-family routers return 

71 (top_k_index, top_k_weights, router_logits) — index 0 is an int64 

72 index tensor, so the default would hook the wrong element. 

73 """ 

74 return MoEBridge( 

75 name="block_sparse_moe", 

76 config=self.cfg, 

77 submodules={ 

78 "gate": MoERouterBridge( 

79 name="router", logits_index=-1, indices_index=0, weights_index=1 

80 ) 

81 }, 

82 ) 

83 

84 def _build_component_mapping(self) -> dict: 

85 """Build the full component mapping for dense Granite.""" 

86 return { 

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

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

89 # HF multiplies each sublayer output by residual_multiplier before the 

90 # residual add, so hook_attn_out / hook_mlp_out must expose the scaled 

91 # contribution, not the raw module output. 

92 "blocks": ScaledResidualBlockBridge( 

93 name="model.layers", 

94 submodules={ 

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

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

97 "attn": self._build_attention_bridge(), 

98 "mlp": self._build_mlp_bridge(), 

99 }, 

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

101 ), 

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

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

104 } 

105 

106 def apply_output_logits_transform(self, logits: torch.Tensor) -> torch.Tensor: 

107 """Match Granite's ``lm_head / logits_scaling`` output path.""" 

108 scaling = float(getattr(self.cfg, "logits_scaling", 1.0)) 

109 return super().apply_output_logits_transform(logits / scaling)