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

1"""Qwen3Next architecture adapter. 

2 

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

43 linear-attn layers per 1 full-attn layer. Extends Qwen3 base with 

5optional attention mapping, MoE MLP, and fold_ln disabled. 

6""" 

7 

8from typing import Any 

9 

10from transformer_lens.model_bridge.generalized_components import ( 

11 LinearBridge, 

12 MoEBridge, 

13 MoERouterBridge, 

14) 

15from transformer_lens.model_bridge.supported_architectures.qwen3 import ( 

16 Qwen3ArchitectureAdapter, 

17) 

18 

19 

20class Qwen3NextArchitectureAdapter(Qwen3ArchitectureAdapter): 

21 """Hybrid linear-attention + full-attention with sparse MoE MLP. 

22 

23 Same hybrid design as Qwen3.5 but with MoE instead of dense MLP. 

24 """ 

25 

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

27 # q_proj stays 2x-wide through weight processing: HF and the attention bridge both 

28 # split [query|gate] per head at forward time; slicing the gate out changes outputs. 

29 setattr(cfg, "gated_q_proj", True) 

30 # HF's Qwen3NextRMSNorm scales by (1 + weight), not weight (cf. Gemma). 

31 setattr(cfg, "rmsnorm_uses_offset", True) 

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

33 

34 def _build_mlp_bridge(self): 

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

36 return MoEBridge( 

37 name="mlp", 

38 config=self.cfg, 

39 sparse_required=("gate",), 

40 submodules={ 

41 # Plain-tensor SparseMoeBlock: router observability comes from 

42 # the gate submodule. Dense fallback layers (mlp_only_layers / 

43 # decoder_sparse_step) bind the projections below (#1645). 

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

45 "shared_expert": self._gated_mlp(name="shared_expert", optional=True), 

46 "shared_expert_gate": LinearBridge(name="shared_expert_gate", optional=True), 

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

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

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

50 }, 

51 )