Coverage for transformer_lens/model_bridge/supported_architectures/deepseek_v2.py: 100%

1"""DeepSeek V2 architecture adapter. 

2 

3Supports DeepSeek-V2, DeepSeek-V2-Lite, and DeepSeek-Coder-V2 models 

4(all use DeepseekV2ForCausalLM). 

5 

6Key features: 

7- Multi-Head Latent Attention (MLA): Q and KV compressed via LoRA-style projections. 

8 DeepSeek-V2-Lite sets q_lora_rank=None, skipping Q compression and using a direct 

9 q_proj instead — MLAAttentionBridge.forward handles both paths automatically. 

10- Mixture of Experts (MoE) with shared experts on most layers 

11- Dense MLP on first `first_k_dense_replace` layers 

12""" 

13 

14from typing import Any 

15 

16from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter 

17from transformer_lens.model_bridge.generalized_components import ( 

18 EmbeddingBridge, 

19 LinearBridge, 

20 MLAAttentionBridge, 

21 MLABlockBridge, 

22 MoEBridge, 

23 RMSNormalizationBridge, 

24 RotaryEmbeddingBridge, 

25 UnembeddingBridge, 

26) 

27from transformer_lens.model_bridge.generalized_components.base import ( 

28 GeneralizedComponent, 

29) 

30 

31 

32class DeepSeekV2ArchitectureAdapter(ArchitectureAdapter): 

33 """Architecture adapter for DeepSeek V2 / V2-Lite / Coder-V2 models. 

34 

35 Uses RMSNorm, MLA with compressed Q/KV projections (or direct Q projection 

36 when q_lora_rank is None), partial RoPE, MoE on most layers (dense MLP on 

37 first few), and no biases. 

38 """ 

39 

40 _testing_eager = None 

41 

42 def __init__(self, cfg: Any) -> None: 

43 super().__init__(cfg) 

44 

45 self.cfg.normalization_type = "RMS" 

46 self.cfg.positional_embedding_type = "rotary" 

47 self.cfg.gated_mlp = True 

48 self.cfg.final_rms = True 

49 self.cfg.uses_rms_norm = True 

50 

51 # MLA has no per-head q/k/v to fold into; skip LN folding. 

52 self.supports_fold_ln = False 

53 

54 self.weight_processing_conversions = {} 

55 

56 self.component_mapping = { 

57 "embed": EmbeddingBridge(name="model.embed_tokens"), 

58 "rotary_emb": RotaryEmbeddingBridge(name="model.rotary_emb", config=self.cfg), 

59 "blocks": MLABlockBridge( 

60 name="model.layers", 

61 submodules={ 

62 "ln1": RMSNormalizationBridge(name="input_layernorm", config=self.cfg), 

63 "ln2": RMSNormalizationBridge(name="post_attention_layernorm", config=self.cfg), 

64 "attn": MLAAttentionBridge( 

65 name="self_attn", 

66 config=self.cfg, 

67 submodules={ 

68 # V2-full (q_lora_rank set): two-stage LoRA Q compression. 

69 # These are absent in V2-Lite — marked optional so bridge 

70 # setup skips them gracefully. The actual forward call is 

71 # handled inside MLAAttentionBridge which checks q_lora_rank. 

72 "q_a_proj": LinearBridge(name="q_a_proj", optional=True), 

73 # q_a_layernorm is a norm inside the attention block; its 

74 # forward is called directly by MLAAttentionBridge, so a 

75 # plain GeneralizedComponent (with optional support) suffices. 

76 "q_a_layernorm": GeneralizedComponent( 

77 name="q_a_layernorm", optional=True 

78 ), 

79 "q_b_proj": LinearBridge(name="q_b_proj", optional=True), 

80 # V2-Lite only: direct Q projection, no compression. 

81 "q_proj": LinearBridge(name="q_proj", optional=True), 

82 # KV path — always present across all V2 variants. 

83 "kv_a_proj_with_mqa": LinearBridge(name="kv_a_proj_with_mqa"), 

84 "kv_a_layernorm": RMSNormalizationBridge( 

85 name="kv_a_layernorm", config=self.cfg 

86 ), 

87 "kv_b_proj": LinearBridge(name="kv_b_proj"), 

88 "o": LinearBridge(name="o_proj"), 

89 }, 

90 ), 

91 # Layers before first_k_dense_replace are DeepseekV2MLP: 

92 # the MoE parts are skipped and the dense_* projections bind, 

93 # which is what gives those layers gated-MLP neuron hooks. 

94 "mlp": self._build_mlp_bridge(), 

95 }, 

96 ), 

97 "ln_final": RMSNormalizationBridge(name="model.norm", config=self.cfg), 

98 "unembed": UnembeddingBridge(name="lm_head"), 

99 } 

100 

101 def _build_mlp_bridge(self): 

102 """Routed MoE with optional shared experts; Youtu (all-dense) overrides. 

103 

104 Dense-prefix layers (idx < first_k_dense_replace) bind as gated MLPs 

105 with neuron-basis hook_pre/hook_pre_linear/hook_post (#1645). 

106 """ 

107 return MoEBridge( 

108 name="mlp", 

109 config=self.cfg, 

110 sparse_required=("gate",), 

111 submodules={ 

112 # Router is a custom Module, not nn.Linear. 

113 "gate": GeneralizedComponent(name="gate", optional=True), 

114 "shared_experts": self._gated_mlp(name="shared_experts", optional=True), 

115 # Dense-layer projections (present only on layers before 

116 # first_k_dense_replace); their presence is what makes 

117 # MoEBridge bind gated-MLP neuron hooks there (#1645). 

118 "dense_gate": LinearBridge(name="gate_proj", optional=True), 

119 "dense_in": LinearBridge(name="up_proj", optional=True), 

120 "dense_out": LinearBridge(name="down_proj", optional=True), 

121 }, 

122 )