Add Lfm2 support

transformer_lens/benchmarks/component_outputs.py

@@ -606,7 +606,12 @@ def _test_component(
             # This is needed for model-specific inputs like position_embeddings or attention_mask
             shared_inputs = None
             if (
-                ("attn" in component_path or "mlp" in component_path or "rotary" in component_path)
+                (
+                    "attn" in component_path
+                    or "mlp" in component_path
+                    or "rotary" in component_path
+                    or "conv" in component_path
+                )
                 and hasattr(bridge_component, "get_random_inputs")
                 and callable(getattr(bridge_component, "get_random_inputs"))
             ):

transformer_lens/factories/architecture_adapter_factory.py

@@ -44,6 +44,7 @@
     HubertArchitectureAdapter,
     HunYuanDenseV1ArchitectureAdapter,
     InternLM2ArchitectureAdapter,
+    Lfm2ArchitectureAdapter,
     Lfm2MoeArchitectureAdapter,
     LLaDAArchitectureAdapter,
     LlamaArchitectureAdapter,
@@ -138,6 +139,7 @@
     "LlavaForConditionalGeneration": LlavaArchitectureAdapter,
     "LlavaNextForConditionalGeneration": LlavaNextArchitectureAdapter,
     "LlavaOnevisionForConditionalGeneration": LlavaOnevisionArchitectureAdapter,
+    "Lfm2ForCausalLM": Lfm2ArchitectureAdapter,
     "Lfm2MoeForCausalLM": Lfm2MoeArchitectureAdapter,
     "Mamba2ForCausalLM": Mamba2ArchitectureAdapter,
     "MambaForCausalLM": MambaArchitectureAdapter,

transformer_lens/model_bridge/generalized_components/__init__.py

@@ -64,6 +64,9 @@
 from transformer_lens.model_bridge.generalized_components.joint_qkv_position_embeddings_attention import (
     JointQKVPositionEmbeddingsAttentionBridge,
 )
+from transformer_lens.model_bridge.generalized_components.lfm2_gated_short_conv import (
+    Lfm2ShortConvBridge,
+)
 from transformer_lens.model_bridge.generalized_components.linear import LinearBridge
 from transformer_lens.model_bridge.generalized_components.mla_attention import (
     MLAAttentionBridge,
@@ -127,6 +130,7 @@
     "AudioFeatureExtractorBridge",
     "BlockBridge",
     "DelegatedAttentionBlockBridge",
+    "Lfm2ShortConvBridge",
     "MLABlockBridge",
     "ParallelBlockBridge",
     "BloomBlockBridge",

transformer_lens/model_bridge/generalized_components/depthwise_conv1d.py

@@ -34,6 +34,12 @@ def forward(self, input: torch.Tensor, *args: Any, **kwargs: Any) -> torch.Tenso
         output = self.hook_out(output)
         return output
 
+    def get_random_inputs(self, batch_size=2, seq_len=8, device=None, dtype=None):
+        device = device or torch.device("cpu")
+        dtype = dtype or torch.float32
+        channels = self.original_component.in_channels  # exact, from the wrapped Conv1d
+        return {"args": (torch.randn(batch_size, channels, seq_len, device=device, dtype=dtype),)}
+
     def __repr__(self) -> str:
         if self.original_component is not None:
             try:

transformer_lens/model_bridge/generalized_components/lfm2_gated_short_conv.py

@@ -0,0 +1,29 @@
+"""LiquidAI LFM2 gated short-convolution mixer bridge."""
+
+from transformer_lens.model_bridge.generalized_components.base import (
+    GeneralizedComponent,
+)
+
+
+class Lfm2ShortConvBridge(GeneralizedComponent):
+    """Wrapper around LFM2's double-gated short-convolution mixer.
+
+    Delegates the forward to HF's ``Lfm2ShortConv`` (preserving its fast CUDA /
+    slow PyTorch dispatch and cache handling) and hooks the residual-stream
+    input/output. Inner in_proj / conv / out_proj are spliced in as submodules,
+    so their hooks fire during HF's own forward.
+
+    Decode-step caveat: on stateful generation HF's conv path reads
+    ``conv.weight`` directly instead of calling ``self.conv(...)``, so
+    ``conv.hook_out`` fires only on prefill — see DepthwiseConv1DBridge.
+
+    CUDA caveat: Hooks surrounding the conv1D operation only fire on the hf
+    "slow path" i.e. if not on CUDA / fast path not available / torch dynamo
+    compiling.
+    """
+
+    hook_aliases = {
+        "hook_in_proj": "in.hook_out",
+        "hook_conv": "conv.hook_out",
+        "hook_gated": "out.hook_in",
+    }

transformer_lens/model_bridge/supported_architectures/__init__.py

@@ -103,6 +103,9 @@
 from transformer_lens.model_bridge.supported_architectures.internlm2 import (
     InternLM2ArchitectureAdapter,
 )
+from transformer_lens.model_bridge.supported_architectures.lfm2 import (
+    Lfm2ArchitectureAdapter,
+)
 from transformer_lens.model_bridge.supported_architectures.lfm2_moe import (
     Lfm2MoeArchitectureAdapter,
 )
@@ -275,6 +278,7 @@
     "LlavaArchitectureAdapter",
     "LlavaNextArchitectureAdapter",
     "LlavaOnevisionArchitectureAdapter",
+    "Lfm2ArchitectureAdapter",
     "Lfm2MoeArchitectureAdapter",
     "MambaArchitectureAdapter",
     "Mamba2ArchitectureAdapter",

transformer_lens/model_bridge/supported_architectures/lfm2.py

@@ -0,0 +1,124 @@
+"""Lfm2 architecture adapter."""
+
+from typing import Any
+
+from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter
+from transformer_lens.model_bridge.generalized_components import (
+    BlockBridge,
+    DepthwiseConv1DBridge,
+    EmbeddingBridge,
+    GatedMLPBridge,
+    Lfm2ShortConvBridge,
+    LinearBridge,
+    PositionEmbeddingsAttentionBridge,
+    RMSNormalizationBridge,
+    RotaryEmbeddingBridge,
+    UnembeddingBridge,
+)
+
+
+class Lfm2ArchitectureAdapter(ArchitectureAdapter):
+    """Architecture adapter for Lfm2 models."""
+
+    def __init__(self, cfg: Any) -> None:
+        """Initialize the Lfm2 architecture adapter."""
+        super().__init__(cfg)
+
+        self.cfg.normalization_type = "RMS"
+        self.cfg.positional_embedding_type = "rotary"
+        self.cfg.final_rms = True
+        self.cfg.gated_mlp = True
+        self.cfg.attn_only = False
+        self.cfg.uses_rms_norm = True
+        self.cfg.act_fn = "silu"
+
+        self.cfg.attn_implementation = "eager"
+
+        if hasattr(cfg, "n_key_value_heads") and cfg.n_key_value_heads is not None:
+            self.cfg.n_key_value_heads = cfg.n_key_value_heads
+
+        self.weight_processing_conversions = {
+            **self._qkvo_weight_conversions(),
+        }
+
+        self.component_mapping = {
+            "embed": EmbeddingBridge(name="model.embed_tokens"),
+            "rotary_emb": RotaryEmbeddingBridge(name="model.rotary_emb"),
+            "blocks": BlockBridge(
+                name="model.layers",
+                submodules={
+                    "ln1": RMSNormalizationBridge(
+                        name="operator_norm",
+                        config=self.cfg,
+                    ),
+                    "ln2": RMSNormalizationBridge(
+                        name="ffn_norm",
+                        config=self.cfg,
+                    ),
+                    "attn": PositionEmbeddingsAttentionBridge(
+                        name="self_attn",
+                        config=self.cfg,
+                        optional=True,
+                        submodules={
+                            "q": LinearBridge(name="q_proj"),
+                            "k": LinearBridge(name="k_proj"),
+                            "v": LinearBridge(name="v_proj"),
+                            "o": LinearBridge(name="out_proj"),
+                            "q_norm": RMSNormalizationBridge(name="q_layernorm", config=self.cfg),
+                            "k_norm": RMSNormalizationBridge(name="k_layernorm", config=self.cfg),
+                        },
+                        requires_attention_mask=True,
+                        requires_position_embeddings=True,
+                    ),
+                    "conv": Lfm2ShortConvBridge(
+                        name="conv",
+                        config=self.cfg,
+                        optional=True,
+                        submodules={
+                            "in": LinearBridge(name="in_proj"),
+                            "conv": DepthwiseConv1DBridge(name="conv"),
+                            "out": LinearBridge(name="out_proj"),
+                        },
+                    ),
+                    "mlp": GatedMLPBridge(
+                        name="feed_forward",
+                        config=self.cfg,
+                        submodules={
+                            "gate": LinearBridge(name="w1"),
+                            "in": LinearBridge(name="w3"),
+                            "out": LinearBridge(name="w2"),
+                        },
+                    ),
+                },
+            ),
+            "ln_final": RMSNormalizationBridge(name="model.embedding_norm", config=self.cfg),
+            "unembed": UnembeddingBridge(name="lm_head", config=self.cfg),
+        }
+
+    def setup_component_testing(self, hf_model: Any, bridge_model: Any = None) -> None:
+        """Set up model-specific references for component testing."""
+        # Get rotary embedding instance from the HF model
+        rotary_emb = hf_model.model.rotary_emb
+
+        # Set attention implementation on HF model to eager (vs sdpa default)
+        if hasattr(hf_model, "config") and hasattr(hf_model.config, "_attn_implementation"):
+            hf_model.config._attn_implementation = "eager"
+
+        if hasattr(hf_model, "model") and hasattr(hf_model.model, "layers"):
+            for layer in hf_model.model.layers:
+                if hasattr(layer, "self_attn") and hasattr(layer.self_attn, "config"):
+                    layer.self_attn.config._attn_implementation = "eager"
+
+        # Set rotary_emb on actual bridge instances
+        if bridge_model is not None and hasattr(bridge_model, "blocks"):
+            for block in bridge_model.blocks:
+                if hasattr(block, "attn"):
+                    block.attn.set_rotary_emb(rotary_emb)
+
+        # Set on template for get_generalized_component() calls
+        # Find the first attention layer (LFM2 layer 0 is conv, not attn)
+        layer_types = getattr(self.cfg, "layer_types", None)
+        if layer_types is not None and "full_attention" in layer_types:
+            first_attn_idx = layer_types.index("full_attention")
+            attn_bridge = self.get_generalized_component(f"blocks.{first_attn_idx}.attn")
+            attn_bridge.set_rotary_emb(rotary_emb)

transformer_lens/tools/model_registry/__init__.py

@@ -89,6 +89,7 @@
     "LlavaForConditionalGeneration",
     "LlavaNextForConditionalGeneration",
     "LlavaOnevisionForConditionalGeneration",
+    "Lfm2ForCausalLM",
     "Lfm2MoeForCausalLM",
     "MambaForCausalLM",
     "Mamba2ForCausalLM",
@@ -170,6 +171,7 @@
     "LlavaForConditionalGeneration": ["llava-hf"],
     "LlavaNextForConditionalGeneration": ["llava-hf"],
     "LlavaOnevisionForConditionalGeneration": ["llava-hf"],
+    "Lfm2ForCausalLM": ["LiquidAI"],
     "Lfm2MoeForCausalLM": ["LiquidAI"],
     "Mamba2ForCausalLM": ["state-spaces"],
     "MambaForCausalLM": ["state-spaces"],

transformer_lens/tools/model_registry/data/supported_models.json

@@ -6,10 +6,52 @@
     "min_downloads": 500,
     "scan_duration_seconds": 8.1
   },
-  "total_architectures": 72,
-  "total_models": 13147,
-  "total_verified": 1030,
+  "total_architectures": 73,
+  "total_models": 13150,
+  "total_verified": 1031,
   "models": [
+    {
+      "architecture_id": "Lfm2ForCausalLM",
+      "model_id": "LiquidAI/LFM2.5-230M",
+      "status": 3,
+      "verified_date": "2026-07-12",
+      "metadata": null,
+      "note": "Below threshold: P1=50.0% < 100.0% (failed: all_components); P3=90.0% but required tests failed: log \u2014 8/98 components failed (8 critical)",
+      "phase1_score": 50.0,
+      "phase2_score": 100.0,
+      "phase3_score": 90.0,
+      "phase4_score": 98.8,
+      "phase7_score": null,
+      "phase8_score": null
+    },
+    {
+      "architecture_id": "Lfm2ForCausalLM",
+      "model_id": "LiquidAI/LFM2-350M",
+      "status": 3,
+      "verified_date": "2026-07-12",
+      "metadata": null,
+      "note": "Below threshold: P1=50.0% < 100.0% (failed: all_components); P3=90.0% but required tests failed: log \u2014 10/114 components failed (10 critical)",
+      "phase1_score": 50.0,
+      "phase2_score": 100.0,
+      "phase3_score": 90.0,
+      "phase4_score": 98.7,
+      "phase7_score": null,
+      "phase8_score": null
+    },
+    {
+      "architecture_id": "Lfm2ForCausalLM",
+      "model_id": "LiquidAI/LFM2-1.2B",
+      "status": 1,
+      "verified_date": "2026-07-15",
+      "metadata": null,
+      "note": "Full verification completed",
+      "phase1_score": 100.0,
+      "phase2_score": 100.0,
+      "phase3_score": 100.0,
+      "phase4_score": 99.3,
+      "phase7_score": null,
+      "phase8_score": null
+    },
     {
       "architecture_id": "FalconH1ForCausalLM",
       "model_id": "tiiuae/Falcon-H1-Tiny-90M-Instruct",