Coverage for transformer_lens/model_bridge/supported_architectures/phi3.py: 81%
119 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-01 16:23 +0000
1"""Phi-3 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.base_tensor_conversion import (
12 BaseTensorConversion,
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.compat import patch_dynamic_cache_v5
19from transformer_lens.model_bridge.generalized_components import (
20 BlockBridge,
21 EmbeddingBridge,
22 JointGateUpMLPBridge,
23 JointQKVPositionEmbeddingsAttentionBridge,
24 LinearBridge,
25 RMSNormalizationBridge,
26 RotaryEmbeddingBridge,
27 UnembeddingBridge,
28)
29from transformer_lens.utilities.quantization import require_readable_weight
32class _SizedSplitConversion(BaseTensorConversion):
33 """Split a tensor using explicit sizes (for GQA where Q/K/V have different dimensions)."""
35 def __init__(self, sizes: list[int], index: int, dim: int = 0):
36 super().__init__()
37 self.sizes = sizes
38 self.index = index
39 self.dim = dim
41 def handle_conversion(self, input_value: torch.Tensor, *full_context: Any) -> torch.Tensor:
42 parts = torch.split(input_value, self.sizes, dim=self.dim)
43 return parts[self.index]
46class Phi3ArchitectureAdapter(ArchitectureAdapter):
47 """Architecture adapter for Phi-3 models."""
49 _testing_eager = None
51 def __init__(self, cfg: Any) -> None:
52 """Initialize the Phi-3 architecture adapter.
54 Args:
55 cfg: The configuration object.
56 """
57 super().__init__(cfg)
59 self._set_rms_rotary_defaults()
61 # Standard fold_ln can't handle joint qkv/gate_up projections (shape mismatch).
62 # LN folding is handled in preprocess_weights() instead.
63 self.supports_fold_ln = False
65 # GQA: Q has n_heads * d_head, K/V have n_kv_heads * d_head.
66 # cfg.d_head honours an explicit HF head_dim, which need not equal
67 # d_model // n_heads.
68 d_head = cfg.d_head
69 n_kv_heads = cfg.n_key_value_heads or cfg.n_heads
70 q_size = cfg.n_heads * d_head
71 kv_size = n_kv_heads * d_head
72 qkv_sizes = [q_size, kv_size, kv_size]
74 self.weight_processing_conversions = {
75 "blocks.{i}.attn.q": ParamProcessingConversion(
76 tensor_conversion=_SizedSplitConversion(qkv_sizes, 0),
77 source_key="model.layers.{i}.self_attn.qkv_proj.weight",
78 ),
79 "blocks.{i}.attn.k": ParamProcessingConversion(
80 tensor_conversion=_SizedSplitConversion(qkv_sizes, 1),
81 source_key="model.layers.{i}.self_attn.qkv_proj.weight",
82 ),
83 "blocks.{i}.attn.v": ParamProcessingConversion(
84 tensor_conversion=_SizedSplitConversion(qkv_sizes, 2),
85 source_key="model.layers.{i}.self_attn.qkv_proj.weight",
86 ),
87 "blocks.{i}.attn.o": ParamProcessingConversion(
88 tensor_conversion=RearrangeTensorConversion("m (n h) -> n h m", n=self.cfg.n_heads),
89 source_key="model.layers.{i}.self_attn.o_proj.weight",
90 ),
91 "blocks.{i}.mlp.in": ParamProcessingConversion(
92 tensor_conversion=SplitTensorConversion(1, 2),
93 source_key="model.layers.{i}.mlp.gate_up_proj.weight",
94 ),
95 "blocks.{i}.mlp.gate": ParamProcessingConversion(
96 tensor_conversion=SplitTensorConversion(0, 2),
97 source_key="model.layers.{i}.mlp.gate_up_proj.weight",
98 ),
99 }
101 # Set up component mapping
102 self.component_mapping = {
103 "embed": EmbeddingBridge(name="model.embed_tokens"),
104 "rotary_emb": RotaryEmbeddingBridge(name="model.rotary_emb"),
105 "blocks": BlockBridge(
106 name="model.layers",
107 submodules={
108 "ln1": RMSNormalizationBridge(name="input_layernorm", config=self.cfg),
109 "ln2": RMSNormalizationBridge(name="post_attention_layernorm", config=self.cfg),
110 "attn": JointQKVPositionEmbeddingsAttentionBridge(
111 name="self_attn",
112 config=self.cfg,
113 split_qkv_matrix=self._split_phi3_qkv,
114 submodules={
115 "qkv": LinearBridge(name="qkv_proj"),
116 "o": LinearBridge(name="o_proj"),
117 },
118 ),
119 "mlp": JointGateUpMLPBridge(
120 name="mlp",
121 config=self.cfg,
122 split_gate_up_matrix=self._split_gate_up,
123 submodules={
124 "out": LinearBridge(name="down_proj"),
125 },
126 ),
127 },
128 ),
129 "ln_final": RMSNormalizationBridge(name="model.norm", config=self.cfg),
130 "unembed": UnembeddingBridge(name="lm_head"),
131 }
133 @staticmethod
134 def _split_gate_up(
135 original_mlp_component: Any,
136 ) -> tuple[torch.nn.Module, torch.nn.Module]:
137 """Split Phi-3's fused gate_up_proj into separate gate and up Linear modules."""
138 # This override, not the guarded default, is what Phi-3, GLM and GLM-4V
139 # install — so the guard has to be here too. FP8 is the case that needs
140 # it: tensor_split and nn.Parameter both accept it without complaint.
141 fused_weight = require_readable_weight(
142 original_mlp_component.gate_up_proj.weight,
143 operation="split a fused gate/up projection at boot",
144 owner=original_mlp_component.gate_up_proj,
145 )
146 gate_w, up_w = torch.tensor_split(fused_weight, 2, dim=0)
147 d_model = fused_weight.shape[1]
148 d_mlp = gate_w.shape[0]
150 has_bias = (
151 hasattr(original_mlp_component.gate_up_proj, "bias")
152 and original_mlp_component.gate_up_proj.bias is not None
153 )
154 gate_b: torch.Tensor | None
155 up_b: torch.Tensor | None
156 if has_bias: 156 ↛ 157line 156 didn't jump to line 157 because the condition on line 156 was never true
157 gate_b, up_b = torch.tensor_split(original_mlp_component.gate_up_proj.bias, 2, dim=0)
158 else:
159 gate_b = up_b = None
161 gate_proj = torch.nn.Linear(d_model, d_mlp, bias=has_bias)
162 gate_proj.weight = torch.nn.Parameter(gate_w)
163 if gate_b is not None: 163 ↛ 164line 163 didn't jump to line 164 because the condition on line 163 was never true
164 gate_proj.bias = torch.nn.Parameter(gate_b)
166 up_proj = torch.nn.Linear(d_model, d_mlp, bias=has_bias)
167 up_proj.weight = torch.nn.Parameter(up_w)
168 if up_b is not None: 168 ↛ 169line 168 didn't jump to line 169 because the condition on line 168 was never true
169 up_proj.bias = torch.nn.Parameter(up_b)
171 return gate_proj, up_proj
173 def _split_phi3_qkv(
174 self, original_attention_component: Any
175 ) -> tuple[torch.nn.Module, torch.nn.Module, torch.nn.Module]:
176 """Split Phi-3's fused qkv_proj into separate Q, K, V linear modules."""
177 qkv_weight = require_readable_weight(
178 original_attention_component.qkv_proj.weight,
179 operation="split a fused QKV projection at boot",
180 owner=original_attention_component.qkv_proj,
181 )
182 d_model = qkv_weight.shape[1]
184 # GQA: Q has n_heads * d_head, K/V have n_kv_heads * d_head each.
185 # cfg.d_head honours an explicit HF head_dim, which need not equal
186 # d_model // n_heads.
187 d_head = self.cfg.d_head
188 n_kv_heads = self.cfg.n_key_value_heads or self.cfg.n_heads
189 q_size = self.cfg.n_heads * d_head
190 kv_size = n_kv_heads * d_head
191 q_weight, k_weight, v_weight = torch.split(qkv_weight, [q_size, kv_size, kv_size], dim=0)
193 has_bias = (
194 hasattr(original_attention_component.qkv_proj, "bias")
195 and original_attention_component.qkv_proj.bias is not None
196 )
197 q_bias: torch.Tensor | None
198 k_bias: torch.Tensor | None
199 v_bias: torch.Tensor | None
200 if has_bias:
201 q_bias, k_bias, v_bias = torch.split(
202 original_attention_component.qkv_proj.bias, [q_size, kv_size, kv_size], dim=0
203 )
204 else:
205 q_bias = k_bias = v_bias = None
207 q_linear = torch.nn.Linear(d_model, q_weight.shape[0], bias=has_bias)
208 q_linear.weight = torch.nn.Parameter(q_weight)
209 if q_bias is not None:
210 q_linear.bias = torch.nn.Parameter(q_bias)
212 k_linear = torch.nn.Linear(d_model, k_weight.shape[0], bias=has_bias)
213 k_linear.weight = torch.nn.Parameter(k_weight)
214 if k_bias is not None:
215 k_linear.bias = torch.nn.Parameter(k_bias)
217 v_linear = torch.nn.Linear(d_model, v_weight.shape[0], bias=has_bias)
218 v_linear.weight = torch.nn.Parameter(v_weight)
219 if v_bias is not None:
220 v_linear.bias = torch.nn.Parameter(v_bias)
222 return q_linear, k_linear, v_linear
224 def prepare_loading(self, model_name: str, model_kwargs: dict) -> None:
225 """Patch cached Phi-3 remote code for transformers v5 compatibility."""
226 uses_remote_code = model_kwargs.get("trust_remote_code", False)
227 if not uses_remote_code:
228 return
230 config = model_kwargs.get("config")
231 if config is not None:
232 rope_scaling = getattr(config, "rope_scaling", None)
233 if isinstance(rope_scaling, dict) and rope_scaling.get("rope_type") == "default":
234 config.rope_scaling = None
236 patch_dynamic_cache_v5()
238 def preprocess_weights(self, state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
239 """Fold layer norms into joint QKV/gate_up projections.
241 Standard fold_ln can't handle joint projections (shape mismatch on round-trip),
242 so we scale the full joint weights directly.
243 """
244 fold_ln = getattr(self, "_fold_ln_requested", True)
245 if not fold_ln:
246 return state_dict
248 n_layers = self.cfg.n_layers
250 for i in range(n_layers):
251 ln1_key = f"blocks.{i}.ln1.weight"
252 ln2_key = f"blocks.{i}.ln2.weight"
254 # Fold ln1 into qkv_proj
255 if ln1_key in state_dict: 255 ↛ 270line 255 didn't jump to line 270 because the condition on line 255 was always true
256 ln1_w = state_dict[ln1_key].float()
257 for qkv_key in [
258 f"blocks.{i}.attn.q.weight",
259 f"blocks.{i}.attn.k.weight",
260 f"blocks.{i}.attn.v.weight",
261 ]:
262 if qkv_key in state_dict: 262 ↛ 257line 262 didn't jump to line 257 because the condition on line 262 was always true
263 orig_dtype = state_dict[qkv_key].dtype
264 state_dict[qkv_key] = (state_dict[qkv_key].float() * ln1_w[None, :]).to(
265 orig_dtype
266 )
267 state_dict[ln1_key] = torch.ones_like(state_dict[ln1_key])
269 # Fold ln2 into gate_up_proj
270 if ln2_key in state_dict: 270 ↛ 250line 270 didn't jump to line 250 because the condition on line 270 was always true
271 ln2_w = state_dict[ln2_key].float()
272 for mlp_key in [
273 f"blocks.{i}.mlp.gate.weight",
274 f"blocks.{i}.mlp.in.weight",
275 ]:
276 if mlp_key in state_dict: 276 ↛ 272line 276 didn't jump to line 272 because the condition on line 276 was always true
277 orig_dtype = state_dict[mlp_key].dtype
278 state_dict[mlp_key] = (state_dict[mlp_key].float() * ln2_w[None, :]).to(
279 orig_dtype
280 )
281 state_dict[ln2_key] = torch.ones_like(state_dict[ln2_key])
283 # Fold ln_final into unembed
284 ln_final_key = "ln_final.weight"
285 unembed_key = "unembed.weight"
286 if ln_final_key in state_dict and unembed_key in state_dict: 286 ↛ 296line 286 didn't jump to line 296 because the condition on line 286 was always true
287 ln_final_w = state_dict[ln_final_key].float()
288 unembed_w = state_dict[unembed_key].float()
289 orig_dtype = state_dict[unembed_key].dtype
290 if unembed_w.shape[-1] == ln_final_w.shape[0]: 290 ↛ 292line 290 didn't jump to line 292 because the condition on line 290 was always true
291 state_dict[unembed_key] = (unembed_w * ln_final_w[None, :]).to(orig_dtype)
292 elif unembed_w.shape[0] == ln_final_w.shape[0]:
293 state_dict[unembed_key] = (unembed_w * ln_final_w[:, None]).to(orig_dtype)
294 state_dict[ln_final_key] = torch.ones_like(state_dict[ln_final_key])
296 return state_dict