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 

10import torch 

11 

12from transformer_lens.model_bridge.generalized_components import ( 

13 LinearBridge, 

14 MoEBridge, 

15 MoERouterBridge, 

16) 

17from transformer_lens.model_bridge.supported_architectures.qwen3 import ( 

18 Qwen3ArchitectureAdapter, 

19) 

20 

21 

22class Qwen3NextArchitectureAdapter(Qwen3ArchitectureAdapter): 

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

24 

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

26 """ 

27 

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

29 setattr(cfg, "gated_q_proj", True) 

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

31 

32 def _build_mlp_bridge(self): 

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

34 return MoEBridge( 

35 name="mlp", 

36 config=self.cfg, 

37 sparse_required=("gate",), 

38 submodules={ 

39 # Plain-tensor SparseMoeBlock: router observability comes from 

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

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

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

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

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

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

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

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

48 }, 

49 ) 

50 

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

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

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