Coverage for transformer_lens/model_bridge/supported_architectures/deepseek_v4.py: 78%
72 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-21 19:27 +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)
31from transformer_lens.utilities.attn_implementation import force_eager_attention
34class DeepseekV4HyperConnectionBridge(GeneralizedComponent):
35 """Bridge an mHC module without discarding its three distinct outputs.
37 ``hook_in`` sees the full ``[batch, pos, hc_mult, d_model]`` residual stack.
38 ``hook_post`` and ``hook_comb`` expose the learned expansion and stream-mix
39 weights, while ``hook_out`` exposes the collapsed conventional residual that
40 enters attention or the MLP.
41 """
43 def __init__(
44 self,
45 name: str,
46 config: Optional[Any] = None,
47 submodules: Optional[Dict[str, GeneralizedComponent]] = None,
48 ) -> None:
49 super().__init__(name, config, submodules=submodules or {})
50 self.hook_post = HookPoint()
51 self.hook_comb = HookPoint()
53 def forward(self, *args: Any, **kwargs: Any) -> Any:
54 """Run the native mHC module and hook each returned tensor separately."""
55 if self.original_component is None: 55 ↛ 56line 55 didn't jump to line 56 because the condition on line 55 was never true
56 raise RuntimeError(
57 f"Original component not set for {self.name}. Call set_original_component() first."
58 )
60 if args and isinstance(args[0], torch.Tensor): 60 ↛ 62line 60 didn't jump to line 62 because the condition on line 60 was always true
61 args = (self.hook_in(args[0]),) + args[1:]
62 elif isinstance(kwargs.get("hidden_streams"), torch.Tensor):
63 kwargs["hidden_streams"] = self.hook_in(kwargs["hidden_streams"])
65 output = self.original_component(*args, **kwargs)
66 if not isinstance(output, tuple) or len(output) != 3: 66 ↛ 67line 66 didn't jump to line 67 because the condition on line 66 was never true
67 raise RuntimeError(
68 f"DeepSeek V4 hyper-connection {self.name} returned an unexpected output"
69 )
71 post, comb, collapsed = output
72 return self.hook_post(post), self.hook_comb(comb), self.hook_out(collapsed)
75class DeepseekV4CompressorBridge(GeneralizedComponent):
76 """Bridge CSA/HCA compression and expose compressed KV plus block bias."""
78 def __init__(
79 self,
80 name: str,
81 config: Optional[Any] = None,
82 submodules: Optional[Dict[str, GeneralizedComponent]] = None,
83 optional: bool = False,
84 ) -> None:
85 super().__init__(
86 name,
87 config,
88 submodules=submodules or {},
89 optional=optional,
90 )
91 self.hook_block_bias = HookPoint()
93 def forward(self, *args: Any, **kwargs: Any) -> Any:
94 """Run the native compressor, preserving and hooking both outputs."""
95 if self.original_component is None: 95 ↛ 96line 95 didn't jump to line 96 because the condition on line 95 was never true
96 raise RuntimeError(
97 f"Original component not set for {self.name}. Call set_original_component() first."
98 )
100 if args and isinstance(args[0], torch.Tensor): 100 ↛ 102line 100 didn't jump to line 102 because the condition on line 100 was always true
101 args = (self.hook_in(args[0]),) + args[1:]
102 elif isinstance(kwargs.get("hidden_states"), torch.Tensor):
103 kwargs["hidden_states"] = self.hook_in(kwargs["hidden_states"])
105 output = self.original_component(*args, **kwargs)
106 if not isinstance(output, tuple) or len(output) != 2: 106 ↛ 107line 106 didn't jump to line 107 because the condition on line 106 was never true
107 raise RuntimeError(f"DeepSeek V4 compressor {self.name} returned an unexpected output")
109 compressed_kv, block_bias = output
110 compressed_kv = self.hook_out(compressed_kv)
111 if isinstance(block_bias, torch.Tensor): 111 ↛ 113line 111 didn't jump to line 113 because the condition on line 111 was always true
112 block_bias = self.hook_block_bias(block_bias)
113 return compressed_kv, block_bias
116class DeepseekV4BlockBridge(BlockBridge):
117 """Block bridge whose input/output hooks carry the full mHC stream stack.
119 Standard residual aliases are intentionally omitted: V4's block boundary is
120 four-dimensional, and presenting it as a conventional single residual stream
121 would make otherwise-valid patching code silently target the wrong tensor.
122 The collapsed attention/MLP inputs are available at ``attn_hc.hook_out`` and
123 ``mlp_hc.hook_out`` respectively.
124 """
126 hook_aliases: dict[str, str | list[str]] = {}
127 hook_out_is_single_residual_stream: bool = False
128 maintain_native_attention: bool = True
131def _compressor_bridge(cfg: Any) -> DeepseekV4CompressorBridge:
132 """Build the common CSA/HCA compressor mapping, including optional indexer."""
133 return DeepseekV4CompressorBridge(
134 name="compressor",
135 config=cfg,
136 optional=True,
137 submodules={
138 "kv_proj": LinearBridge(name="kv_proj"),
139 "gate_proj": LinearBridge(name="gate_proj"),
140 "kv_norm": RMSNormalizationBridge(name="kv_norm", config=cfg),
141 "rotary_emb": RotaryEmbeddingBridge(name="rotary_emb", config=cfg),
142 "indexer": GeneralizedComponent(
143 name="indexer",
144 optional=True,
145 submodules={
146 "kv_proj": LinearBridge(name="kv_proj"),
147 "gate_proj": LinearBridge(name="gate_proj"),
148 "kv_norm": RMSNormalizationBridge(name="kv_norm", config=cfg),
149 "q_b_proj": LinearBridge(name="q_b_proj"),
150 "scorer": GeneralizedComponent(
151 name="scorer",
152 submodules={
153 "weights_proj": LinearBridge(name="weights_proj"),
154 },
155 ),
156 "rotary_emb": RotaryEmbeddingBridge(name="rotary_emb", config=cfg),
157 },
158 ),
159 },
160 )
163class DeepSeekV4ArchitectureAdapter(ArchitectureAdapter):
164 """Adapter for ``DeepseekV4ForCausalLM`` (Flash and Pro variants)."""
166 # The isolated component harness assumes a three-dimensional residual. V4's
167 # mHC stack is four-dimensional, so parity is covered by integration tests and
168 # verify_models' whole-model hook/text phases instead of isolated Phase 1.
169 applicable_phases: list[int] = [2, 4]
171 def __init__(self, cfg: Any) -> None:
172 super().__init__(cfg)
174 self.cfg.normalization_type = "RMS"
175 self.cfg.uses_rms_norm = True
176 self.cfg.final_rms = True
177 self.cfg.rmsnorm_uses_offset = False
178 self.cfg.positional_embedding_type = "rotary"
179 self.cfg.gated_mlp = True
180 self.cfg.attn_implementation = "eager"
182 # Folding/centering assumes one additive residual stream. Applying either
183 # transform to mHC's learned collapse/expand path is not basis preserving.
184 self.supports_fold_ln = False
185 self.supports_center_writing_weights = False
186 self.weight_processing_conversions = {}
188 def hyper_connection(name: str) -> DeepseekV4HyperConnectionBridge:
189 return DeepseekV4HyperConnectionBridge(
190 name=name,
191 config=self.cfg,
192 submodules={
193 "input_norm": GeneralizedComponent(name="input_norm"),
194 },
195 )
197 attention = GeneralizedComponent(
198 name="self_attn",
199 submodules={
200 "q_a_proj": LinearBridge(name="q_a_proj"),
201 "q_a_norm": RMSNormalizationBridge(name="q_a_norm", config=self.cfg),
202 "q_b_proj": LinearBridge(name="q_b_proj"),
203 "q_b_norm": GeneralizedComponent(name="q_b_norm"),
204 "kv_proj": LinearBridge(name="kv_proj"),
205 "kv_norm": RMSNormalizationBridge(name="kv_norm", config=self.cfg),
206 "compressor": _compressor_bridge(self.cfg),
207 "o_a_proj": GeneralizedComponent(name="o_a_proj"),
208 "o_b_proj": LinearBridge(name="o_b_proj"),
209 },
210 )
212 mlp = MoEBridge(
213 name="mlp",
214 config=self.cfg,
215 submodules={
216 "gate": GeneralizedComponent(name="gate"),
217 "experts": GeneralizedComponent(name="experts"),
218 "shared_experts": GatedMLPBridge(
219 name="shared_experts",
220 config=self.cfg,
221 submodules={
222 "gate": LinearBridge(name="gate_proj"),
223 "in": LinearBridge(name="up_proj"),
224 "out": LinearBridge(name="down_proj"),
225 },
226 ),
227 },
228 )
230 self.component_mapping = {
231 "embed": EmbeddingBridge(name="model.embed_tokens"),
232 "rotary_emb": RotaryEmbeddingBridge(name="model.rotary_emb", config=self.cfg),
233 "blocks": DeepseekV4BlockBridge(
234 name="model.layers",
235 config=self.cfg,
236 submodules={
237 "attn_hc": hyper_connection("attn_hc"),
238 "ln1": RMSNormalizationBridge(name="input_layernorm", config=self.cfg),
239 "attn": attention,
240 "mlp_hc": hyper_connection("ffn_hc"),
241 "ln2": RMSNormalizationBridge(name="post_attention_layernorm", config=self.cfg),
242 "mlp": mlp,
243 },
244 ),
245 "hc_head": GeneralizedComponent(
246 name="model.hc_head",
247 submodules={
248 "input_norm": GeneralizedComponent(name="input_norm"),
249 },
250 ),
251 "ln_final": RMSNormalizationBridge(name="model.norm", config=self.cfg),
252 "unembed": UnembeddingBridge(name="lm_head"),
253 }
255 def prepare_loading(self, model_name: str, model_kwargs: dict) -> None:
256 """Force eager attention so the delegated attention path is deterministic."""
257 model_kwargs["attn_implementation"] = "eager"
259 def prepare_model(self, hf_model: Any) -> None:
260 """Force eager attention on a pre-loaded model before installing bridges."""
261 force_eager_attention(hf_model, per_layer=True)