Coverage for transformer_lens/model_bridge/generalized_components/rms_normalization.py: 91%

1"""RMS Normalization bridge component implementation. 

2 

3RMSNorm (Root Mean Square Layer Normalization) is used in models like T5, LLaMA, Mistral, etc. 

4Unlike LayerNorm, RMSNorm doesn't center the inputs (no mean subtraction) and has no bias. 

5""" 

6from __future__ import annotations 

7 

8from typing import TYPE_CHECKING, Any, Dict, Optional 

9 

10from transformer_lens.model_bridge.generalized_components.normalization import ( 

11 NormalizationBridge, 

12) 

13 

14if TYPE_CHECKING: 

15 from transformer_lens.model_bridge.generalized_components.base import ( 

16 GeneralizedComponent, 

17 ) 

18 

19 

20class RMSNormalizationBridge(NormalizationBridge): 

21 """RMS Normalization bridge for models that use RMSNorm (T5, LLaMA, etc). 

22 

23 RMSNorm differs from LayerNorm in two ways: 

24 1. No mean centering (no subtraction of mean) 

25 2. No bias term (only weight/scale parameter) 

26 

27 This bridge does a simple pass-through to the original HuggingFace component 

28 with hooks on input and output. 

29 """ 

30 

31 property_aliases = {"w": "weight"} 

32 

33 def __init__( 

34 self, 

35 name: str, 

36 config: Any, 

37 submodules: Optional[Dict[str, "GeneralizedComponent"]] = None, 

38 use_native_layernorm_autograd: bool = True, 

39 optional: bool = False, 

40 ): 

41 """Initialize the RMS normalization bridge. 

42 

43 Args: 

44 name: The name of this component 

45 config: Configuration object 

46 submodules: Dictionary of GeneralizedComponent submodules to register 

47 use_native_layernorm_autograd: Use HF's RMSNorm implementation for exact numerical match 

48 optional: If True, setup skips this subtree when absent (hybrid architectures) 

49 """ 

50 super().__init__( 

51 name, 

52 config, 

53 submodules=submodules or {}, 

54 use_native_layernorm_autograd=use_native_layernorm_autograd, 

55 optional=optional, 

56 ) 

57 if self.config is not None and (not hasattr(self.config, "uses_rms_norm")): 57 ↛ 58line 57 didn't jump to line 58 because the condition on line 57 was never true

58 self.config.uses_rms_norm = True