Coverage for transformer_lens/model_bridge/supported_architectures/neox.py: 100%
55 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-01 16:23 +0000
1"""NeoX architecture adapter."""
3from typing import Any
5import torch
7from transformer_lens.conversion_utils.conversion_steps import (
8 RearrangeTensorConversion,
9 SplitTensorConversion,
10)
11from transformer_lens.conversion_utils.conversion_steps.chain_tensor_conversion import (
12 ChainTensorConversion,
13)
14from transformer_lens.conversion_utils.param_processing_conversion import (
15 ParamProcessingConversion,
16)
17from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter
18from transformer_lens.model_bridge.generalized_components import (
19 BlockBridge,
20 EmbeddingBridge,
21 JointQKVPositionEmbeddingsAttentionBridge,
22 LinearBridge,
23 MLPBridge,
24 NormalizationBridge,
25 ParallelBlockBridge,
26 RotaryEmbeddingBridge,
27 UnembeddingBridge,
28)
31class NeoxArchitectureAdapter(ArchitectureAdapter):
32 """Architecture adapter for NeoX models."""
34 def __init__(self, cfg: Any) -> None:
35 """Initialize the NeoX architecture adapter.
37 Args:
38 cfg: The configuration object.
39 """
40 super().__init__(cfg)
42 # Set config variables for weight processing
43 self.cfg.normalization_type = "LN"
44 self.cfg.positional_embedding_type = "rotary"
45 self.cfg.final_rms = False
46 self.cfg.gated_mlp = False
47 self.cfg.attn_only = False
49 # GPTNeoX ships both parallel (Pythia, HF's default) and sequential
50 # variants. Hardcoding parallel drops hook_resid_mid on sequential
51 # checkpoints that genuinely have a post-attention residual.
52 # HF-booted configs carry use_parallel_residual; a caller-supplied
53 # TransformerBridgeConfig only has parallel_attn_mlp, so fall back to
54 # it before defaulting to HF's True.
55 use_parallel_residual = getattr(
56 cfg, "use_parallel_residual", getattr(cfg, "parallel_attn_mlp", True)
57 )
58 self.cfg.parallel_attn_mlp = use_parallel_residual
59 block_cls = ParallelBlockBridge if use_parallel_residual else BlockBridge
61 # NeoX/Pythia models were not trained with BOS tokens
62 self.cfg.default_prepend_bos = False
64 self.weight_processing_conversions = {
65 "blocks.{i}.attn.q": ParamProcessingConversion(
66 tensor_conversion=ChainTensorConversion(
67 [
68 SplitTensorConversion(0, 3),
69 RearrangeTensorConversion(
70 "(head d_head) d_model -> head d_model d_head",
71 head=self.cfg.n_heads,
72 d_head=self.cfg.d_model // self.cfg.n_heads,
73 ),
74 ]
75 ),
76 source_key="gpt_neox.layers.{i}.attention.query_key_value.weight",
77 ),
78 "blocks.{i}.attn.k": ParamProcessingConversion(
79 tensor_conversion=ChainTensorConversion(
80 [
81 SplitTensorConversion(1, 3),
82 RearrangeTensorConversion(
83 "(head d_head) d_model -> head d_model d_head",
84 head=self.cfg.n_heads,
85 d_head=self.cfg.d_model // self.cfg.n_heads,
86 ),
87 ]
88 ),
89 source_key="gpt_neox.layers.{i}.attention.query_key_value.weight",
90 ),
91 "blocks.{i}.attn.v": ParamProcessingConversion(
92 tensor_conversion=ChainTensorConversion(
93 [
94 SplitTensorConversion(2, 3),
95 RearrangeTensorConversion(
96 "(head d_head) d_model -> head d_model d_head",
97 head=self.cfg.n_heads,
98 d_head=self.cfg.d_model // self.cfg.n_heads,
99 ),
100 ]
101 ),
102 source_key="gpt_neox.layers.{i}.attention.query_key_value.weight",
103 ),
104 "blocks.{i}.attn.b_Q": ParamProcessingConversion(
105 tensor_conversion=ChainTensorConversion(
106 [
107 SplitTensorConversion(0, 3),
108 RearrangeTensorConversion(
109 "(head d_head) -> head d_head",
110 head=self.cfg.n_heads,
111 ),
112 ]
113 ),
114 source_key="gpt_neox.layers.{i}.attention.query_key_value.bias",
115 ),
116 "blocks.{i}.attn.b_K": ParamProcessingConversion(
117 tensor_conversion=ChainTensorConversion(
118 [
119 SplitTensorConversion(1, 3),
120 RearrangeTensorConversion(
121 "(head d_head) -> head d_head",
122 head=self.cfg.n_heads,
123 ),
124 ]
125 ),
126 source_key="gpt_neox.layers.{i}.attention.query_key_value.bias",
127 ),
128 "blocks.{i}.attn.b_V": ParamProcessingConversion(
129 tensor_conversion=ChainTensorConversion(
130 [
131 SplitTensorConversion(2, 3),
132 RearrangeTensorConversion(
133 "(head d_head) -> head d_head",
134 head=self.cfg.n_heads,
135 ),
136 ]
137 ),
138 source_key="gpt_neox.layers.{i}.attention.query_key_value.bias",
139 ),
140 "blocks.{i}.attn.o": ParamProcessingConversion(
141 tensor_conversion=RearrangeTensorConversion(
142 "d_model (head d_head) -> head d_head d_model",
143 head=self.cfg.n_heads,
144 d_head=self.cfg.d_model // self.cfg.n_heads,
145 ),
146 source_key="gpt_neox.layers.{i}.attention.dense.weight",
147 ),
148 }
150 self.component_mapping = {
151 "embed": EmbeddingBridge(name="gpt_neox.embed_in"),
152 "rotary_emb": RotaryEmbeddingBridge(name="gpt_neox.rotary_emb"),
153 "blocks": block_cls(
154 name="gpt_neox.layers",
155 submodules={
156 "ln1": NormalizationBridge(
157 name="input_layernorm",
158 config=self.cfg,
159 use_native_layernorm_autograd=True,
160 ),
161 "ln2": NormalizationBridge(
162 name="post_attention_layernorm",
163 config=self.cfg,
164 use_native_layernorm_autograd=True,
165 ),
166 "attn": JointQKVPositionEmbeddingsAttentionBridge(
167 name="attention",
168 config=self.cfg,
169 split_qkv_matrix=self.split_qkv_matrix,
170 requires_attention_mask=True, # GPTNeoX/StableLM requires attention_mask
171 submodules={
172 "qkv": LinearBridge(name="query_key_value"),
173 "o": LinearBridge(name="dense"),
174 },
175 ),
176 "mlp": MLPBridge(
177 name="mlp",
178 submodules={
179 "in": LinearBridge(name="dense_h_to_4h"),
180 "out": LinearBridge(name="dense_4h_to_h"),
181 },
182 ),
183 },
184 ),
185 "ln_final": NormalizationBridge(
186 name="gpt_neox.final_layer_norm",
187 config=self.cfg,
188 use_native_layernorm_autograd=True,
189 ),
190 "unembed": UnembeddingBridge(name="embed_out"),
191 }
193 def split_qkv_matrix(
194 self, original_attention_component: Any
195 ) -> tuple[torch.nn.Linear, torch.nn.Linear, torch.nn.Linear]:
196 """Split the QKV matrix into separate linear transformations.
198 GPT-NeoX/StableLM uses an interleaved QKV format where the weights are stored as
199 [Q_h0, K_h0, V_h0, Q_h1, K_h1, V_h1, ...] - i.e., Q, K, V are interleaved per head.
201 The weight shape is [n_heads * 3 * d_head, d_model] and the output is reshaped
202 by HuggingFace as [batch, seq, n_heads, 3*d_head] then split on the last dim.
204 Args:
205 original_attention_component: The original attention layer component
207 Returns:
208 Tuple of nn.Linear modules for Q, K, and V transformations
209 """
210 assert original_attention_component is not None
211 assert original_attention_component.query_key_value is not None
213 qkv_weights = original_attention_component.query_key_value.weight
214 assert isinstance(qkv_weights, torch.Tensor)
216 n_heads = self.cfg.n_heads
217 d_head = self.cfg.d_head
218 d_model = self.cfg.d_model
220 # Weight shape: [n_heads * 3 * d_head, d_model]
221 # Reshape to [n_heads, 3 * d_head, d_model] to access Q, K, V per head
222 W_reshaped = qkv_weights.view(n_heads, 3 * d_head, d_model)
224 # Extract Q, K, V weights for all heads and flatten back
225 W_Q = W_reshaped[:, :d_head, :].reshape(n_heads * d_head, d_model)
226 W_K = W_reshaped[:, d_head : 2 * d_head, :].reshape(n_heads * d_head, d_model)
227 W_V = W_reshaped[:, 2 * d_head :, :].reshape(n_heads * d_head, d_model)
229 # Handle bias - same interleaved format
230 qkv_bias = original_attention_component.query_key_value.bias
231 assert isinstance(qkv_bias, torch.Tensor)
233 # Bias shape: [n_heads * 3 * d_head]
234 # Reshape to [n_heads, 3 * d_head] to access Q, K, V per head
235 b_reshaped = qkv_bias.view(n_heads, 3 * d_head)
236 b_Q = b_reshaped[:, :d_head].reshape(n_heads * d_head)
237 b_K = b_reshaped[:, d_head : 2 * d_head].reshape(n_heads * d_head)
238 b_V = b_reshaped[:, 2 * d_head :].reshape(n_heads * d_head)
240 # Create nn.Linear modules
241 # Weight shape for nn.Linear is [out_features, in_features]
242 W_Q_transformation = torch.nn.Linear(d_model, n_heads * d_head, bias=True)
243 W_Q_transformation.weight = torch.nn.Parameter(W_Q)
244 W_Q_transformation.bias = torch.nn.Parameter(b_Q)
246 W_K_transformation = torch.nn.Linear(d_model, n_heads * d_head, bias=True)
247 W_K_transformation.weight = torch.nn.Parameter(W_K)
248 W_K_transformation.bias = torch.nn.Parameter(b_K)
250 W_V_transformation = torch.nn.Linear(d_model, n_heads * d_head, bias=True)
251 W_V_transformation.weight = torch.nn.Parameter(W_V)
252 W_V_transformation.bias = torch.nn.Parameter(b_V)
254 return W_Q_transformation, W_K_transformation, W_V_transformation
256 def setup_component_testing(self, hf_model: Any, bridge_model: Any = None) -> None:
257 """Set up rotary embedding references for GPT-NeoX/StableLM component testing.
259 GPT-NeoX models use RoPE (Rotary Position Embeddings) which need to be
260 set on all attention bridge instances for component testing.
262 Args:
263 hf_model: The HuggingFace GPT-NeoX model instance
264 bridge_model: The TransformerBridge model (if available, set rotary_emb on actual instances)
265 """
266 # Get rotary embedding instance from model level
267 # In GPT-NeoX/StableLM, rotary_emb is at the model level
268 rotary_emb = hf_model.gpt_neox.rotary_emb
270 # Set rotary_emb on actual bridge instances in bridge_model if available
271 if bridge_model is not None and hasattr(bridge_model, "blocks"):
272 # Set on each layer's actual attention bridge instance
273 for block in bridge_model.blocks:
274 if hasattr(block, "attn"):
275 block.attn.set_rotary_emb(rotary_emb)