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

14 GatedMLPBridge, 

15 LinearBridge, 

16 MoEBridge, 

17 MoERouterBridge, 

18 PositionEmbeddingsAttentionBridge, 

19 RMSNormalizationBridge, 

20 RotaryEmbeddingBridge, 

21 ScaledResidualBlockBridge, 

22 UnembeddingBridge, 

23) 

24 

25 

26class GraniteArchitectureAdapter(ArchitectureAdapter): 

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

28 

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

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

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

32 ScaledResidualBlockBridge accounts for residual_multiplier so 

33 hook_attn_out / hook_mlp_out expose the scaled residual contributions. 

34 

35 Optional Parameters (may not exist in state_dict): 

36 ------------------------------------------------- 

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

38 

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

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

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

42 """ 

43 

44 _testing_hybrid = True 

45 _testing_eager = None 

46 

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

48 """Initialize the Granite architecture adapter.""" 

49 super().__init__(cfg) 

50 

51 self._setup_common_config(cfg) 

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

53 self.component_mapping = self._build_component_mapping() 

54 

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

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

57 self._set_rms_rotary_defaults() 

58 self.cfg.default_prepend_bos = False 

59 

60 def _build_attention_bridge(self, optional: bool = False) -> PositionEmbeddingsAttentionBridge: 

61 """Build the standard Granite attention bridge.""" 

62 return PositionEmbeddingsAttentionBridge( 

63 name="self_attn", 

64 config=self.cfg, 

65 optional=optional, 

66 submodules={ 

67 "q": LinearBridge(name="q_proj"), 

68 "k": LinearBridge(name="k_proj"), 

69 "v": LinearBridge(name="v_proj"), 

70 "o": LinearBridge(name="o_proj"), 

71 }, 

72 requires_attention_mask=True, 

73 requires_position_embeddings=True, 

74 ) 

75 

76 def _build_mlp_bridge(self) -> GatedMLPBridge: 

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

78 return self._gated_mlp() 

79 

80 def _build_moe_bridge(self) -> MoEBridge: 

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

82 

83 logits_index=-1: GraniteMoe-family routers return 

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

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

86 """ 

87 return MoEBridge( 

88 name="block_sparse_moe", 

89 config=self.cfg, 

90 submodules={"gate": MoERouterBridge(name="router", logits_index=-1)}, 

91 ) 

92 

93 def _build_component_mapping(self) -> dict: 

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

95 return { 

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

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

98 # HF multiplies each sublayer output by residual_multiplier before the 

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

100 # contribution, not the raw module output. 

101 "blocks": ScaledResidualBlockBridge( 

102 name="model.layers", 

103 submodules={ 

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

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

106 "attn": self._build_attention_bridge(), 

107 "mlp": self._build_mlp_bridge(), 

108 }, 

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

110 ), 

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

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

113 } 

114 

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

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

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

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