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-21 19:27 +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 # setup_component_testing knobs: rotary_emb lives on hf_model.gpt_neox;
35 # eager forcing is intentionally skipped for NeoX.
36 _testing_lm_attr = "gpt_neox"
37 _testing_eager = None
39 def __init__(self, cfg: Any) -> None:
40 """Initialize the NeoX architecture adapter.
42 Args:
43 cfg: The configuration object.
44 """
45 super().__init__(cfg)
47 # Set config variables for weight processing
48 self.cfg.normalization_type = "LN"
49 self.cfg.positional_embedding_type = "rotary"
50 self.cfg.final_rms = False
51 self.cfg.gated_mlp = False
52 self.cfg.attn_only = False
54 # GPTNeoX ships both parallel (Pythia, HF's default) and sequential
55 # variants. Hardcoding parallel drops hook_resid_mid on sequential
56 # checkpoints that genuinely have a post-attention residual.
57 # HF-booted configs carry use_parallel_residual; a caller-supplied
58 # TransformerBridgeConfig only has parallel_attn_mlp, so fall back to
59 # it before defaulting to HF's True.
60 use_parallel_residual = getattr(
61 cfg, "use_parallel_residual", getattr(cfg, "parallel_attn_mlp", True)
62 )
63 self.cfg.parallel_attn_mlp = use_parallel_residual
64 block_cls = ParallelBlockBridge if use_parallel_residual else BlockBridge
66 # NeoX/Pythia models were not trained with BOS tokens
67 self.cfg.default_prepend_bos = False
69 self.weight_processing_conversions = {
70 "blocks.{i}.attn.q": ParamProcessingConversion(
71 tensor_conversion=ChainTensorConversion(
72 [
73 SplitTensorConversion(0, 3),
74 RearrangeTensorConversion(
75 "(head d_head) d_model -> head d_model d_head",
76 head=self.cfg.n_heads,
77 d_head=self.cfg.d_model // self.cfg.n_heads,
78 ),
79 ]
80 ),
81 source_key="gpt_neox.layers.{i}.attention.query_key_value.weight",
82 ),
83 "blocks.{i}.attn.k": ParamProcessingConversion(
84 tensor_conversion=ChainTensorConversion(
85 [
86 SplitTensorConversion(1, 3),
87 RearrangeTensorConversion(
88 "(head d_head) d_model -> head d_model d_head",
89 head=self.cfg.n_heads,
90 d_head=self.cfg.d_model // self.cfg.n_heads,
91 ),
92 ]
93 ),
94 source_key="gpt_neox.layers.{i}.attention.query_key_value.weight",
95 ),
96 "blocks.{i}.attn.v": ParamProcessingConversion(
97 tensor_conversion=ChainTensorConversion(
98 [
99 SplitTensorConversion(2, 3),
100 RearrangeTensorConversion(
101 "(head d_head) d_model -> head d_model d_head",
102 head=self.cfg.n_heads,
103 d_head=self.cfg.d_model // self.cfg.n_heads,
104 ),
105 ]
106 ),
107 source_key="gpt_neox.layers.{i}.attention.query_key_value.weight",
108 ),
109 "blocks.{i}.attn.b_Q": ParamProcessingConversion(
110 tensor_conversion=ChainTensorConversion(
111 [
112 SplitTensorConversion(0, 3),
113 RearrangeTensorConversion(
114 "(head d_head) -> head d_head",
115 head=self.cfg.n_heads,
116 ),
117 ]
118 ),
119 source_key="gpt_neox.layers.{i}.attention.query_key_value.bias",
120 ),
121 "blocks.{i}.attn.b_K": ParamProcessingConversion(
122 tensor_conversion=ChainTensorConversion(
123 [
124 SplitTensorConversion(1, 3),
125 RearrangeTensorConversion(
126 "(head d_head) -> head d_head",
127 head=self.cfg.n_heads,
128 ),
129 ]
130 ),
131 source_key="gpt_neox.layers.{i}.attention.query_key_value.bias",
132 ),
133 "blocks.{i}.attn.b_V": ParamProcessingConversion(
134 tensor_conversion=ChainTensorConversion(
135 [
136 SplitTensorConversion(2, 3),
137 RearrangeTensorConversion(
138 "(head d_head) -> head d_head",
139 head=self.cfg.n_heads,
140 ),
141 ]
142 ),
143 source_key="gpt_neox.layers.{i}.attention.query_key_value.bias",
144 ),
145 "blocks.{i}.attn.o": ParamProcessingConversion(
146 tensor_conversion=RearrangeTensorConversion(
147 "d_model (head d_head) -> head d_head d_model",
148 head=self.cfg.n_heads,
149 d_head=self.cfg.d_model // self.cfg.n_heads,
150 ),
151 source_key="gpt_neox.layers.{i}.attention.dense.weight",
152 ),
153 }
155 self.component_mapping = {
156 "embed": EmbeddingBridge(name="gpt_neox.embed_in"),
157 "rotary_emb": RotaryEmbeddingBridge(name="gpt_neox.rotary_emb"),
158 "blocks": block_cls(
159 name="gpt_neox.layers",
160 submodules={
161 "ln1": NormalizationBridge(
162 name="input_layernorm",
163 config=self.cfg,
164 use_native_layernorm_autograd=True,
165 ),
166 "ln2": NormalizationBridge(
167 name="post_attention_layernorm",
168 config=self.cfg,
169 use_native_layernorm_autograd=True,
170 ),
171 "attn": JointQKVPositionEmbeddingsAttentionBridge(
172 name="attention",
173 config=self.cfg,
174 split_qkv_matrix=self.split_qkv_matrix,
175 requires_attention_mask=True, # GPTNeoX/StableLM requires attention_mask
176 submodules={
177 "qkv": LinearBridge(name="query_key_value"),
178 "o": LinearBridge(name="dense"),
179 },
180 ),
181 "mlp": MLPBridge(
182 name="mlp",
183 submodules={
184 "in": LinearBridge(name="dense_h_to_4h"),
185 "out": LinearBridge(name="dense_4h_to_h"),
186 },
187 ),
188 },
189 ),
190 "ln_final": NormalizationBridge(
191 name="gpt_neox.final_layer_norm",
192 config=self.cfg,
193 use_native_layernorm_autograd=True,
194 ),
195 "unembed": UnembeddingBridge(name="lm_head"),
196 }
198 def prepare_model(self, hf_model: Any) -> None:
199 """Fix up the unembed target once the real HF module tree is available.
201 transformers >= 5.14 renamed ``GPTNeoXForCausalLM.embed_out`` to
202 ``lm_head``; the repo's locked 5.13.0 still exposes ``embed_out``. The
203 component_mapping is built in ``__init__`` before any HF model exists,
204 so the ``lm_head`` default above can't be hasattr-checked until now.
205 """
206 super().prepare_model(hf_model)
207 if not hasattr(hf_model, "lm_head") and hasattr(hf_model, "embed_out"):
208 self.components["unembed"].name = "embed_out"
210 def split_qkv_matrix(
211 self, original_attention_component: Any
212 ) -> tuple[torch.nn.Linear, torch.nn.Linear, torch.nn.Linear]:
213 """Split the QKV matrix into separate linear transformations.
215 GPT-NeoX/StableLM uses an interleaved QKV format where the weights are stored as
216 [Q_h0, K_h0, V_h0, Q_h1, K_h1, V_h1, ...] - i.e., Q, K, V are interleaved per head.
218 The weight shape is [n_heads * 3 * d_head, d_model] and the output is reshaped
219 by HuggingFace as [batch, seq, n_heads, 3*d_head] then split on the last dim.
221 Args:
222 original_attention_component: The original attention layer component
224 Returns:
225 Tuple of nn.Linear modules for Q, K, and V transformations
226 """
227 assert original_attention_component is not None
228 assert original_attention_component.query_key_value is not None
230 qkv_weights = original_attention_component.query_key_value.weight
231 assert isinstance(qkv_weights, torch.Tensor)
233 n_heads = self.cfg.n_heads
234 d_head = self.cfg.d_head
235 d_model = self.cfg.d_model
237 # Weight shape: [n_heads * 3 * d_head, d_model]
238 # Reshape to [n_heads, 3 * d_head, d_model] to access Q, K, V per head
239 W_reshaped = qkv_weights.view(n_heads, 3 * d_head, d_model)
241 # Extract Q, K, V weights for all heads and flatten back
242 W_Q = W_reshaped[:, :d_head, :].reshape(n_heads * d_head, d_model)
243 W_K = W_reshaped[:, d_head : 2 * d_head, :].reshape(n_heads * d_head, d_model)
244 W_V = W_reshaped[:, 2 * d_head :, :].reshape(n_heads * d_head, d_model)
246 # Handle bias - same interleaved format
247 qkv_bias = original_attention_component.query_key_value.bias
248 assert isinstance(qkv_bias, torch.Tensor)
250 # Bias shape: [n_heads * 3 * d_head]
251 # Reshape to [n_heads, 3 * d_head] to access Q, K, V per head
252 b_reshaped = qkv_bias.view(n_heads, 3 * d_head)
253 b_Q = b_reshaped[:, :d_head].reshape(n_heads * d_head)
254 b_K = b_reshaped[:, d_head : 2 * d_head].reshape(n_heads * d_head)
255 b_V = b_reshaped[:, 2 * d_head :].reshape(n_heads * d_head)
257 # Create nn.Linear modules
258 # Weight shape for nn.Linear is [out_features, in_features]
259 W_Q_transformation = torch.nn.Linear(d_model, n_heads * d_head, bias=True)
260 W_Q_transformation.weight = torch.nn.Parameter(W_Q)
261 W_Q_transformation.bias = torch.nn.Parameter(b_Q)
263 W_K_transformation = torch.nn.Linear(d_model, n_heads * d_head, bias=True)
264 W_K_transformation.weight = torch.nn.Parameter(W_K)
265 W_K_transformation.bias = torch.nn.Parameter(b_K)
267 W_V_transformation = torch.nn.Linear(d_model, n_heads * d_head, bias=True)
268 W_V_transformation.weight = torch.nn.Parameter(W_V)
269 W_V_transformation.bias = torch.nn.Parameter(b_V)
271 return W_Q_transformation, W_K_transformation, W_V_transformation