Coverage for transformer_lens/model_bridge/supported_architectures/deepseek_v4.py: 79%
77 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-08-11 18:50 +0000
1"""DeepSeek V4 architecture adapter.
3DeepSeek V4 replaces V2/V3's MLA path with a hybrid local/compressed attention
4stack and keeps ``hc_mult`` residual streams alive between blocks through
5manifold-constrained Hyper-Connections (mHC). The adapter delegates those
6architecture-specific calculations to Transformers while exposing the modules
7that are useful for interpretability: mHC collapse/mix tensors, compressed KV
8states and masks, Lightning Indexer selections, attention projections, and MoE
9routing/expert outputs.
10"""
12from typing import Any, Dict, Optional
14import torch
16from transformer_lens.hook_points import HookPoint
17from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter
18from transformer_lens.model_bridge.generalized_components import (
19 BlockBridge,
20 EmbeddingBridge,
21 GatedMLPBridge,
22 LinearBridge,
23 MoEBridge,
24 RMSNormalizationBridge,
25 RotaryEmbeddingBridge,
26 UnembeddingBridge,
27)
28from transformer_lens.model_bridge.generalized_components.base import (
29 GeneralizedComponent,
30)
33class DeepseekV4HyperConnectionBridge(GeneralizedComponent):
34 """Bridge an mHC module without discarding its three distinct outputs.
36 ``hook_in`` sees the full ``[batch, pos, hc_mult, d_model]`` residual stack.
37 ``hook_post`` and ``hook_comb`` expose the learned expansion and stream-mix
38 weights, while ``hook_out`` exposes the collapsed conventional residual that
39 enters attention or the MLP.
40 """
42 def __init__(
43 self,
44 name: str,
45 config: Optional[Any] = None,
46 submodules: Optional[Dict[str, GeneralizedComponent]] = None,
47 ) -> None:
48 super().__init__(name, config, submodules=submodules or {})
49 self.hook_post = HookPoint()
50 self.hook_comb = HookPoint()
52 def forward(self, *args: Any, **kwargs: Any) -> Any:
53 """Run the native mHC module and hook each returned tensor separately."""
54 if self.original_component is None: 54 ↛ 55line 54 didn't jump to line 55 because the condition on line 54 was never true
55 raise RuntimeError(
56 f"Original component not set for {self.name}. Call set_original_component() first."
57 )
59 if args and isinstance(args[0], torch.Tensor): 59 ↛ 61line 59 didn't jump to line 61 because the condition on line 59 was always true
60 args = (self.hook_in(args[0]),) + args[1:]
61 elif isinstance(kwargs.get("hidden_streams"), torch.Tensor):
62 kwargs["hidden_streams"] = self.hook_in(kwargs["hidden_streams"])
64 output = self.original_component(*args, **kwargs)
65 if not isinstance(output, tuple) or len(output) != 3: 65 ↛ 66line 65 didn't jump to line 66 because the condition on line 65 was never true
66 raise RuntimeError(
67 f"DeepSeek V4 hyper-connection {self.name} returned an unexpected output"
68 )
70 post, comb, collapsed = output
71 return self.hook_post(post), self.hook_comb(comb), self.hook_out(collapsed)
74class DeepseekV4CompressorBridge(GeneralizedComponent):
75 """Bridge CSA/HCA compression and expose compressed KV plus block bias."""
77 def __init__(
78 self,
79 name: str,
80 config: Optional[Any] = None,
81 submodules: Optional[Dict[str, GeneralizedComponent]] = None,
82 optional: bool = False,
83 ) -> None:
84 super().__init__(
85 name,
86 config,
87 submodules=submodules or {},
88 optional=optional,
89 )
90 self.hook_block_bias = HookPoint()
92 def forward(self, *args: Any, **kwargs: Any) -> Any:
93 """Run the native compressor, preserving and hooking both outputs."""
94 if self.original_component is None: 94 ↛ 95line 94 didn't jump to line 95 because the condition on line 94 was never true
95 raise RuntimeError(
96 f"Original component not set for {self.name}. Call set_original_component() first."
97 )
99 if args and isinstance(args[0], torch.Tensor): 99 ↛ 101line 99 didn't jump to line 101 because the condition on line 99 was always true
100 args = (self.hook_in(args[0]),) + args[1:]
101 elif isinstance(kwargs.get("hidden_states"), torch.Tensor):
102 kwargs["hidden_states"] = self.hook_in(kwargs["hidden_states"])
104 output = self.original_component(*args, **kwargs)
105 if not isinstance(output, tuple) or len(output) != 2: 105 ↛ 106line 105 didn't jump to line 106 because the condition on line 105 was never true
106 raise RuntimeError(f"DeepSeek V4 compressor {self.name} returned an unexpected output")
108 compressed_kv, block_bias = output
109 compressed_kv = self.hook_out(compressed_kv)
110 if isinstance(block_bias, torch.Tensor): 110 ↛ 112line 110 didn't jump to line 112 because the condition on line 110 was always true
111 block_bias = self.hook_block_bias(block_bias)
112 return compressed_kv, block_bias
115class DeepseekV4BlockBridge(BlockBridge):
116 """Block bridge whose input/output hooks carry the full mHC stream stack.
118 Standard residual aliases are intentionally omitted: V4's block boundary is
119 four-dimensional, and presenting it as a conventional single residual stream
120 would make otherwise-valid patching code silently target the wrong tensor.
121 The collapsed attention/MLP inputs are available at ``attn_hc.hook_out`` and
122 ``mlp_hc.hook_out`` respectively.
123 """
125 hook_aliases: dict[str, str | list[str]] = {}
126 hook_out_is_single_residual_stream: bool = False
127 maintain_native_attention: bool = True
130def _compressor_bridge(cfg: Any) -> DeepseekV4CompressorBridge:
131 """Build the common CSA/HCA compressor mapping, including optional indexer."""
132 return DeepseekV4CompressorBridge(
133 name="compressor",
134 config=cfg,
135 optional=True,
136 submodules={
137 "kv_proj": LinearBridge(name="kv_proj"),
138 "gate_proj": LinearBridge(name="gate_proj"),
139 "kv_norm": RMSNormalizationBridge(name="kv_norm", config=cfg),
140 "rotary_emb": RotaryEmbeddingBridge(name="rotary_emb", config=cfg),
141 "indexer": GeneralizedComponent(
142 name="indexer",
143 optional=True,
144 submodules={
145 "kv_proj": LinearBridge(name="kv_proj"),
146 "gate_proj": LinearBridge(name="gate_proj"),
147 "kv_norm": RMSNormalizationBridge(name="kv_norm", config=cfg),
148 "q_b_proj": LinearBridge(name="q_b_proj"),
149 "scorer": GeneralizedComponent(
150 name="scorer",
151 submodules={
152 "weights_proj": LinearBridge(name="weights_proj"),
153 },
154 ),
155 "rotary_emb": RotaryEmbeddingBridge(name="rotary_emb", config=cfg),
156 },
157 ),
158 },
159 )
162class DeepSeekV4ArchitectureAdapter(ArchitectureAdapter):
163 """Adapter for ``DeepseekV4ForCausalLM`` (Flash and Pro variants)."""
165 # The isolated component harness assumes a three-dimensional residual. V4's
166 # mHC stack is four-dimensional, so parity is covered by integration tests and
167 # verify_models' whole-model hook/text phases instead of isolated Phase 1.
168 applicable_phases: list[int] = [2, 4]
170 def __init__(self, cfg: Any) -> None:
171 super().__init__(cfg)
173 self.cfg.normalization_type = "RMS"
174 self.cfg.uses_rms_norm = True
175 self.cfg.final_rms = True
176 self.cfg.rmsnorm_uses_offset = False
177 self.cfg.positional_embedding_type = "rotary"
178 self.cfg.gated_mlp = True
179 self.cfg.attn_implementation = "eager"
181 # Folding/centering assumes one additive residual stream. Applying either
182 # transform to mHC's learned collapse/expand path is not basis preserving.
183 self.supports_fold_ln = False
184 self.supports_center_writing_weights = False
185 self.weight_processing_conversions = {}
187 def hyper_connection(name: str) -> DeepseekV4HyperConnectionBridge:
188 return DeepseekV4HyperConnectionBridge(
189 name=name,
190 config=self.cfg,
191 submodules={
192 "input_norm": GeneralizedComponent(name="input_norm"),
193 },
194 )
196 attention = GeneralizedComponent(
197 name="self_attn",
198 submodules={
199 "q_a_proj": LinearBridge(name="q_a_proj"),
200 "q_a_norm": RMSNormalizationBridge(name="q_a_norm", config=self.cfg),
201 "q_b_proj": LinearBridge(name="q_b_proj"),
202 "q_b_norm": GeneralizedComponent(name="q_b_norm"),
203 "kv_proj": LinearBridge(name="kv_proj"),
204 "kv_norm": RMSNormalizationBridge(name="kv_norm", config=self.cfg),
205 "compressor": _compressor_bridge(self.cfg),
206 "o_a_proj": GeneralizedComponent(name="o_a_proj"),
207 "o_b_proj": LinearBridge(name="o_b_proj"),
208 },
209 )
211 mlp = MoEBridge(
212 name="mlp",
213 config=self.cfg,
214 submodules={
215 "gate": GeneralizedComponent(name="gate"),
216 "experts": GeneralizedComponent(name="experts"),
217 "shared_experts": GatedMLPBridge(
218 name="shared_experts",
219 config=self.cfg,
220 submodules={
221 "gate": LinearBridge(name="gate_proj"),
222 "in": LinearBridge(name="up_proj"),
223 "out": LinearBridge(name="down_proj"),
224 },
225 ),
226 },
227 )
229 self.component_mapping = {
230 "embed": EmbeddingBridge(name="model.embed_tokens"),
231 "rotary_emb": RotaryEmbeddingBridge(name="model.rotary_emb", config=self.cfg),
232 "blocks": DeepseekV4BlockBridge(
233 name="model.layers",
234 config=self.cfg,
235 submodules={
236 "attn_hc": hyper_connection("attn_hc"),
237 "ln1": RMSNormalizationBridge(name="input_layernorm", config=self.cfg),
238 "attn": attention,
239 "mlp_hc": hyper_connection("ffn_hc"),
240 "ln2": RMSNormalizationBridge(name="post_attention_layernorm", config=self.cfg),
241 "mlp": mlp,
242 },
243 ),
244 "hc_head": GeneralizedComponent(
245 name="model.hc_head",
246 submodules={
247 "input_norm": GeneralizedComponent(name="input_norm"),
248 },
249 ),
250 "ln_final": RMSNormalizationBridge(name="model.norm", config=self.cfg),
251 "unembed": UnembeddingBridge(name="lm_head"),
252 }
254 def prepare_loading(self, model_name: str, model_kwargs: dict) -> None:
255 """Force eager attention so the delegated attention path is deterministic."""
256 model_kwargs["attn_implementation"] = "eager"
258 def prepare_model(self, hf_model: Any) -> None:
259 """Force eager attention on a pre-loaded model before installing bridges."""
260 if hasattr(hf_model, "config"): 260 ↛ 262line 260 didn't jump to line 262 because the condition on line 260 was always true
261 hf_model.config._attn_implementation = "eager"
262 model = getattr(hf_model, "model", None)
263 if model is not None and hasattr(model, "layers"): 263 ↛ exitline 263 didn't return from function 'prepare_model' because the condition on line 263 was always true
264 for layer in model.layers:
265 if hasattr(layer, "self_attn") and hasattr(layer.self_attn, "config"): 265 ↛ 264line 265 didn't jump to line 264 because the condition on line 265 was always true
266 layer.self_attn.config._attn_implementation = "eager"