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 submodules={ 

40 "gate": Qwen2MoeRouterBridge(name="gate"), 

41 "experts": MoEBridge(name="experts", config=self.cfg), 

42 "shared_expert": self._gated_mlp(name="shared_expert"), 

43 "shared_expert_gate": LinearBridge(name="shared_expert_gate"), 

44 }, 

45 )