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

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 RMSNormalizationBridge, 

12 RotaryEmbeddingBridge, 

13 UnembeddingBridge, 

14) 

15from transformer_lens.model_bridge.generalized_components.position_embeddings_attention import ( 

16 PositionEmbeddingsAttentionBridge, 

17) 

18 

19logger = logging.getLogger(__name__) 

20 

21 

22class ApertusArchitectureAdapter(ArchitectureAdapter): 

23 """Architecture adapter for Apertus models. 

24 

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

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

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

28 

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

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

31 

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

33 - attention_layernorm (instead of input_layernorm) 

34 - feedforward_layernorm (instead of post_attention_layernorm) 

35 """ 

36 

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

38 """Initialize the Apertus architecture adapter.""" 

39 super().__init__(cfg) 

40 

41 self._set_rms_rotary_defaults(gated=False) 

42 

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

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

45 self.cfg.attn_implementation = "eager" 

46 

47 self.weight_processing_conversions = { 

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

49 **self._qkvo_weight_conversions(), 

50 } 

51 

52 # Set up component mapping 

53 # Apertus uses attention_layernorm / feedforward_layernorm instead of the 

54 # typical input_layernorm / post_attention_layernorm names. 

55 self.component_mapping = { 

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

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

58 "blocks": BlockBridge( 

59 name="model.layers", 

60 submodules={ 

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

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

63 "attn": PositionEmbeddingsAttentionBridge( 

64 name="self_attn", 

65 config=self.cfg, 

66 submodules={ 

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

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

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

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

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

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

73 }, 

74 ), 

75 "mlp": self._ungated_mlp(), 

76 }, 

77 ), 

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

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

80 } 

81 

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

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

84 

85 Transformers v5 uses meta tensors during from_pretrained, but 

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

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

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

89 patch can be removed. 

90 

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

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

93 """ 

94 try: 

95 from transformers.activations import XIELUActivation 

96 except ImportError: 

97 return 

98 

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

100 return 

101 

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

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

104 return 

105 

106 _orig_init = XIELUActivation.__init__ 

107 _orig_forward = XIELUActivation.forward 

108 

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

110 try: 

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

112 except NotImplementedError: 

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

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

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

116 import torch 

117 

118 torch.nn.Module.__init__(self) 

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

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

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

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

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

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

125 self.alpha_p = torch.nn.Parameter( 

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

127 ) 

128 self.alpha_n = torch.nn.Parameter( 

129 torch.log( 

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

131 ).unsqueeze(0) 

132 ) 

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

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

135 self._beta_scalar = None 

136 self._eps_scalar = None 

137 self._xielu_cuda_obj = None 

138 

139 def _patched_forward(self, x): 

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

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

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

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

144 return _orig_forward(self, x) 

145 

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

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

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

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

150 

151 @staticmethod 

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

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

154 import inspect 

155 

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

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

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