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

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

2 

3Hybrid linear-attention (GatedDeltaNet) + full-attention with sparse MoE MLP 

4(256 experts, top-8 routing, shared expert in public checkpoints). Same hybrid 

5design as Qwen3.5 dense and the same MoE block family as Qwen3-Next. 

6 

7Two adapters: text-only ``Qwen3_5MoeForCausalLM`` and the vision-language 

8``Qwen3_5MoeForConditionalGeneration`` (text backbone nested under 

9``model.language_model`` plus the Qwen3.5 vision tower). 

10""" 

11 

12from transformer_lens.model_bridge.generalized_components import ( 

13 LinearBridge, 

14 MoEBridge, 

15 MoERouterBridge, 

16) 

17from transformer_lens.model_bridge.supported_architectures.qwen3_5 import ( 

18 Qwen3_5ArchitectureAdapter, 

19) 

20from transformer_lens.model_bridge.supported_architectures.qwen3_5_multimodal import ( 

21 Qwen3_5MultimodalArchitectureAdapter, 

22) 

23 

24 

25def _sparse_moe_mlp(adapter): 

26 """Qwen3.5 sparse MoE block: 3-tuple router, batched experts, shared expert.""" 

27 return MoEBridge( 

28 name="mlp", 

29 config=adapter.cfg, 

30 submodules={ 

31 "gate": MoERouterBridge(name="gate"), 

32 "experts": MoEBridge(name="experts", config=adapter.cfg), 

33 "shared_expert": adapter._gated_mlp(name="shared_expert"), 

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

35 }, 

36 ) 

37 

38 

39class Qwen3_5MoeArchitectureAdapter(Qwen3_5ArchitectureAdapter): 

40 """Text-only Qwen3.5-MoE: hybrid GatedDeltaNet + full attention, sparse MoE MLP.""" 

41 

42 _multimodal_arch_name: str = "Qwen3_5MoeForConditionalGeneration" 

43 

44 def _build_mlp_bridge(self): 

45 """Sparse MoE MLP (router + batched experts + shared expert).""" 

46 return _sparse_moe_mlp(self) 

47 

48 

49class Qwen3_5MoeMultimodalArchitectureAdapter(Qwen3_5MultimodalArchitectureAdapter): 

50 """Vision-language adapter for Qwen3_5MoeForConditionalGeneration. 

51 

52 Reuses the Qwen3.5 multimodal wiring (language model under 

53 ``model.language_model`` + vision tower) with the MLP swapped for sparse MoE. 

54 """ 

55 

56 def _build_mlp_bridge(self): 

57 """Sparse MoE MLP (router + batched experts + shared expert).""" 

58 return _sparse_moe_mlp(self)