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

1"""Llama 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 GatedMLPBridge, 

10 LinearBridge, 

11 PositionEmbeddingsAttentionBridge, 

12 RMSNormalizationBridge, 

13 RotaryEmbeddingBridge, 

14 UnembeddingBridge, 

15) 

16 

17 

18class LlamaArchitectureAdapter(ArchitectureAdapter): 

19 """Architecture adapter for Llama models. 

20 

21 Optional Parameters (may not exist in state_dict): 

22 ------------------------------------------------- 

23 LLaMA models do NOT have biases on attention and MLP projections: 

24 

25 - blocks.{i}.attn.b_Q - No bias on query projection 

26 - blocks.{i}.attn.b_K - No bias on key projection 

27 - blocks.{i}.attn.b_V - No bias on value projection 

28 - blocks.{i}.attn.b_O - No bias on output projection 

29 - blocks.{i}.mlp.b_in - No bias on MLP input (up_proj) 

30 - blocks.{i}.mlp.b_gate - No bias on MLP gate projection 

31 - blocks.{i}.mlp.b_out - No bias on MLP output (down_proj) 

32 - blocks.{i}.ln1.b - RMSNorm has no bias 

33 - blocks.{i}.ln2.b - RMSNorm has no bias 

34 - ln_final.b - RMSNorm has no bias 

35 

36 Weight processing must handle these missing biases gracefully using 

37 ProcessWeights._safe_get_tensor() or by checking for None values. 

38 """ 

39 

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

41 """Initialize the Llama architecture adapter.""" 

42 super().__init__(cfg) 

43 

44 # Set config variables for weight processing 

45 self.cfg.normalization_type = "RMS" 

46 self.cfg.positional_embedding_type = "rotary" 

47 self.cfg.final_rms = True 

48 self.cfg.gated_mlp = True 

49 self.cfg.attn_only = False 

50 

51 self.default_config = { 

52 "d_model": cfg.d_model, 

53 "d_head": cfg.d_model // cfg.n_heads, 

54 "n_heads": cfg.n_heads, 

55 "n_layers": cfg.n_layers, 

56 "d_vocab": cfg.d_vocab, 

57 } 

58 

59 # Add GQA support for Llama 3.1, 3.2, and later models 

60 # Must set directly on cfg, not just in default_config 

61 if hasattr(cfg, "n_key_value_heads") and cfg.n_key_value_heads is not None: 

62 self.default_config["n_key_value_heads"] = cfg.n_key_value_heads 

63 self.cfg.n_key_value_heads = cfg.n_key_value_heads 

64 

65 self.cfg.uses_rms_norm = True 

66 

67 self.weight_processing_conversions = { 

68 **self._qkvo_weight_conversions(), 

69 } 

70 

71 self.component_mapping = { 

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

73 "rotary_emb": RotaryEmbeddingBridge(name="model.rotary_emb"), 

74 "blocks": BlockBridge( 

75 name="model.layers", 

76 submodules={ 

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

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

79 "attn": PositionEmbeddingsAttentionBridge( 

80 name="self_attn", 

81 config=self.cfg, 

82 submodules={ 

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

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

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

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

87 }, 

88 requires_attention_mask=True, 

89 requires_position_embeddings=True, 

90 ), 

91 "mlp": GatedMLPBridge( 

92 name="mlp", 

93 config=self.cfg, 

94 submodules={ 

95 "gate": LinearBridge(name="gate_proj"), 

96 "in": LinearBridge(name="up_proj"), 

97 "out": LinearBridge(name="down_proj"), 

98 }, 

99 ), 

100 }, 

101 ), 

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

103 "unembed": UnembeddingBridge(name="lm_head", config=self.cfg), 

104 } 

105 

106 def setup_component_testing(self, hf_model: Any, bridge_model: Any = None) -> None: 

107 """Set up rotary embedding references for Llama component testing. 

108 

109 Llama uses RoPE (Rotary Position Embeddings). We set the rotary_emb reference 

110 on all attention bridge instances for component testing. 

111 

112 Args: 

113 hf_model: The HuggingFace Llama model instance 

114 bridge_model: The TransformerBridge model (if available, set rotary_emb on actual instances) 

115 """ 

116 # Get rotary embedding instance from the model 

117 rotary_emb = hf_model.model.rotary_emb 

118 

119 # Set rotary_emb on actual bridge instances in bridge_model if available 

120 if bridge_model is not None and hasattr(bridge_model, "blocks"): 

121 # Set on each layer's actual attention bridge instance 

122 for block in bridge_model.blocks: 

123 if hasattr(block, "attn"): 

124 block.attn.set_rotary_emb(rotary_emb) 

125 

126 # Also set on the template for get_generalized_component() calls 

127 attn_bridge = self.get_generalized_component("blocks.0.attn") 

128 attn_bridge.set_rotary_emb(rotary_emb)