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

1"""Mixtral architecture adapter.""" 

2 

3from typing import Any 

4 

5from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter 

6from transformer_lens.model_bridge.generalized_components import ( 

7 BlockBridge, 

8 EmbeddingBridge, 

9 LinearBridge, 

10 MoEBridge, 

11 MoERouterBridge, 

12 PositionEmbeddingsAttentionBridge, 

13 RMSNormalizationBridge, 

14 RotaryEmbeddingBridge, 

15 UnembeddingBridge, 

16) 

17 

18 

19class MixtralArchitectureAdapter(ArchitectureAdapter): 

20 """Architecture adapter for Mixtral models. 

21 

22 Mixtral uses a pre-norm architecture with RMSNorm, rotary position embeddings 

23 (RoPE), and a Sparse Mixture of Experts MLP. Key features: 

24 

25 - Pre-norm: RMSNorm applied BEFORE attention and BEFORE MLP. 

26 - Rotary embeddings: stored at model.rotary_emb and passed per-forward-call. 

27 - Sparse MoE: batched expert parameters (gate_up_proj, down_proj as 3D tensors). 

28 - MixtralAttention.forward() requires position_embeddings and attention_mask args. 

29 - Optional GQA (n_key_value_heads may differ from n_heads). 

30 """ 

31 

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

33 """Initialize the Mixtral architecture adapter.""" 

34 super().__init__(cfg) 

35 

36 self._set_rms_rotary_defaults(final_rms=False) 

37 

38 n_kv_heads = ( 

39 self.cfg.n_key_value_heads 

40 if hasattr(self.cfg, "n_key_value_heads") and self.cfg.n_key_value_heads is not None 

41 else self.cfg.n_heads 

42 ) 

43 

44 self.weight_processing_conversions = { 

45 **self._qkvo_weight_conversions(), 

46 } 

47 

48 # Set up component mapping 

49 self.component_mapping = { 

50 "embed": EmbeddingBridge(name="model.embed_tokens"), 

51 "rotary_emb": RotaryEmbeddingBridge(name="model.rotary_emb", config=self.cfg), 

52 "blocks": BlockBridge( 

53 name="model.layers", 

54 submodules={ 

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

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

57 # MixtralAttention.forward() requires position_embeddings and 

58 # attention_mask as positional arguments (not optional kwargs). 

59 "attn": PositionEmbeddingsAttentionBridge( 

60 name="self_attn", 

61 config=self.cfg, 

62 submodules={ 

63 "q": LinearBridge(name="q_proj"), 

64 "k": LinearBridge(name="k_proj"), 

65 "v": LinearBridge(name="v_proj"), 

66 "o": LinearBridge(name="o_proj"), 

67 }, 

68 requires_attention_mask=True, 

69 requires_position_embeddings=True, 

70 ), 

71 # Mixtral uses batched expert parameters (gate_up_proj, down_proj 

72 # as 3D tensors) rather than a ModuleList of individual experts. 

73 # MoEBridge wraps the entire MLP module and delegates to HF's 

74 # native forward pass. 5.13 renamed the decoder-layer attr 

75 # block_sparse_moe -> mlp. 

76 "mlp": MoEBridge( 

77 name="mlp", 

78 config=self.cfg, 

79 submodules={ 

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

81 }, 

82 ), 

83 }, 

84 ), 

85 "ln_final": RMSNormalizationBridge(name="model.norm", config=self.cfg), 

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

87 }