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

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 typing import Any 

13 

14import torch 

15 

16from transformer_lens.model_bridge.generalized_components import ( 

17 LinearBridge, 

18 MoEBridge, 

19 MoERouterBridge, 

20) 

21from transformer_lens.model_bridge.supported_architectures.qwen3 import ( 

22 Qwen3ArchitectureAdapter, 

23) 

24from transformer_lens.model_bridge.supported_architectures.qwen3_5_multimodal import ( 

25 Qwen3_5MultimodalArchitectureAdapter, 

26) 

27 

28 

29def _sparse_moe_mlp(adapter): 

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

31 return MoEBridge( 

32 name="mlp", 

33 config=adapter.cfg, 

34 submodules={ 

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

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

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

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

39 }, 

40 ) 

41 

42 

43class Qwen3_5MoeArchitectureAdapter(Qwen3ArchitectureAdapter): 

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

45 

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

47 setattr(cfg, "gated_q_proj", True) 

48 super().__init__(cfg, hybrid=True) 

49 

50 def _build_mlp_bridge(self): 

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

52 return _sparse_moe_mlp(self) 

53 

54 def prepare_loading(self, model_name: str, model_kwargs: dict) -> None: 

55 """Swap to ``text_config`` so AutoModelForCausalLM loads the text-only model 

56 (checkpoints ship the ConditionalGeneration architecture).""" 

57 config = model_kwargs.get("config") 

58 if config is not None and hasattr(config, "text_config"): 58 ↛ exitline 58 didn't return from function 'prepare_loading' because the condition on line 58 was always true

59 model_kwargs["config"] = config.text_config 

60 

61 def prepare_model(self, hf_model: Any) -> None: 

62 """Reject full multimodal Qwen3.5-MoE models on this text-only adapter.""" 

63 config = getattr(hf_model, "config", None) 

64 architectures = getattr(config, "architectures", []) or [] 

65 class_name = type(hf_model).__name__ 

66 

67 is_conditional_generation = ( 

68 class_name == "Qwen3_5MoeForConditionalGeneration" 

69 or "Qwen3_5MoeForConditionalGeneration" in architectures 

70 ) 

71 still_has_top_level_multimodal_config = hasattr(config, "text_config") 

72 if is_conditional_generation or still_has_top_level_multimodal_config: 

73 raise ValueError( 

74 "This adapter is text-only. Pass a Qwen3_5MoeForCausalLM / " 

75 "Qwen3_5MoeTextConfig model, or load by model id with " 

76 "TransformerBridge.boot_transformers(...) so Qwen3_5MoeForConditionalGeneration " 

77 "checkpoints route to the multimodal adapter automatically." 

78 ) 

79 

80 def preprocess_weights(self, state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: 

81 """Slice query half from gated q_proj.weight for weight-space analysis.""" 

82 return self._preprocess_gated_q_proj(state_dict, self.cfg.n_heads, self.cfg.d_head) 

83 

84 

85class Qwen3_5MoeMultimodalArchitectureAdapter(Qwen3_5MultimodalArchitectureAdapter): 

86 """Vision-language adapter for Qwen3_5MoeForConditionalGeneration. 

87 

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

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

90 """ 

91 

92 def _build_mlp_bridge(self): 

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

94 return _sparse_moe_mlp(self)