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 GeneralizedComponent, 

29) 

30 

31 

32class _Llama4SharedExpertBridge(CloneOutputUnderGradMixin, GatedMLPBridge): 

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

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

35 

36 

37class _Llama4MoEBridge(MoEBridge): 

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

39 

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

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

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

43 """ 

44 

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

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

47 raise RuntimeError( 

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

49 ) 

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

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

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

53 else: 

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

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

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

57 if isinstance(output, tuple): 

58 flat = output[0] 

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

60 self.hook_router_scores(output[1]) 

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

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

63 assert isinstance(output, torch.Tensor) 

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

65 

66 

67class Llama4ArchitectureAdapter(ArchitectureAdapter): 

68 """Architecture adapter for Llama4ForCausalLM models.""" 

69 

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

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

72 super().__init__(cfg) 

73 

74 self._set_rms_rotary_defaults() 

75 self.cfg.attn_implementation = "eager" 

76 

77 self.weight_processing_conversions = { 

78 **self._qkvo_weight_conversions(), 

79 } 

80 

81 self.component_mapping = { 

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

83 "blocks": BlockBridge( 

84 name="model.layers", 

85 submodules={ 

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

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

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

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

90 "attn": AttentionBridge( 

91 name="self_attn", 

92 config=self.cfg, 

93 submodules={ 

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

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

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

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

98 }, 

99 maintain_native_attention=True, 

100 requires_attention_mask=True, 

101 ), 

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

103 # unwrapped; MoEBridge.hook_router_scores captures logits. 

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

105 # its projections map as optional dense_* submodules. 

106 "mlp": _Llama4MoEBridge( 

107 name="feed_forward", 

108 config=self.cfg, 

109 sparse_required=("router",), 

110 submodules={ 

111 # HF creates the router unconditionally on MoE 

112 # layers; mapping it makes a rename loud and gives 

113 # the layer real router observability. 

114 "router": GeneralizedComponent(name="router", optional=True), 

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

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

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

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

119 "shared_expert": _Llama4SharedExpertBridge( 

120 name="shared_expert", 

121 config=self.cfg, 

122 optional=True, 

123 submodules={ 

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

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

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

127 }, 

128 ), 

129 }, 

130 ), 

131 }, 

132 ), 

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

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

135 }