Coverage for transformer_lens/model_bridge/supported_architectures/apertus.py: 93%

1"""Apertus architecture adapter.""" 

2 

3import logging 

4from typing import Any 

5 

6from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter 

7from transformer_lens.model_bridge.generalized_components import ( 

8 BlockBridge, 

9 EmbeddingBridge, 

10 LinearBridge, 

11 MLPBridge, 

12 RMSNormalizationBridge, 

13 RotaryEmbeddingBridge, 

14 UnembeddingBridge, 

15) 

16from transformer_lens.model_bridge.generalized_components.position_embeddings_attention import ( 

17 PositionEmbeddingsAttentionBridge, 

18) 

19 

20logger = logging.getLogger(__name__) 

21 

22 

23class ApertusArchitectureAdapter(ArchitectureAdapter): 

24 """Architecture adapter for Apertus models. 

25 

26 Apertus uses a pre-norm architecture with RMSNorm, Q/K normalization in attention, 

27 rotary position embeddings (RoPE with LLaMA-3 scaling), grouped query attention (GQA), 

28 non-gated MLP (XiELU activation), and no biases on any projections. 

29 

30 Similar to Qwen3 (pre-norm RMSNorm, QK-norm, GQA, RoPE) but uses a non-gated MLP 

31 (up_proj -> XiELU -> down_proj) instead of gated MLP. 

32 

33 Note: Apertus uses different layer norm names than most Llama-family models: 

34 - attention_layernorm (instead of input_layernorm) 

35 - feedforward_layernorm (instead of post_attention_layernorm) 

36 """ 

37 

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

39 """Initialize the Apertus architecture adapter.""" 

40 super().__init__(cfg) 

41 

42 # Set config variables for weight processing 

43 self.cfg.normalization_type = "RMS" 

44 self.cfg.positional_embedding_type = "rotary" 

45 self.cfg.final_rms = True 

46 self.cfg.gated_mlp = False 

47 self.cfg.attn_only = False 

48 self.cfg.uses_rms_norm = True 

49 

50 # Use eager attention to support output_attentions for hook_attn_scores and hook_pattern 

51 # SDPA doesn't support output_attentions, which is required for HookedTransformer compatibility 

52 self.cfg.attn_implementation = "eager" 

53 

54 self.weight_processing_conversions = { 

55 # Q/K/V weight conversions - handle GQA (Grouped Query Attention) 

56 **self._qkvo_weight_conversions(), 

57 } 

58 

59 # Set up component mapping 

60 # Apertus uses attention_layernorm / feedforward_layernorm instead of the 

61 # typical input_layernorm / post_attention_layernorm names. 

62 self.component_mapping = { 

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

64 "rotary_emb": RotaryEmbeddingBridge(name="model.rotary_emb", config=self.cfg), 

65 "blocks": BlockBridge( 

66 name="model.layers", 

67 submodules={ 

68 "ln1": RMSNormalizationBridge(name="attention_layernorm", config=self.cfg), 

69 "ln2": RMSNormalizationBridge(name="feedforward_layernorm", config=self.cfg), 

70 "attn": PositionEmbeddingsAttentionBridge( 

71 name="self_attn", 

72 config=self.cfg, 

73 submodules={ 

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

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

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

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

78 "q_norm": RMSNormalizationBridge(name="q_norm", config=self.cfg), 

79 "k_norm": RMSNormalizationBridge(name="k_norm", config=self.cfg), 

80 }, 

81 ), 

82 "mlp": MLPBridge( 

83 name="mlp", 

84 submodules={ 

85 "in": LinearBridge(name="up_proj"), 

86 "out": LinearBridge(name="down_proj"), 

87 }, 

88 ), 

89 }, 

90 ), 

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

92 "unembed": UnembeddingBridge(name="lm_head"), 

93 } 

94 

95 def prepare_loading(self, model_name: str, model_kwargs: dict) -> None: 

96 """Patch XIELUActivation to defer eager .item() calls for meta tensor compat. 

97 

98 Transformers v5 uses meta tensors during from_pretrained, but 

99 XIELUActivation.__init__ eagerly calls .item() on beta/eps buffers to 

100 precompute _beta_scalar/_eps_scalar for the CUDA kernel path. This fails 

101 on meta device. Once upstream fixes this (transformers PR #43473), this 

102 patch can be removed. 

103 

104 Instead of reimplementing __init__, we wrap it to catch the meta tensor 

105 failure and defer scalar computation to forward() time. 

106 """ 

107 try: 

108 from transformers.activations import XIELUActivation 

109 except ImportError: 

110 return 

111 

112 if getattr(XIELUActivation, "_apertus_patched", False): 

113 return 

114 

115 # Check if upstream already defers scalar computation (fix landed) 

116 if not self._xielu_needs_patch(XIELUActivation): 116 ↛ 117line 116 didn't jump to line 117 because the condition on line 116 was never true

117 return 

118 

119 _orig_init = XIELUActivation.__init__ 

120 _orig_forward = XIELUActivation.forward 

121 

122 def _patched_init(self, *args, **kwargs): 

123 try: 

124 _orig_init(self, *args, **kwargs) 

125 except NotImplementedError: 

126 # Meta device — re-run without the .item() calls 

127 _orig_init.__wrapped_meta = True # type: ignore[attr-defined] 

128 # Call nn.Module.__init__ and replicate only the tensor setup 

129 import torch 

130 

131 torch.nn.Module.__init__(self) 

132 alpha_p_init = kwargs.get("alpha_p_init", 0.8) 

133 alpha_n_init = kwargs.get("alpha_n_init", 0.8) 

134 beta = kwargs.get("beta", 0.5) 

135 eps = kwargs.get("eps", -1e-6) 

136 dtype = kwargs.get("dtype", torch.bfloat16) 

137 self.with_vector_loads = kwargs.get("with_vector_loads", False) 

138 self.alpha_p = torch.nn.Parameter( 

139 torch.log(torch.expm1(torch.tensor(alpha_p_init, dtype=dtype))).unsqueeze(0) 

140 ) 

141 self.alpha_n = torch.nn.Parameter( 

142 torch.log( 

143 torch.expm1(torch.tensor(alpha_n_init - beta, dtype=dtype)) 

144 ).unsqueeze(0) 

145 ) 

146 self.register_buffer("beta", torch.tensor(beta, dtype=dtype)) 

147 self.register_buffer("eps", torch.tensor(eps, dtype=dtype)) 

148 self._beta_scalar = None 

149 self._eps_scalar = None 

150 self._xielu_cuda_obj = None 

151 

152 def _patched_forward(self, x): 

153 """Lazily compute scalars on first real forward pass.""" 

154 if self._beta_scalar is None: 154 ↛ 157line 154 didn't jump to line 157 because the condition on line 154 was always true

155 self._beta_scalar = float(self.beta.detach().cpu().float().item()) 

156 self._eps_scalar = float(self.eps.detach().cpu().float().item()) 

157 return _orig_forward(self, x) 

158 

159 XIELUActivation.__init__ = _patched_init # type: ignore[method-assign] 

160 XIELUActivation.forward = _patched_forward # type: ignore[method-assign] 

161 XIELUActivation._apertus_patched = True # type: ignore[attr-defined] 

162 logger.debug("Patched XIELUActivation for meta tensor compatibility") 

163 

164 @staticmethod 

165 def _xielu_needs_patch(cls: type) -> bool: 

166 """Check whether XIELUActivation still eagerly calls .item() in __init__.""" 

167 import inspect 

168 

169 src = inspect.getsource(cls.__init__) # type: ignore[misc] 

170 # If __init__ still has the eager .item() / float() pattern, patch needed 

171 return "_beta_scalar" in src and ".item()" in src