Coverage for transformer_lens/model_bridge/supported_architectures/llama4.py: 85%

1"""Llama 4 (text) architecture adapter. 

2 

3Meta's Llama 4 text decoder (``Llama4ForCausalLM``): llama-style RMS-norm 

4blocks whose attention adds complex-valued interleaved RoPE, NoPE layers 

5with temperature tuning, post-RoPE weightless L2 QK-norm, and chunked 

6attention masks — so attention stays delegated to HF. The feed-forward is 

7either a sparse MoE (batched 3D experts + top-k sigmoid router + shared 

8expert) or a dense gated MLP on non-MoE layers. 

9""" 

10 

11from typing import Any 

12 

13import torch 

14 

15from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter 

16from transformer_lens.model_bridge.generalized_components import ( 

17 AttentionBridge, 

18 BlockBridge, 

19 EmbeddingBridge, 

20 GatedMLPBridge, 

21 LinearBridge, 

22 MoEBridge, 

23 RMSNormalizationBridge, 

24 UnembeddingBridge, 

25) 

26from transformer_lens.model_bridge.generalized_components.base import ( 

27 CloneOutputUnderGradMixin, 

28) 

29 

30 

31class _Llama4SharedExpertBridge(CloneOutputUnderGradMixin, GatedMLPBridge): 

32 """Llama4TextMoe accumulates routed output into the shared-expert result 

33 with an in-place ``add_``; clone under grad (see mixin).""" 

34 

35 

36class _Llama4MoEBridge(MoEBridge): 

37 """MoEBridge that fires hook_out in [batch, seq, d_model]. 

38 

39 Llama4TextMoe flattens to [batch * seq, d_model] internally and the 

40 decoder layer views the result back, so hooks are fired on an 

41 input-shaped view and the HF-native flat shape is returned. 

42 """ 

43 

44 def forward(self, *args: Any, **kwargs: Any) -> Any: 

45 if self.original_component is None: 45 ↛ 46line 45 didn't jump to line 46 because the condition on line 45 was never true

46 raise RuntimeError( 

47 f"Original component not set for {self.name}. Call set_original_component() first." 

48 ) 

49 if len(args) > 0: 49 ↛ 53line 49 didn't jump to line 53 because the condition on line 49 was always true

50 hidden = self.hook_in(args[0]) 

51 args = (hidden,) + args[1:] 

52 else: 

53 hidden = self.hook_in(kwargs["hidden_states"]) 

54 kwargs = {**kwargs, "hidden_states": hidden} 

55 output = self.original_component(*args, **kwargs) 

56 if isinstance(output, tuple): 

57 flat = output[0] 

58 if len(output) > 1: 58 ↛ 60line 58 didn't jump to line 60 because the condition on line 58 was always true

59 self.hook_router_scores(output[1]) 

60 hooked = self.hook_out(flat.view(hidden.shape)) 

61 return (hooked.view(flat.shape),) + output[1:] 

62 assert isinstance(output, torch.Tensor) 

63 return self.hook_out(output.view(hidden.shape)).view(output.shape) 

64 

65 

66class Llama4ArchitectureAdapter(ArchitectureAdapter): 

67 """Architecture adapter for Llama4ForCausalLM models.""" 

68 

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

70 """Initialize the Llama 4 architecture adapter.""" 

71 super().__init__(cfg) 

72 

73 self._set_rms_rotary_defaults() 

74 self.cfg.attn_implementation = "eager" 

75 

76 self.weight_processing_conversions = { 

77 **self._qkvo_weight_conversions(), 

78 } 

79 

80 self.component_mapping = { 

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

82 "blocks": BlockBridge( 

83 name="model.layers", 

84 submodules={ 

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

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

87 # Complex-tensor RoPE, NoPE temperature tuning, L2 QK-norm, 

88 # and chunked masks live in HF's forward; keep it native. 

89 "attn": AttentionBridge( 

90 name="self_attn", 

91 config=self.cfg, 

92 submodules={ 

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

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

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

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

97 }, 

98 maintain_native_attention=True, 

99 requires_attention_mask=True, 

100 ), 

101 # The router returns a (scores, logits) tuple, so it stays 

102 # unwrapped; MoEBridge.hook_router_scores captures logits. 

103 # Non-MoE layers hold a dense gated MLP under the same name; 

104 # its projections map as optional dense_* submodules. 

105 "mlp": _Llama4MoEBridge( 

106 name="feed_forward", 

107 config=self.cfg, 

108 submodules={ 

109 # Dense-layer projections (absent on MoE layers). 

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

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

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

113 "shared_expert": _Llama4SharedExpertBridge( 

114 name="shared_expert", 

115 config=self.cfg, 

116 optional=True, 

117 submodules={ 

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

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

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

121 }, 

122 ), 

123 }, 

124 ), 

125 }, 

126 ), 

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

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

129 }