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

1"""Qwen3 architecture adapter. 

2 

3Base adapter for the Qwen3 model family. Provides shared config setup, 

4attention bridge construction, and setup_component_testing used by 

5Qwen3, Qwen3.5, and Qwen3Next variants. 

6""" 

7 

8from typing import Any 

9 

10from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter 

11from transformer_lens.model_bridge.generalized_components import ( 

12 AttentionBridge, 

13 BlockBridge, 

14 EmbeddingBridge, 

15 RMSNormalizationBridge, 

16 RotaryEmbeddingBridge, 

17 UnembeddingBridge, 

18) 

19from transformer_lens.model_bridge.generalized_components.gated_delta_net import ( 

20 GatedDeltaNetBridge, 

21) 

22 

23 

24class Qwen3ArchitectureAdapter(ArchitectureAdapter): 

25 """Architecture adapter for Qwen3 dense models. 

26 

27 RMSNorm, RoPE, GQA, Q/K head norms, gated MLP. No biases. 

28 Serves as base class for Qwen3.5 and Qwen3Next hybrid variants. 

29 """ 

30 

31 _testing_hybrid = True 

32 

33 def __init__(self, cfg: Any, *, hybrid: bool = False, lm_prefix: str = "model") -> None: 

34 super().__init__(cfg) 

35 self._setup_qwen3_config(cfg) 

36 if hybrid: 

37 self.supports_fold_ln = False 

38 self.weight_processing_conversions: dict = {} 

39 else: 

40 self.weight_processing_conversions = {**self._qkvo_weight_conversions()} 

41 self.component_mapping = self._build_component_mapping(hybrid=hybrid, lm_prefix=lm_prefix) 

42 

43 def _setup_qwen3_config(self, cfg: Any) -> None: 

44 """Config shared across all Qwen3 variants (dense, hybrid, MoE).""" 

45 self._set_rms_rotary_defaults() 

46 self.cfg.default_prepend_bos = False 

47 self.cfg.attn_implementation = "eager" 

48 

49 def _build_attention_bridge(self, optional: bool = False) -> AttentionBridge: 

50 """Standard Qwen3 attention bridge with Q/K norms.""" 

51 return self._qkvo_attention_bridge( 

52 optional=optional, 

53 extra_submodules={ 

54 "q_norm": RMSNormalizationBridge(name="q_norm", config=self.cfg), 

55 "k_norm": RMSNormalizationBridge(name="k_norm", config=self.cfg), 

56 }, 

57 # None omits the flags, keeping the bridge class's own defaults. 

58 requires_attention_mask=None, 

59 requires_position_embeddings=None, 

60 ) 

61 

62 def _build_mlp_bridge(self): 

63 """Dense gated MLP (gate_proj + up_proj -> down_proj). Override for MoE.""" 

64 return self._gated_mlp() 

65 

66 def _build_linear_attn_bridge(self, optional: bool = False) -> GatedDeltaNetBridge: 

67 """GatedDeltaNet linear-attention bridge for hybrid variants.""" 

68 return GatedDeltaNetBridge( 

69 name="linear_attn", 

70 config=self.cfg, 

71 optional=optional, 

72 ) 

73 

74 def _build_component_mapping(self, *, hybrid: bool = False, lm_prefix: str = "model") -> dict: 

75 """Parametric component mapping. hybrid=True adds optional linear_attn; lm_prefix 

76 nests the text model (``model``, or ``model.language_model`` for multimodal). lm_head 

77 stays top-level. 

78 """ 

79 block_submodules: dict = { 

80 "ln1": RMSNormalizationBridge(name="input_layernorm", config=self.cfg), 

81 "ln2": RMSNormalizationBridge(name="post_attention_layernorm", config=self.cfg), 

82 "attn": self._build_attention_bridge(optional=hybrid), 

83 "mlp": self._build_mlp_bridge(), 

84 } 

85 if hybrid: 

86 block_submodules["linear_attn"] = self._build_linear_attn_bridge(optional=True) 

87 return { 

88 "embed": EmbeddingBridge(name=f"{lm_prefix}.embed_tokens"), 

89 "rotary_emb": RotaryEmbeddingBridge(name=f"{lm_prefix}.rotary_emb", config=self.cfg), 

90 "blocks": BlockBridge(name=f"{lm_prefix}.layers", submodules=block_submodules), 

91 "ln_final": RMSNormalizationBridge(name=f"{lm_prefix}.norm", config=self.cfg), 

92 "unembed": UnembeddingBridge(name="lm_head"), 

93 }