Coverage for transformer_lens/model_bridge/supported_architectures/deepseek_v2.py: 100%
39 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-21 19:27 +0000
1"""DeepSeek V2 adapter and the shared DeepSeek-MLA family base (DeepSeek V2/V3,
2GLM-MoE-DSA, GLM-4.7-Flash; Youtu via V2); per-member differences are declarative.
4DeepSeek V2 support covers DeepSeek-V2, DeepSeek-V2-Lite, and DeepSeek-Coder-V2
5(all use DeepseekV2ForCausalLM).
6"""
8from typing import Any, Dict, Optional, Type
10from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter
11from transformer_lens.model_bridge.generalized_components import (
12 EmbeddingBridge,
13 LinearBridge,
14 MLAAttentionBridge,
15 MLABlockBridge,
16 MoEBridge,
17 RMSNormalizationBridge,
18 RotaryEmbeddingBridge,
19 UnembeddingBridge,
20)
21from transformer_lens.model_bridge.generalized_components.base import (
22 GeneralizedComponent,
23)
26class DeepSeekMLAFamilyArchitectureAdapter(ArchitectureAdapter):
27 """Shared base for DeepSeek-style Multi-head Latent Attention + MoE decoders.
29 Members share LoRA-compressed attention (q_a/q_b and kv_a/kv_b projections with
30 decoupled RoPE), RMSNorm, no biases, MLA weights kept in HF layout (no QKVO
31 rearrangements, no LN folding), and a routed MoE whose router/shared experts are
32 absent on dense layers. Divergence is declarative:
34 - ``q_lora_optional``: some checkpoints (DeepSeek-V2-Lite; GigaChat3 on the V3
35 class) set q_lora_rank=None — no compressed-Q pair, direct q_proj instead — so
36 the whole Q path is optional; ``_build_q_a_layernorm`` picks the norm's type.
37 Families whose checkpoints always compress Q keep the projections REQUIRED so
38 a genuinely missing weight fails loudly; never relax them family-wide.
39 - ``attention_cls``: GLM-MoE-DSA swaps in its sparse-attention bridge.
40 - ``_build_mlp_bridge`` / ``_build_router`` seams: Youtu maps all-dense MLPs;
41 GLM-4.7-Flash uses a MoERouterBridge router.
42 """
44 attention_cls: Type[MLAAttentionBridge] = MLAAttentionBridge
45 q_lora_optional: bool = False
46 # Sets cfg.attn_implementation="eager" (member-specific WHY at each override).
47 eager_attention: bool = False
48 # Sets cfg.default_prepend_bos; None leaves the config default untouched.
49 prepend_bos: Optional[bool] = None
51 def __init__(self, cfg: Any) -> None:
52 """Apply the family config knobs and build the shared mapping."""
53 super().__init__(cfg)
55 self._set_rms_rotary_defaults()
56 if self.eager_attention:
57 self.cfg.attn_implementation = "eager"
58 if self.prepend_bos is not None:
59 self.cfg.default_prepend_bos = self.prepend_bos
61 # MLA has no per-head q/k/v to fold into; skip LN folding.
62 self.supports_fold_ln = False
64 # MLA weights keep their HF layout; no QKVO rearrangements apply.
65 self.weight_processing_conversions = {}
67 self.component_mapping = self._build_component_mapping()
69 def _build_component_mapping(self) -> dict:
70 """PRE-NORM mapping: ln1 = input_layernorm (before attention),
71 ln2 = post_attention_layernorm (before MLP)."""
72 return {
73 "embed": EmbeddingBridge(name="model.embed_tokens"),
74 "rotary_emb": RotaryEmbeddingBridge(name="model.rotary_emb", config=self.cfg),
75 "blocks": MLABlockBridge(
76 name="model.layers",
77 submodules={
78 "ln1": RMSNormalizationBridge(name="input_layernorm", config=self.cfg),
79 "ln2": RMSNormalizationBridge(name="post_attention_layernorm", config=self.cfg),
80 "attn": self.attention_cls(
81 name="self_attn",
82 config=self.cfg,
83 submodules=self._build_attention_submodules(),
84 ),
85 "mlp": self._build_mlp_bridge(),
86 },
87 ),
88 "ln_final": RMSNormalizationBridge(name="model.norm", config=self.cfg),
89 "unembed": UnembeddingBridge(name="lm_head"),
90 }
92 def _build_attention_submodules(self) -> Dict[str, GeneralizedComponent]:
93 """MLA projection submodules; the Q path follows ``q_lora_optional``."""
94 optional = self.q_lora_optional
95 submodules: Dict[str, GeneralizedComponent] = {
96 "q_a_proj": LinearBridge(name="q_a_proj", optional=optional),
97 "q_a_layernorm": self._build_q_a_layernorm(optional),
98 "q_b_proj": LinearBridge(name="q_b_proj", optional=optional),
99 }
100 if optional:
101 # Direct Q projection, mutually exclusive with the compressed pair above;
102 # MLAAttentionBridge.forward picks the path from q_lora_rank.
103 submodules["q_proj"] = LinearBridge(name="q_proj", optional=True)
104 submodules.update(
105 {
106 # KV path — always compressed, present in every family member.
107 "kv_a_proj_with_mqa": LinearBridge(name="kv_a_proj_with_mqa"),
108 "kv_a_layernorm": RMSNormalizationBridge(name="kv_a_layernorm", config=self.cfg),
109 "kv_b_proj": LinearBridge(name="kv_b_proj"),
110 "o": LinearBridge(name="o_proj"),
111 }
112 )
113 return submodules
115 def _build_q_a_layernorm(self, optional: bool) -> GeneralizedComponent:
116 """Compressed-Q norm. Its forward is called directly by MLAAttentionBridge, so
117 on the optional (direct-Q) path a plain GeneralizedComponent suffices; V3
118 overrides to keep the full norm bridge."""
119 if optional:
120 return GeneralizedComponent(name="q_a_layernorm", optional=True)
121 return RMSNormalizationBridge(name="q_a_layernorm", config=self.cfg)
123 def _build_router(self) -> GeneralizedComponent:
124 """Router is a custom Module, not nn.Linear; absent on dense layers."""
125 return GeneralizedComponent(name="gate", optional=True)
127 def _build_mlp_bridge(self):
128 """Routed MoE with optional shared experts — on dense layers (e.g. idx <
129 first_k_dense_replace) router and shared_experts are absent, so setup skips
130 the optional submodules; Youtu (all-dense) overrides.
132 Dense-prefix layers bind as gated MLPs with neuron-basis
133 hook_pre/hook_pre_linear/hook_post (#1645).
134 """
135 return MoEBridge(
136 name="mlp",
137 config=self.cfg,
138 sparse_required=("gate",),
139 submodules={
140 "gate": self._build_router(),
141 "shared_experts": self._gated_mlp(name="shared_experts", optional=True),
142 # Dense-layer projections (present only on the dense layers of
143 # this interleaved stack); their presence is what makes MoEBridge
144 # bind gated-MLP neuron hooks there (#1645).
145 "dense_gate": LinearBridge(name="gate_proj", optional=True),
146 "dense_in": LinearBridge(name="up_proj", optional=True),
147 "dense_out": LinearBridge(name="down_proj", optional=True),
148 },
149 )
152class DeepSeekV2ArchitectureAdapter(DeepSeekMLAFamilyArchitectureAdapter):
153 """Architecture adapter for DeepSeek V2 / V2-Lite / Coder-V2 models.
155 Uses RMSNorm, MLA with compressed Q/KV projections (or direct Q projection
156 when q_lora_rank is None), partial RoPE, MoE on most layers (dense MLP on
157 first few), and no biases.
158 """
160 _testing_eager = None
162 # V2-Lite sets q_lora_rank=None: no q_a_proj/q_b_proj in the state_dict,
163 # direct q_proj instead — MLAAttentionBridge.forward handles both paths.
164 q_lora_optional = True