Coverage for transformer_lens/model_bridge/supported_architectures/vaultgemma.py: 100%

1"""VaultGemma architecture adapter. 

2 

3Google's VaultGemma (``VaultGemmaForCausalLM``, native in transformers): 

4the only fully DP-SGD-pretrained open LLM — a Gemma-2 recipe trained under 

5differential privacy. Structurally it is Gemma 2 with the two post-norms 

6removed (blocks keep only input_layernorm and pre_feedforward_layernorm); 

7the scaled word embedding and attention logit soft-cap carry over, so this 

8subclasses the Gemma 2 adapter and rebuilds the block entry without 

9ln1_post/ln2_post. 

10""" 

11 

12from transformer_lens.model_bridge.generalized_components import RMSNormalizationBridge 

13from transformer_lens.model_bridge.supported_architectures.gemma2 import ( 

14 Gemma2ArchitectureAdapter, 

15) 

16 

17 

18class VaultGemmaArchitectureAdapter(Gemma2ArchitectureAdapter): 

19 """Architecture adapter for VaultGemmaForCausalLM models.""" 

20 

21 # Compat mode's gemma path assumes the post-norm sandwich this offset-RMS variant 

22 # removed, so its stored-processed-weights forward diverges; P3 is dropped and compat 

23 # is gated off so enable_compatibility_mode() raises rather than silently diverging. 

24 applicable_phases: list[int] = [1, 2, 4] 

25 supports_compatibility_mode: bool = False 

26 

27 def _block_norms(self): 

28 """Gemma 2 minus the post-norms: only input and pre-feedforward norms.""" 

29 return { 

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

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

32 }