Coverage for transformer_lens/pretrained/weight_conversions/gemma.py: 96%

1import einops 

2import torch 

3 

4from transformer_lens.config.hooked_transformer_config import HookedTransformerConfig 

5 

6 

7def convert_gemma_weights(gemma, cfg: HookedTransformerConfig): 

8 state_dict = {} 

9 

10 assert cfg.n_key_value_heads is not None # keep mypy happy 

11 assert cfg.d_mlp is not None # keep mypy happy 

12 

13 # Multimodal detection must survive transformers 5.x, which moved 

14 # Gemma3ForConditionalGeneration's language_model under .model — the old 

15 # top-level probe silently reported text-only there and converted the 

16 # multimodal model down the wrong path. 

17 language_model = getattr(gemma, "language_model", None) 

18 if language_model is None: 

19 language_model = getattr(getattr(gemma, "model", None), "language_model", None) 

20 

21 if language_model is not None: 

22 # Vision tower is skipped entirely; only the text transformer converts. 

23 base_model = getattr(language_model, "model", language_model) 

24 else: 

25 # Text-only Gemma3ForCausalLM has .model wrapper 

26 base_model = gemma.model 

27 

28 # Gemma Models scale embeddings by multiplying by sqrt(d_model), use hidden state type to match 

29 # HF implementation 

30 state_dict["embed.W_E"] = base_model.embed_tokens.weight * torch.tensor( 

31 cfg.d_model**0.5, dtype=cfg.dtype 

32 ) 

33 

34 # Gemma has no biases anywhere 

35 for l in range(cfg.n_layers): 

36 # GemmaRMSNorm adds 1 to weights before multiplying by input, keep RMS calcs in float32 

37 state_dict[f"blocks.{l}.ln1.w"] = base_model.layers[ 

38 l 

39 ].input_layernorm.weight.float() + torch.ones_like( 

40 base_model.layers[l].input_layernorm.weight, dtype=torch.float32 

41 ) 

42 if cfg.use_normalization_before_and_after: 42 ↛ 50line 42 didn't jump to line 50 because the condition on line 42 was always true

43 # Only applies for Gemma 2 

44 state_dict[f"blocks.{l}.ln1_post.w"] = base_model.layers[ 

45 l 

46 ].post_attention_layernorm.weight.float() + torch.ones_like( 

47 base_model.layers[l].input_layernorm.weight, dtype=torch.float32 

48 ) 

49 

50 W_Q = base_model.layers[l].self_attn.q_proj.weight 

51 W_K = base_model.layers[l].self_attn.k_proj.weight 

52 W_V = base_model.layers[l].self_attn.v_proj.weight 

53 W_Q = einops.rearrange(W_Q, "(n h) m->n m h", n=cfg.n_heads) 

54 W_K = einops.rearrange(W_K, "(n h) m->n m h", n=cfg.n_key_value_heads) 

55 W_V = einops.rearrange(W_V, "(n h) m->n m h", n=cfg.n_key_value_heads) 

56 state_dict[f"blocks.{l}.attn.W_Q"] = W_Q 

57 state_dict[f"blocks.{l}.attn._W_K"] = W_K 

58 state_dict[f"blocks.{l}.attn._W_V"] = W_V 

59 

60 # Load q_norm and k_norm if they exist (Gemma 3) 

61 # Gemma3RMSNorm adds 1 to weights in forward(), so we pre-add it here 

62 if cfg.use_qk_norm: 

63 state_dict[f"blocks.{l}.attn.q_norm.w"] = base_model.layers[ 

64 l 

65 ].self_attn.q_norm.weight.float() + torch.ones_like( 

66 base_model.layers[l].self_attn.q_norm.weight, dtype=torch.float32 

67 ) 

68 state_dict[f"blocks.{l}.attn.k_norm.w"] = base_model.layers[ 

69 l 

70 ].self_attn.k_norm.weight.float() + torch.ones_like( 

71 base_model.layers[l].self_attn.k_norm.weight, dtype=torch.float32 

72 ) 

73 

74 state_dict[f"blocks.{l}.attn.b_Q"] = torch.zeros( 

75 cfg.n_heads, cfg.d_head, dtype=cfg.dtype, device=W_Q.device 

76 ) 

77 state_dict[f"blocks.{l}.attn._b_K"] = torch.zeros( 

78 cfg.n_key_value_heads, cfg.d_head, dtype=cfg.dtype, device=W_K.device 

79 ) 

80 state_dict[f"blocks.{l}.attn._b_V"] = torch.zeros( 

81 cfg.n_key_value_heads, cfg.d_head, dtype=cfg.dtype, device=W_V.device 

82 ) 

83 

84 W_O = base_model.layers[l].self_attn.o_proj.weight 

85 W_O = einops.rearrange(W_O, "m (n h)->n h m", n=cfg.n_heads) 

86 state_dict[f"blocks.{l}.attn.W_O"] = W_O 

87 

88 state_dict[f"blocks.{l}.attn.b_O"] = torch.zeros( 

89 cfg.d_model, dtype=cfg.dtype, device=W_O.device 

90 ) 

91 

92 # GemmaRMSNorm adds 1 to weights before multiplying by input, keep RMS calcs in float32 

93 if not cfg.use_normalization_before_and_after: 93 ↛ 95line 93 didn't jump to line 95 because the condition on line 93 was never true

94 # Only applies for Gemma 1. Confusingly post_attention_layernorm is applied to mlp_input in Gemma 1 and attn_out in Gemma 2 

95 state_dict[f"blocks.{l}.ln2.w"] = base_model.layers[ 

96 l 

97 ].post_attention_layernorm.weight.float() + torch.ones_like( 

98 base_model.norm.weight, dtype=torch.float32 

99 ) 

100 else: 

101 # Only applies for Gemma 2 

102 state_dict[f"blocks.{l}.ln2.w"] = base_model.layers[ 

103 l 

104 ].pre_feedforward_layernorm.weight.float() + torch.ones_like( 

105 base_model.layers[l].pre_feedforward_layernorm.weight, dtype=torch.float32 

106 ) 

107 state_dict[f"blocks.{l}.ln2_post.w"] = base_model.layers[ 

108 l 

109 ].post_feedforward_layernorm.weight.float() + torch.ones_like( 

110 base_model.layers[l].post_feedforward_layernorm.weight, dtype=torch.float32 

111 ) 

112 

113 state_dict[f"blocks.{l}.mlp.W_in"] = base_model.layers[l].mlp.up_proj.weight.T 

114 state_dict[f"blocks.{l}.mlp.W_gate"] = base_model.layers[l].mlp.gate_proj.weight.T 

115 state_dict[f"blocks.{l}.mlp.b_in"] = torch.zeros( 

116 cfg.d_mlp, dtype=cfg.dtype, device=base_model.layers[l].mlp.up_proj.weight.device 

117 ) 

118 

119 state_dict[f"blocks.{l}.mlp.W_out"] = base_model.layers[l].mlp.down_proj.weight.T 

120 state_dict[f"blocks.{l}.mlp.b_out"] = torch.zeros( 

121 cfg.d_model, dtype=cfg.dtype, device=base_model.layers[l].mlp.down_proj.weight.device 

122 ) 

123 

124 # GemmaRMSNorm adds 1 to weights before multiplying by input, keep RMS calcs in float32 

125 state_dict["ln_final.w"] = base_model.norm.weight.float() + torch.ones_like( 

126 base_model.norm.weight, dtype=torch.float32 

127 ) 

128 

129 # For multimodal models, lm_head might not exist or be tied to embeddings 

130 if hasattr(gemma, "lm_head"): 

131 state_dict["unembed.W_U"] = gemma.lm_head.weight.T 

132 unembed_device = gemma.lm_head.weight.device 

133 else: 

134 # Multimodal models might use tied embeddings 

135 state_dict["unembed.W_U"] = base_model.embed_tokens.weight.T 

136 unembed_device = base_model.embed_tokens.weight.device 

137 state_dict["unembed.b_U"] = torch.zeros(cfg.d_vocab, dtype=cfg.dtype, device=unembed_device) 

138 

139 return state_dict