Coverage for transformer_lens/model_bridge/supported_architectures/gemma3.py: 46%

1"""Gemma3 architecture adapter.""" 

2 

3 

4from typing import Any 

5 

6from transformer_lens.conversion_utils.conversion_steps import TransposeTensorConversion 

7from transformer_lens.conversion_utils.param_processing_conversion import ( 

8 ParamProcessingConversion, 

9) 

10from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter 

11from transformer_lens.model_bridge.generalized_components import ( 

12 BlockBridge, 

13 EmbeddingBridge, 

14 LinearBridge, 

15 RMSNormalizationBridge, 

16 RotaryEmbeddingBridge, 

17 UnembeddingBridge, 

18) 

19from transformer_lens.model_bridge.generalized_components.position_embeddings_attention import ( 

20 PositionEmbeddingsAttentionBridge, 

21) 

22 

23 

24def _enable_native_qk_norm_autograd(bridge_model: Any) -> None: 

25 """Component tests grad-check through the wrapped Q/K norms; delegate them 

26 to HF's native layernorm autograd (shared by text and multimodal Gemma-3).""" 

27 if bridge_model is None or not hasattr(bridge_model, "blocks"): 

28 return 

29 for block in bridge_model.blocks: 

30 hf_attn = getattr(getattr(block, "attn", None), "original_component", None) 

31 if hf_attn is None: 

32 continue 

33 if hasattr(hf_attn, "q_norm"): 

34 hf_attn.q_norm.use_native_layernorm_autograd = True 

35 if hasattr(hf_attn, "k_norm"): 

36 hf_attn.k_norm.use_native_layernorm_autograd = True 

37 

38 

39class Gemma3ArchitectureAdapter(ArchitectureAdapter): 

40 """Architecture adapter for Gemma3 models.""" 

41 

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

43 """Initialize the Gemma3 architecture adapter.""" 

44 super().__init__(cfg) 

45 

46 self.cfg.gated_mlp = True 

47 

48 self.cfg.uses_rms_norm = True 

49 self.cfg.normalization_type = "RMS" 

50 # Gemma models use (1.0 + weight) in RMSNorm instead of just weight 

51 # See: https://github.com/huggingface/transformers/pull/29402 

52 self.cfg.rmsnorm_uses_offset = True 

53 

54 # Gemma 3 uses rotary positional embeddings (dual RoPE) 

55 self.cfg.positional_embedding_type = "rotary" 

56 

57 # Use eager attention to support output_attentions for hook_attn_scores and hook_pattern 

58 # SDPA doesn't support output_attentions, which is required for HookedTransformer compatibility 

59 self.cfg.attn_implementation = "eager" 

60 

61 self.weight_processing_conversions = { 

62 # Note: Gemma3TextScaledWordEmbedding scales by sqrt(d_model) inside 

63 # its own forward(). Bridge.embed wraps that layer, so embed.hook_out 

64 # already captures the scaled value — no weight pre-scaling and no 

65 # hook_conversion needed (setup_hook_compatibility is a no-op). 

66 # 

67 # Q/K/V weight conversions 

68 **self._qkvo_weight_conversions(), 

69 # Norm weights (incl. q_norm/k_norm) carry no conversion entry: 

70 # cfg.rmsnorm_uses_offset is the only offset signal — weight processing 

71 # and the runtime norm read it directly and apply (1 + w) themselves. 

72 # MLP weight conversions - transpose from [out, in] to [in, out] 

73 "blocks.{i}.mlp.gate.weight": ParamProcessingConversion( 

74 tensor_conversion=TransposeTensorConversion(), 

75 ), 

76 "blocks.{i}.mlp.in.weight": ParamProcessingConversion( 

77 tensor_conversion=TransposeTensorConversion(), 

78 ), 

79 "blocks.{i}.mlp.out.weight": ParamProcessingConversion( 

80 tensor_conversion=TransposeTensorConversion(), 

81 ), 

82 # Unembed weight conversion - transpose from [vocab, d_model] to [d_model, vocab] 

83 "unembed.weight": ParamProcessingConversion( 

84 tensor_conversion=TransposeTensorConversion(), 

85 ), 

86 # Note: Gemma-3 does NOT have biases on attention projections (q/k/v/o_proj.bias are all None) 

87 # No bias conversions needed 

88 } 

89 

90 # Set up component mapping with actual bridge instances 

91 self.component_mapping = { 

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

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

94 "blocks": BlockBridge( 

95 name="model.layers", 

96 submodules={ 

97 # All Gemma-3 normalizations use simple RMSNorm pass-through 

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

99 "ln1_post": RMSNormalizationBridge( 

100 name="post_attention_layernorm", config=self.cfg 

101 ), 

102 "ln2": RMSNormalizationBridge( 

103 name="pre_feedforward_layernorm", config=self.cfg 

104 ), 

105 "ln2_post": RMSNormalizationBridge( 

106 name="post_feedforward_layernorm", config=self.cfg 

107 ), 

108 "attn": PositionEmbeddingsAttentionBridge( 

109 name="self_attn", 

110 config=self.cfg, 

111 submodules={ 

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

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

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

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

116 "q_norm": RMSNormalizationBridge(name="q_norm", config=self.cfg), 

117 "k_norm": RMSNormalizationBridge(name="k_norm", config=self.cfg), 

118 }, 

119 ), 

120 "mlp": self._gated_mlp(), 

121 }, 

122 ), 

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

124 "unembed": UnembeddingBridge(name="lm_head"), 

125 } 

126 

127 def setup_component_testing(self, hf_model: Any, bridge_model: Any = None) -> None: 

128 """Wire local RoPE + eager attention; q/k norms delegate to HF autograd. 

129 

130 Gemma-3 uses dual RoPE (global + local); component tests share the local 

131 instance across all layers (layers on global RoPE accept the tradeoff). 

132 """ 

133 super().setup_component_testing(hf_model, bridge_model) 

134 _enable_native_qk_norm_autograd(bridge_model)