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

1"""Qwen3-VL-MoE architecture adapter. 

2 

3Alibaba's Qwen3-VL-MoE (``Qwen3VLMoeForConditionalGeneration``): the 

4Qwen3-VL layout (DeepStack vision tower, interleaved-mRoPE text decoder) 

5with a sparse-MoE MLP — batched experts plus a parameter-only top-k 

6router returning a (logits, scores, indices) tuple, so the router stays 

7unwrapped and MoEBridge delegates the block. Layers listed in 

8``mlp_only_layers`` hold a dense gated MLP under the same name. 

9""" 

10 

11from typing import Any 

12 

13from transformer_lens.model_bridge.generalized_components import ( 

14 LinearBridge, 

15 MoEBridge, 

16 MoERouterBridge, 

17) 

18from transformer_lens.model_bridge.supported_architectures.qwen3_vl import ( 

19 Qwen3VLArchitectureAdapter, 

20) 

21 

22 

23class Qwen3VLMoeArchitectureAdapter(Qwen3VLArchitectureAdapter): 

24 """Architecture adapter for Qwen3VLMoeForConditionalGeneration models.""" 

25 

26 def _build_mlp_bridge(self) -> Any: 

27 """Sparse MoE block; gate/experts absent on dense mlp_only_layers.""" 

28 return MoEBridge( 

29 name="mlp", 

30 config=self.cfg, 

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

32 submodules={ 

33 "gate": MoERouterBridge(name="gate", optional=True), 

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

35 # Dense layers listed in mlp_only_layers (#1645). 

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

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

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

39 }, 

40 )