Coverage for transformer_lens/model_bridge/supported_architectures/qwen3_5.py: 96%

1"""Qwen3.5 architecture adapter. 

2 

3Hybrid linear-attention (GatedDeltaNet) + full-attention with dense gated MLP. 

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

5optional attention mapping and fold_ln disabled. 

6""" 

7 

8from typing import Any 

9 

10from transformer_lens.model_bridge.supported_architectures.qwen3 import ( 

11 Qwen3ArchitectureAdapter, 

12) 

13 

14 

15class Qwen3_5ArchitectureAdapter(Qwen3ArchitectureAdapter): 

16 """Hybrid linear-attention + full-attention with dense gated MLP. 

17 

18 Inherits Qwen3 config/attention/MLP structure. Differences: 

19 - Attention + linear_attn are optional (per-layer type) 

20 - Gated q_proj: [query|gate] is split at forward time, never in weight space; 

21 AttentionBridge exposes a query-only W_Q analysis view 

22 """ 

23 

24 # Multimodal wrapper architecture this text-only adapter rejects; the MoE 

25 # subclass swaps in its own name. 

26 _multimodal_arch_name: str = "Qwen3_5ForConditionalGeneration" 

27 

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

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

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

31 setattr(cfg, "gated_q_proj", True) 

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

33 setattr(cfg, "rmsnorm_uses_offset", True) 

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

35 

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

37 """Swap the multimodal config for its text_config so AutoModelForCausalLM 

38 loads the text-only model (published checkpoints carry the 

39 ForConditionalGeneration architecture).""" 

40 config = model_kwargs.get("config") 

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

42 model_kwargs["config"] = config.text_config 

43 

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

45 """Reject full multimodal checkpoints on this text-only adapter.""" 

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

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

48 class_name = type(hf_model).__name__ 

49 multimodal_arch = self._multimodal_arch_name 

50 

51 is_conditional_generation = ( 

52 class_name == multimodal_arch or multimodal_arch in architectures 

53 ) 

54 still_has_top_level_multimodal_config = hasattr(config, "text_config") 

55 if is_conditional_generation or still_has_top_level_multimodal_config: 

56 causal_arch = multimodal_arch.replace("ForConditionalGeneration", "ForCausalLM") 

57 text_config_name = multimodal_arch.replace("ForConditionalGeneration", "TextConfig") 

58 raise ValueError( 

59 f"This adapter is text-only. Pass a {causal_arch} / " 

60 f"{text_config_name} model, or load by model id with " 

61 f"TransformerBridge.boot_transformers(...) so {multimodal_arch} " 

62 f"checkpoints route to the multimodal adapter automatically." 

63 )