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

1"""Falcon-H1 parallel-hybrid architecture adapter. 

2 

3Supports ``FalconH1ForCausalLM`` (e.g. ``tiiuae/Falcon-H1-0.5B-Base``). 

4 

5Architecture overview: 

6- Every block runs **GQA attention and a Mamba-2 mixer in parallel** on the same 

7 ``input_layernorm`` output; their results are summed into the residual stream 

8 alongside scalar multipliers. A second norm (``pre_ff_layernorm``) precedes a 

9 SwiGLU feed-forward MLP. This is the parallel hybrid that distinguishes 

10 Falcon-H1 from heterogeneous-layer hybrids like NemotronH (one mixer type per 

11 layer) — here both branches are present in *every* block. 

12- RoPE via a model-level ``model.rotary_emb`` module (very large ``rope_theta``). 

13- ``mamba_rms_norm=false`` on the released checkpoints — the Mamba inner gated 

14 RMSNorm is **absent**, so ``inner_norm`` is declared optional. 

15- ~12 scalar multipliers (``embedding_multiplier``, ``attention_out_multiplier``, 

16 ``key_multiplier``, ``ssm_*_multiplier``, ``mlp_multipliers``, 

17 ``lm_head_multiplier`` …). HF applies all of these in its own forward, so in 

18 raw (passthrough) mode the bridge inherits them for free — no weight folding is 

19 needed for forward parity. Folding would only matter for compatibility mode. 

20 

21Key adapter decisions: 

22- ``BlockBridge`` is the block container: the block has two norms 

23 (``input_layernorm`` + ``pre_ff_layernorm``) and a real post-attention 

24 residual, so transformer-shaped hooks (``hook_resid_mid`` via ``ln2``) are 

25 meaningful — unlike single-norm SSM blocks that need ``SSMBlockBridge``. 

26 ``BlockBridge.forward`` delegates the whole block to HF, so the parallel 

27 attn+mamba combine happens natively. 

28- ``SSM2MixerBridge`` wraps the Mamba-2 mixer (passthrough forward). Its inner 

29 ``in_proj`` / ``conv1d`` / ``out_proj`` are mapped for hookability; the gated 

30 ``inner_norm`` is optional (absent when ``mamba_rms_norm=false``). 

31- ``applicable_phases = []``: ``verify_models`` is transformer-shaped and does 

32 not meaningfully cover SSM hybrids. Correctness is gated by the integration 

33 parity test instead (the standard set for Mamba/Mamba2/NemotronH). 

34- Generation uses the **standard transformer KV-cache path**, not the bridge's 

35 stateful-Mamba loop. Falcon-H1's HF forward carries both attention KV and the 

36 Mamba conv/recurrent state inside one unified ``past_key_values`` cache, so the 

37 default path matches HF bit-for-bit. Setting ``is_stateful`` would instead 

38 route generation through the pure-Mamba ``cache_params`` convention, which this 

39 hybrid does not use, and diverges from HF after the first decode step. 

40""" 

41 

42from typing import Any 

43 

44import torch 

45 

46from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter 

47from transformer_lens.model_bridge.generalized_components import ( 

48 BlockBridge, 

49 DepthwiseConv1DBridge, 

50 EmbeddingBridge, 

51 GatedRMSNormBridge, 

52 LinearBridge, 

53 PositionEmbeddingsAttentionBridge, 

54 RMSNormalizationBridge, 

55 RotaryEmbeddingBridge, 

56 SSM2MixerBridge, 

57 UnembeddingBridge, 

58) 

59from transformer_lens.model_bridge.generalized_components.base import ( 

60 GeneralizedComponent, 

61) 

62 

63 

64def _make_optional(component: GeneralizedComponent) -> GeneralizedComponent: 

65 """Mark a GeneralizedComponent submodule as optional. 

66 

67 Some bridge classes (e.g. ``GatedRMSNormBridge``) do not forward ``optional`` 

68 through their own ``__init__`` even though ``GeneralizedComponent`` supports 

69 it. Setting the attribute directly is safe because ``component_setup.py`` 

70 reads ``getattr(submodule, "optional", False)`` at setup time. 

71 """ 

72 component.optional = True 

73 return component 

74 

75 

76class FalconH1ArchitectureAdapter(ArchitectureAdapter): 

77 """Architecture adapter for ``FalconH1ForCausalLM``. 

78 

79 Parallel hybrid: every block runs GQA attention and a Mamba-2 mixer side by 

80 side, then a SwiGLU MLP. Both branches are mapped on every block so each 

81 sub-path is independently hookable for ablation studies. 

82 """ 

83 

84 _testing_hybrid = True 

85 _testing_eager = None 

86 

87 # verify_models is transformer-shaped and would need a dedicated refactor to 

88 # cover SSM hybrids. Forward-pass correctness lives in the integration test: 

89 # tests/integration/model_bridge/test_falcon_h1_adapter.py 

90 applicable_phases: list[int] = [] 

91 

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

93 super().__init__(cfg) 

94 

95 self._set_rms_rotary_defaults() 

96 # SwiGLU feed-forward (gate_proj / up_proj / down_proj, silu activation). 

97 # FalconH1RMSNorm stores its epsilon as `variance_epsilon` (like Llama). 

98 setattr(self.cfg, "eps_attr", "variance_epsilon") 

99 

100 # Mamba-2 dimensional config. Falcon-H1 specifies the inner SSM width 

101 # directly via `mamba_d_ssm` (= mamba_n_heads * mamba_d_head), unlike 

102 # Mamba2's `expand`-derived intermediate size. conv_dim follows the 

103 # Mamba-2 layout: inner width + 2 groups of state for B and C. 

104 mamba_d_ssm = getattr(cfg, "mamba_d_ssm", self.cfg.d_model) 

105 mamba_n_groups = getattr(cfg, "mamba_n_groups", 1) 

106 mamba_d_state = getattr(cfg, "mamba_d_state", 128) 

107 mamba_n_heads = getattr(cfg, "mamba_n_heads", 0) 

108 conv_dim = mamba_d_ssm + 2 * mamba_n_groups * mamba_d_state 

109 setattr(self.cfg, "mamba_intermediate_size", mamba_d_ssm) 

110 setattr(self.cfg, "conv_dim", conv_dim) 

111 # HF fuses gate, hidden_BC, and dt into one in_proj output; stored so a 

112 # future HF layout change is caught by the integration test. 

113 setattr( 

114 self.cfg, "expected_in_proj_out_features", 2 * mamba_d_ssm + conv_dim + mamba_n_heads 

115 ) 

116 

117 # Raw (passthrough) mode delegates the full forward to HF, which applies 

118 # all of Falcon-H1's scalar multipliers natively — no weight folding for 

119 # forward parity. 

120 self.weight_processing_conversions = {} 

121 

122 self.component_mapping = { 

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

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

125 "blocks": BlockBridge( 

126 name="model.layers", 

127 submodules={ 

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

129 "ln2": RMSNormalizationBridge(name="pre_ff_layernorm", config=self.cfg), 

130 "attn": PositionEmbeddingsAttentionBridge( 

131 name="self_attn", 

132 config=self.cfg, 

133 submodules={ 

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

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

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

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

138 }, 

139 requires_attention_mask=True, 

140 requires_position_embeddings=True, 

141 ), 

142 "mamba": SSM2MixerBridge( 

143 name="mamba", 

144 config=self.cfg, 

145 submodules={ 

146 "in_proj": LinearBridge(name="in_proj"), 

147 "conv1d": DepthwiseConv1DBridge(name="conv1d"), 

148 # HF names the inner gated norm "norm"; TL aliases it 

149 # to "inner_norm" to disambiguate from the block norm. 

150 # Absent when mamba_rms_norm=false, so optional. 

151 "inner_norm": _make_optional(GatedRMSNormBridge(name="norm")), 

152 "out_proj": LinearBridge(name="out_proj"), 

153 }, 

154 ), 

155 "mlp": self._gated_mlp(name="feed_forward"), 

156 }, 

157 ), 

158 "ln_final": RMSNormalizationBridge(name="model.final_layernorm", config=self.cfg), 

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

160 } 

161 

162 def apply_output_logits_transform(self, logits: torch.Tensor) -> torch.Tensor: 

163 """Match Falcon-H1's post-unembedding multiplier.""" 

164 multiplier = float(getattr(self.cfg, "lm_head_multiplier", 1.0)) 

165 return super().apply_output_logits_transform(logits * multiplier)