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

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 typing import Any 

13 

14from transformer_lens.model_bridge.generalized_components import RMSNormalizationBridge 

15from transformer_lens.model_bridge.supported_architectures.gemma2 import ( 

16 Gemma2ArchitectureAdapter, 

17) 

18 

19 

20class VaultGemmaArchitectureAdapter(Gemma2ArchitectureAdapter): 

21 """Architecture adapter for VaultGemmaForCausalLM models.""" 

22 

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

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

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

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

27 supports_compatibility_mode: bool = False 

28 

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

30 """Initialize the VaultGemma architecture adapter.""" 

31 super().__init__(cfg) 

32 

33 # Gemma 2 minus the post-norms: drop the inherited RMS-offset 

34 # conversions for the ln1_post/ln2_post norms this variant removes. 

35 if self.weight_processing_conversions is not None: 35 ↛ exitline 35 didn't return from function '__init__' because the condition on line 35 was always true

36 for dead in ("blocks.{i}.ln1_post.weight", "blocks.{i}.ln2_post.weight"): 

37 self.weight_processing_conversions.pop(dead, None) 

38 

39 def _block_norms(self): 

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

41 return { 

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

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

44 }