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 EmbeddingBridge similarly exposes embedding_multiplier at hook_embed, 

34 while per-head hook_result includes the residual multiplier. 

35 

36 Optional Parameters (may not exist in state_dict): 

37 ------------------------------------------------- 

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

39 

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

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

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

43 """ 

44 

45 _testing_hybrid = True 

46 _testing_eager = None 

47 

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

49 """Initialize the Granite architecture adapter.""" 

50 super().__init__(cfg) 

51 

52 self._setup_common_config(cfg) 

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

54 self.component_mapping = self._build_component_mapping() 

55 

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

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

58 self._set_rms_rotary_defaults() 

59 self.cfg.default_prepend_bos = False 

60 

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

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

63 return self._qkvo_attention_bridge(optional=optional) 

64 

65 def _build_embedding_bridge(self) -> EmbeddingBridge: 

66 """Expose the embedding's scaled residual contribution at hook_out.""" 

67 return EmbeddingBridge( 

68 name="model.embed_tokens", 

69 output_contribution_scale=getattr(self.cfg, "embedding_multiplier", 1.0), 

70 ) 

71 

72 def _build_mlp_bridge(self) -> GatedMLPBridge: 

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

74 return self._gated_mlp() 

75 

76 def _build_moe_bridge(self) -> MoEBridge: 

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

78 

79 logits_index=-1: GraniteMoe-family routers return 

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

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

82 """ 

83 return MoEBridge( 

84 name="block_sparse_moe", 

85 config=self.cfg, 

86 submodules={ 

87 "gate": MoERouterBridge( 

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

89 ) 

90 }, 

91 ) 

92 

93 def _build_component_mapping(self) -> dict: 

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

95 return { 

96 "embed": self._build_embedding_bridge(), 

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)