transformer_lens.model_bridge.supported_architectures.falcon_h1 module

Falcon-H1 parallel-hybrid architecture adapter.

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

Architecture overview: - Every block runs GQA attention and a Mamba-2 mixer in parallel on the same

input_layernorm output; their results are summed into the residual stream alongside scalar multipliers. A second norm (pre_ff_layernorm) precedes a SwiGLU feed-forward MLP. This is the parallel hybrid that distinguishes Falcon-H1 from heterogeneous-layer hybrids like NemotronH (one mixer type per layer) — here both branches are present in every block.

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

  • mamba_rms_norm=false on the released checkpoints — the Mamba inner gated RMSNorm is absent, so inner_norm is declared optional.

  • ~12 scalar multipliers (embedding_multiplier, attention_out_multiplier, key_multiplier, ssm_*_multiplier, mlp_multipliers, lm_head_multiplier …). HF applies all of these in its own forward, so in raw (passthrough) mode the bridge inherits them for free — no weight folding is needed for forward parity. Folding would only matter for compatibility mode.

Key adapter decisions: - BlockBridge is the block container: the block has two norms

(input_layernorm + pre_ff_layernorm) and a real post-attention residual, so transformer-shaped hooks (hook_resid_mid via ln2) are meaningful — unlike single-norm SSM blocks that need SSMBlockBridge. BlockBridge.forward delegates the whole block to HF, so the parallel attn+mamba combine happens natively.

  • SSM2MixerBridge wraps the Mamba-2 mixer (passthrough forward). Its inner in_proj / conv1d / out_proj are mapped for hookability; the gated inner_norm is optional (absent when mamba_rms_norm=false).

  • applicable_phases = []: verify_models is transformer-shaped and does not meaningfully cover SSM hybrids. Correctness is gated by the integration parity test instead (the standard set for Mamba/Mamba2/NemotronH).

  • Generation uses the standard transformer KV-cache path, not the bridge’s stateful-Mamba loop. Falcon-H1’s HF forward carries both attention KV and the Mamba conv/recurrent state inside one unified past_key_values cache, so the default path matches HF bit-for-bit. Setting is_stateful would instead route generation through the pure-Mamba cache_params convention, which this hybrid does not use, and diverges from HF after the first decode step.

class transformer_lens.model_bridge.supported_architectures.falcon_h1.FalconH1ArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for FalconH1ForCausalLM.

Parallel hybrid: every block runs GQA attention and a Mamba-2 mixer side by side, then a SwiGLU MLP. Both branches are mapped on every block so each sub-path is independently hookable for ablation studies.

applicable_phases: list[int] = []
apply_output_logits_transform(logits: Tensor) → Tensor

Match Falcon-H1’s post-unembedding multiplier.

component_mapping: ComponentMapping | None
uses_split_attention: bool
weight_processing_conversions: Dict[str, ParamProcessingConversion | str] | None