Coverage for transformer_lens/model_bridge/supported_architectures/qwen2_moe.py: 86%

1"""Qwen2-MoE architecture adapter.""" 

2 

3from typing import Any 

4 

5from transformer_lens.model_bridge.generalized_components import ( 

6 LinearBridge, 

7 MoEBridge, 

8 MoERouterBridge, 

9) 

10from transformer_lens.model_bridge.supported_architectures.qwen2 import ( 

11 Qwen2ArchitectureAdapter, 

12) 

13 

14 

15class Qwen2MoeRouterBridge(MoERouterBridge): 

16 """Tuple-preserving router bridge for ``Qwen2MoeTopKRouter``.""" 

17 

18 

19class Qwen2MoeArchitectureAdapter(Qwen2ArchitectureAdapter): 

20 """Architecture adapter for Qwen2-MoE models. 

21 

22 Qwen2-MoE uses the Qwen2 attention stack plus a sparse MoE MLP with an 

23 always-on shared expert path. 

24 """ 

25 

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

27 """Initialize the Qwen2-MoE architecture adapter.""" 

28 super().__init__(cfg) 

29 

30 self.cfg.attn_implementation = "eager" 

31 

32 if self.component_mapping is None: 32 ↛ 33line 32 didn't jump to line 33 because the condition on line 32 was never true

33 raise ValueError("Qwen2 component mapping was not initialized") 

34 

35 blocks = self.component_mapping["blocks"] 

36 blocks.submodules["mlp"] = MoEBridge( 

37 name="mlp", 

38 config=self.cfg, 

39 sparse_required=("gate", "experts"), 

40 submodules={ 

41 # mlp_only_layers / decoder_sparse_step let HF build a dense 

42 # Qwen2MoeMLP on some layers: the MoE parts are optional there, 

43 # and the dense projections below make those layers bind 

44 # gated-MLP neuron hooks instead of MoE boundary tensors (#1645). 

45 "gate": Qwen2MoeRouterBridge(name="gate", optional=True), 

46 "experts": MoEBridge(name="experts", config=self.cfg, optional=True), 

47 "shared_expert": self._gated_mlp(name="shared_expert", optional=True), 

48 "shared_expert_gate": LinearBridge(name="shared_expert_gate", optional=True), 

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

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

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

52 }, 

53 )