Coverage for transformer_lens/model_bridge/supported_architectures/gemma2.py: 83%

1"""Gemma2 architecture adapter.""" 

2 

3from typing import Any 

4 

5from transformer_lens.conversion_utils.conversion_steps import TransposeTensorConversion 

6from transformer_lens.conversion_utils.param_processing_conversion import ( 

7 ParamProcessingConversion, 

8) 

9from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter 

10from transformer_lens.model_bridge.generalized_components import ( 

11 BlockBridge, 

12 EmbeddingBridge, 

13 LinearBridge, 

14 PositionEmbeddingsAttentionBridge, 

15 RMSNormalizationBridge, 

16 RotaryEmbeddingBridge, 

17 UnembeddingBridge, 

18) 

19 

20 

21class Gemma2ArchitectureAdapter(ArchitectureAdapter): 

22 """Architecture adapter for Gemma2 models.""" 

23 

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

25 """Initialize the Gemma2 architecture adapter.""" 

26 super().__init__(cfg) 

27 

28 self._set_rms_rotary_defaults() 

29 

30 self.cfg.uses_rms_norm = True 

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

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

33 self.cfg.rmsnorm_uses_offset = True 

34 

35 # Gemma2 uses logit softcapping 

36 if hasattr(self.cfg, "final_logit_softcapping"): 36 ↛ 37line 36 didn't jump to line 37 because the condition on line 36 was never true

37 self.cfg.output_logits_soft_cap = self.cfg.final_logit_softcapping 

38 if hasattr(self.cfg, "attn_logit_softcapping"): 38 ↛ 39line 38 didn't jump to line 39 because the condition on line 38 was never true

39 self.cfg.attn_scores_soft_cap = self.cfg.attn_logit_softcapping 

40 

41 # Note: n_key_value_heads is now automatically mapped from num_key_value_heads 

42 # by map_default_transformer_lens_config() in sources/transformers.py 

43 

44 self.weight_processing_conversions = { 

45 # NOTE: Gemma2 scales embeddings by sqrt(d_model) at RUNTIME inside 

46 # Gemma2TextScaledWordEmbedding.forward() (HF transformers >= 5.0). 

47 # That layer is what bridge.embed wraps, so embed.hook_out already 

48 # captures the scaled value — matching HookedTransformer's hook_embed 

49 # (which uses pre-scaled W_E). We must NOT pre-scale weights here and 

50 # we must NOT install a runtime hook_conversion that re-scales. 

51 **self._qkvo_weight_conversions(), 

52 # Norm weights carry no conversion entry: cfg.rmsnorm_uses_offset is the 

53 # only offset signal — weight processing and the runtime norm read it 

54 # directly and apply (1 + w) themselves. 

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

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

57 tensor_conversion=TransposeTensorConversion(), 

58 ), 

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

60 tensor_conversion=TransposeTensorConversion(), 

61 ), 

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

63 tensor_conversion=TransposeTensorConversion(), 

64 ), 

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

66 "unembed.weight": ParamProcessingConversion( 

67 tensor_conversion=TransposeTensorConversion(), 

68 ), 

69 } 

70 

71 self.component_mapping = { 

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

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

74 "blocks": BlockBridge( 

75 name="model.layers", 

76 config=self.cfg, 

77 submodules={ 

78 # Gemma 2 uses RMSNorm for all normalization layers 

79 **self._block_norms(), 

80 # Gemma 2 uses PositionEmbeddingsAttentionBridge like Gemma 3 

81 "attn": PositionEmbeddingsAttentionBridge( 

82 name="self_attn", 

83 config=self.cfg, 

84 submodules={ 

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

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

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

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

89 }, 

90 requires_attention_mask=True, 

91 requires_position_embeddings=True, 

92 ), 

93 "mlp": self._gated_mlp(), 

94 }, 

95 ), 

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

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

98 } 

99 

100 def _block_norms(self): 

101 """Norm-layout seam; VaultGemma drops the two post-norms.""" 

102 return { 

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

104 "ln1_post": RMSNormalizationBridge(name="post_attention_layernorm", config=self.cfg), 

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

106 "ln2_post": RMSNormalizationBridge(name="post_feedforward_layernorm", config=self.cfg), 

107 }