diff --git a/tests/unit/model_bridge/supported_architectures/test_lfm2_adapter.py b/tests/unit/model_bridge/supported_architectures/test_lfm2_adapter.py new file mode 100644 index 000000000..b80a67cdd --- /dev/null +++ b/tests/unit/model_bridge/supported_architectures/test_lfm2_adapter.py @@ -0,0 +1,431 @@ +"""Unit tests for Lfm2ArchitectureAdapter. + +Tests cover: +- Config attribute validation +- Component mapping structure +- Weight conversion keys and rearrange patterns +- Architecture guards +- Setup component tests +""" + +from types import SimpleNamespace + +import pytest + +from transformer_lens.config import TransformerBridgeConfig +from transformer_lens.conversion_utils.conversion_steps import RearrangeTensorConversion +from transformer_lens.conversion_utils.param_processing_conversion import ( + ParamProcessingConversion, +) +from transformer_lens.model_bridge.generalized_components import ( + BlockBridge, + DepthwiseConv1DBridge, + EmbeddingBridge, + GatedMLPBridge, + Lfm2ShortConvBridge, + LinearBridge, + PositionEmbeddingsAttentionBridge, + RMSNormalizationBridge, + RotaryEmbeddingBridge, + UnembeddingBridge, +) +from transformer_lens.model_bridge.supported_architectures.lfm2 import ( + Lfm2ArchitectureAdapter, +) + +# --------------------------------------------------------------------------- +# Fixtures & Helpers +# --------------------------------------------------------------------------- + + +def _make_cfg( + n_heads: int = 32, + n_key_value_heads: int = 4, + d_model: int = 128, + n_layers: int = 2, + d_mlp: int = 256, + d_vocab: int = 1000, + n_ctx: int = 512, + layer_types: list[str] = ["conv", "full_attention"], +) -> TransformerBridgeConfig: + """Return a minimal TransformerBridgeConfig for Lfm2 adapter tests.""" + cfg = TransformerBridgeConfig( + d_model=d_model, + d_head=d_model // n_heads, + n_layers=n_layers, + n_ctx=n_ctx, + n_heads=n_heads, + n_key_value_heads=n_key_value_heads, + d_vocab=d_vocab, + d_mlp=d_mlp, + architecture="Lfm2ForCausalLM", + ) + + cfg.layer_types = layer_types + + return cfg + + +@pytest.fixture +def cfg() -> TransformerBridgeConfig: + return _make_cfg() + + +@pytest.fixture +def adapter(cfg: TransformerBridgeConfig) -> Lfm2ArchitectureAdapter: + return Lfm2ArchitectureAdapter(cfg) + + +# For rotary embedding and attention implementation check in setup component testing + +layer_types = ["conv", "full_attention"] + + +def _fake_attn(layer_idx: int) -> SimpleNamespace: + """Per-layer self_attn with a mutable .config so the eager flip is observable.""" + return SimpleNamespace( + config=SimpleNamespace(_attn_implementation="sdpa"), + layer_idx=layer_idx, + ) + + +def _fake_hf_model(rotary_emb: object, n_layers: int = 2) -> SimpleNamespace: + """Stub hf_model exposing everything setup_component_testing walks: + - .model.rotary_emb -> rotary wiring + - .config._attn_implementation -> top-level eager flip + - .model.layers[*].self_attn.config._attn_implementation -> per-layer eager flip + """ + return SimpleNamespace( + config=SimpleNamespace(_attn_implementation="sdpa"), + model=SimpleNamespace( + rotary_emb=rotary_emb, + layers=[ + SimpleNamespace(self_attn=_fake_attn(i)) + for i in range(n_layers) + if layer_types[i] == "full_attention" + ], + ), + ) + + +class DummyAttention: + def __init__(self) -> None: + self.rotary_emb = None + + def set_rotary_emb(self, rotary_emb: object) -> None: + self.rotary_emb = rotary_emb + + +class DummyBlock: + def __init__(self, has_attention: bool = True) -> None: + if has_attention: + self.attn = DummyAttention() + + +class DummyBridgeModel: + def __init__(self, blocks: list[DummyBlock]) -> None: + self.blocks = blocks + + +# --------------------------------------------------------------------------- +# Config attribute tests +# --------------------------------------------------------------------------- + + +class TestLfm2AdapterConfig: + """Adapter must set all required config flags to the values Lfm2 expects.""" + + def test_attn_implementation_is_eager(self, adapter: Lfm2ArchitectureAdapter) -> None: + """Set to eager.""" + assert adapter.cfg.attn_implementation == "eager" + + +# --------------------------------------------------------------------------- +# Component mapping structure tests +# --------------------------------------------------------------------------- + + +class TestLfm2AdapterComponentMapping: + """Component mapping must have the correct bridge types and HF module names.""" + + # -- Top-level keys -- + + def test_embed_is_embedding_bridge(self, adapter: Lfm2ArchitectureAdapter) -> None: + assert isinstance(adapter.component_mapping["embed"], EmbeddingBridge) + + def test_embed_name(self, adapter: Lfm2ArchitectureAdapter) -> None: + assert adapter.component_mapping["embed"].name == "model.embed_tokens" + + def test_rotary_emb_is_rotary_embedding_bridge(self, adapter: Lfm2ArchitectureAdapter) -> None: + assert isinstance(adapter.component_mapping["rotary_emb"], RotaryEmbeddingBridge) + + def test_rotary_emb_name(self, adapter: Lfm2ArchitectureAdapter) -> None: + assert adapter.component_mapping["rotary_emb"].name == "model.rotary_emb" + + def test_blocks_is_block_bridge(self, adapter: Lfm2ArchitectureAdapter) -> None: + """Sequential attn/conv & MLP requires BlockBridge, not ParallelBlockBridge.""" + assert isinstance(adapter.component_mapping["blocks"], BlockBridge) + + def test_blocks_name(self, adapter: Lfm2ArchitectureAdapter) -> None: + assert adapter.component_mapping["blocks"].name == "model.layers" + + def test_ln_final_is_rms_normalization_bridge(self, adapter: Lfm2ArchitectureAdapter) -> None: + assert isinstance(adapter.component_mapping["ln_final"], RMSNormalizationBridge) + + def test_ln_final_name(self, adapter: Lfm2ArchitectureAdapter) -> None: + assert adapter.component_mapping["ln_final"].name == "model.embedding_norm" + + def test_unembed_is_unembedding_bridge(self, adapter: Lfm2ArchitectureAdapter) -> None: + assert isinstance(adapter.component_mapping["unembed"], UnembeddingBridge) + + def test_unembed_name(self, adapter: Lfm2ArchitectureAdapter) -> None: + assert adapter.component_mapping["unembed"].name == "lm_head" + + # -- Block submodules -- + + def test_blocks_ln1_is_rms_normalization_bridge(self, adapter: Lfm2ArchitectureAdapter) -> None: + blocks = adapter.component_mapping["blocks"] + assert isinstance(blocks.submodules["ln1"], RMSNormalizationBridge) + + def test_blocks_ln1_name(self, adapter: Lfm2ArchitectureAdapter) -> None: + blocks = adapter.component_mapping["blocks"] + assert blocks.submodules["ln1"].name == "operator_norm" + + def test_blocks_ln2_is_rms_normalization_bridge(self, adapter: Lfm2ArchitectureAdapter) -> None: + blocks = adapter.component_mapping["blocks"] + assert isinstance(blocks.submodules["ln2"], RMSNormalizationBridge) + + def test_blocks_ln2_name(self, adapter: Lfm2ArchitectureAdapter) -> None: + blocks = adapter.component_mapping["blocks"] + assert blocks.submodules["ln2"].name == "ffn_norm" + + def test_attn_is_position_embeddings_attention_bridge( + self, adapter: Lfm2ArchitectureAdapter + ) -> None: + blocks = adapter.component_mapping["blocks"] + assert isinstance(blocks.submodules["attn"], PositionEmbeddingsAttentionBridge) + + def test_attn_name(self, adapter: Lfm2ArchitectureAdapter) -> None: + blocks = adapter.component_mapping["blocks"] + assert blocks.submodules["attn"].name == "self_attn" + + def test_attn_requires_attention_mask_is_true(self, adapter: Lfm2ArchitectureAdapter) -> None: + blocks = adapter.component_mapping["blocks"] + assert blocks.submodules["attn"].requires_attention_mask is True + + def test_attn_requires_position_embeddings_is_true( + self, adapter: Lfm2ArchitectureAdapter + ) -> None: + blocks = adapter.component_mapping["blocks"] + assert blocks.submodules["attn"].requires_position_embeddings is True + + def test_conv_is_lfm2shortconv_bridge(self, adapter: Lfm2ArchitectureAdapter) -> None: + blocks = adapter.component_mapping["blocks"] + assert isinstance(blocks.submodules["conv"], Lfm2ShortConvBridge) + + def test_conv_name(self, adapter: Lfm2ArchitectureAdapter) -> None: + blocks = adapter.component_mapping["blocks"] + assert blocks.submodules["conv"].name == "conv" + + def test_mlp_is_gated_mlp_bridge(self, adapter: Lfm2ArchitectureAdapter) -> None: + blocks = adapter.component_mapping["blocks"] + assert isinstance(blocks.submodules["mlp"], GatedMLPBridge) + + def test_mlp_name(self, adapter: Lfm2ArchitectureAdapter) -> None: + blocks = adapter.component_mapping["blocks"] + assert blocks.submodules["mlp"].name == "feed_forward" + + # -- Attention submodules -- + + @pytest.mark.parametrize("slot", ["q", "k", "v", "o"]) + def test_attn_submodule_is_linear_bridge( + self, adapter: Lfm2ArchitectureAdapter, slot: str + ) -> None: + attn = adapter.component_mapping["blocks"].submodules["attn"] + assert isinstance(attn.submodules[slot], LinearBridge) + + @pytest.mark.parametrize("slot", ["q_norm", "k_norm"]) + def test_attn_submodule_is_rms_normalization_bridge( + self, adapter: Lfm2ArchitectureAdapter, slot: str + ) -> None: + attn = adapter.component_mapping["blocks"].submodules["attn"] + assert isinstance(attn.submodules[slot], RMSNormalizationBridge) + + @pytest.mark.parametrize( + "slot, hf_name", + [ + ("q", "q_proj"), + ("k", "k_proj"), + ("v", "v_proj"), + ("o", "out_proj"), + ("q_norm", "q_layernorm"), + ("k_norm", "k_layernorm"), + ], + ) + def test_attn_submodule_name( + self, adapter: Lfm2ArchitectureAdapter, slot: str, hf_name: str + ) -> None: + attn = adapter.component_mapping["blocks"].submodules["attn"] + assert attn.submodules[slot].name == hf_name + + # -- Conv submodules -- + + def test_conv_in_is_linear_bridge(self, adapter: Lfm2ArchitectureAdapter) -> None: + conv = adapter.component_mapping["blocks"].submodules["conv"] + assert isinstance(conv.submodules["in"], LinearBridge) + + def test_conv_in_name(self, adapter: Lfm2ArchitectureAdapter) -> None: + conv = adapter.component_mapping["blocks"].submodules["conv"] + assert conv.submodules["in"].name == "in_proj" + + def test_conv_conv_is_depthwise_conv1d_bridge(self, adapter: Lfm2ArchitectureAdapter) -> None: + conv = adapter.component_mapping["blocks"].submodules["conv"] + assert isinstance(conv.submodules["conv"], DepthwiseConv1DBridge) + + def test_conv_conv_name(self, adapter: Lfm2ArchitectureAdapter) -> None: + conv = adapter.component_mapping["blocks"].submodules["conv"] + assert conv.submodules["conv"].name == "conv" + + def test_conv_out_is_linear_bridge(self, adapter: Lfm2ArchitectureAdapter) -> None: + conv = adapter.component_mapping["blocks"].submodules["conv"] + assert isinstance(conv.submodules["out"], LinearBridge) + + def test_conv_out_name(self, adapter: Lfm2ArchitectureAdapter) -> None: + conv = adapter.component_mapping["blocks"].submodules["conv"] + assert conv.submodules["out"].name == "out_proj" + + # -- MLP submodules -- + + def test_mlp_gate_is_linear_bridge(self, adapter: Lfm2ArchitectureAdapter) -> None: + mlp = adapter.component_mapping["blocks"].submodules["mlp"] + assert isinstance(mlp.submodules["gate"], LinearBridge) + + def test_mlp_gate_name(self, adapter: Lfm2ArchitectureAdapter) -> None: + mlp = adapter.component_mapping["blocks"].submodules["mlp"] + assert mlp.submodules["gate"].name == "w1" + + def test_mlp_in_is_linear_bridge(self, adapter: Lfm2ArchitectureAdapter) -> None: + mlp = adapter.component_mapping["blocks"].submodules["mlp"] + assert isinstance(mlp.submodules["in"], LinearBridge) + + def test_mlp_in_name(self, adapter: Lfm2ArchitectureAdapter) -> None: + mlp = adapter.component_mapping["blocks"].submodules["mlp"] + assert mlp.submodules["in"].name == "w3" + + def test_mlp_out_is_linear_bridge(self, adapter: Lfm2ArchitectureAdapter) -> None: + mlp = adapter.component_mapping["blocks"].submodules["mlp"] + assert isinstance(mlp.submodules["out"], LinearBridge) + + def test_mlp_out_name(self, adapter: Lfm2ArchitectureAdapter) -> None: + mlp = adapter.component_mapping["blocks"].submodules["mlp"] + assert mlp.submodules["out"].name == "w2" + + +# --------------------------------------------------------------------------- +# Weight processing conversion tests +# --------------------------------------------------------------------------- + + +class TestLfm2AdapterWeightConversions: + """Adapter must define exactly the four QKVO weight conversions.""" + + def test_conversion_keys_present(self, adapter: Lfm2ArchitectureAdapter) -> None: + """Lfm2 has 4 weight matrices (QKVO) per attention layer""" + assert adapter.weight_processing_conversions.keys() == { + "blocks.{i}.attn.q.weight", + "blocks.{i}.attn.k.weight", + "blocks.{i}.attn.v.weight", + "blocks.{i}.attn.o.weight", + } + + @pytest.mark.parametrize("slot", ["q", "k", "v"]) + def test_qkv_weight_uses_split_heads_pattern( + self, adapter: Lfm2ArchitectureAdapter, slot: str + ) -> None: + conv = adapter.weight_processing_conversions[f"blocks.{{i}}.attn.{slot}.weight"] + expected = adapter.cfg.n_key_value_heads if slot in ["k", "v"] else adapter.cfg.n_heads + assert isinstance(conv, ParamProcessingConversion) + assert isinstance(conv.tensor_conversion, RearrangeTensorConversion) + assert conv.tensor_conversion.pattern == "(n h) m -> n m h" + assert conv.tensor_conversion.axes_lengths["n"] == expected + + def test_o_uses_merge_heads_pattern(self, adapter: Lfm2ArchitectureAdapter) -> None: + conv = adapter.weight_processing_conversions["blocks.{i}.attn.o.weight"] + assert isinstance(conv, ParamProcessingConversion) + assert isinstance(conv.tensor_conversion, RearrangeTensorConversion) + assert conv.tensor_conversion.pattern == "m (n h) -> n h m" + assert conv.tensor_conversion.axes_lengths["n"] == adapter.cfg.n_heads + + +# --------------------------------------------------------------------------- +# Architecture guards +# --------------------------------------------------------------------------- + + +class TestLfm2ArchitectureGuards: + """Guard against accidental introduction of features Lfm2 does not have.""" + + def test_no_pos_embed_component(self, adapter: Lfm2ArchitectureAdapter) -> None: + """Lfm2 uses rotary embeddings, so there is no learned positional embedding.""" + assert "pos_embed" not in adapter.component_mapping + + +# --------------------------------------------------------------------------- +# Setup component testing tests +# --------------------------------------------------------------------------- + + +class TestLfm2SetupComponentTesting: + """setup_component_testing must wire Lfm2's shared rotary embedding into attention bridges.""" + + def test_setup_flips_top_level_attn_implementation_to_eager( + self, adapter: Lfm2ArchitectureAdapter + ) -> None: + """HF reference defaults to sdpa; setup must flip the top-level config to eager.""" + hf = _fake_hf_model(object()) + assert hf.config._attn_implementation == "sdpa" + + adapter.setup_component_testing(hf) + + assert hf.config._attn_implementation == "eager" + + def test_setup_flips_per_layer_attn_implementation_to_eager( + self, adapter: Lfm2ArchitectureAdapter + ) -> None: + """Each already-built attn layer caches its own config; setup must flip all of them.""" + hf = _fake_hf_model(object(), n_layers=2) + assert all(l.self_attn.config._attn_implementation == "sdpa" for l in hf.model.layers) + + adapter.setup_component_testing(hf) + + for layer in hf.model.layers: + assert layer.self_attn.config._attn_implementation == "eager" + + def test_sets_rotary_emb_on_template_attention(self, adapter: Lfm2ArchitectureAdapter) -> None: + rotary_emb = object() + attn_template = adapter.get_generalized_component("blocks.0.attn") + assert isinstance(attn_template, PositionEmbeddingsAttentionBridge) + assert attn_template._rotary_emb is None + + adapter.setup_component_testing(_fake_hf_model(rotary_emb)) + + assert attn_template._rotary_emb is rotary_emb + + def test_sets_rotary_emb_on_each_bridge_model_attention( + self, adapter: Lfm2ArchitectureAdapter + ) -> None: + rotary_emb = object() + bridge_model = DummyBridgeModel([DummyBlock(), DummyBlock(), DummyBlock()]) + + adapter.setup_component_testing(_fake_hf_model(rotary_emb), bridge_model=bridge_model) + + for block in bridge_model.blocks: + assert block.attn.rotary_emb is rotary_emb + + def test_skips_bridge_blocks_without_attention(self, adapter: Lfm2ArchitectureAdapter) -> None: + rotary_emb = object() + bridge_model = DummyBridgeModel([DummyBlock(), DummyBlock(has_attention=False)]) + + adapter.setup_component_testing(_fake_hf_model(rotary_emb), bridge_model=bridge_model) + + assert bridge_model.blocks[0].attn.rotary_emb is rotary_emb diff --git a/transformer_lens/benchmarks/component_outputs.py b/transformer_lens/benchmarks/component_outputs.py index f23315dee..2bdb23bf9 100644 --- a/transformer_lens/benchmarks/component_outputs.py +++ b/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")) ): diff --git a/transformer_lens/factories/architecture_adapter_factory.py b/transformer_lens/factories/architecture_adapter_factory.py index 1821059c8..bfe1b90ff 100644 --- a/transformer_lens/factories/architecture_adapter_factory.py +++ b/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, diff --git a/transformer_lens/model_bridge/generalized_components/__init__.py b/transformer_lens/model_bridge/generalized_components/__init__.py index 07a6b2110..1ab4685eb 100644 --- a/transformer_lens/model_bridge/generalized_components/__init__.py +++ b/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", diff --git a/transformer_lens/model_bridge/generalized_components/depthwise_conv1d.py b/transformer_lens/model_bridge/generalized_components/depthwise_conv1d.py index acb304df2..5f8623225 100644 --- a/transformer_lens/model_bridge/generalized_components/depthwise_conv1d.py +++ b/transformer_lens/model_bridge/generalized_components/depthwise_conv1d.py @@ -1,5 +1,5 @@ """Bridge for Mamba-style depthwise causal Conv1d (distinct from GPT-2's Conv1D linear).""" -from typing import Any +from typing import Any, Optional, cast import torch @@ -34,6 +34,24 @@ 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: int = 2, + seq_len: int = 8, + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, + ): + if self.original_component is None: + raise RuntimeError( + f"Original component not set for {self.name}. " + "Call set_original_component() first." + ) + device = device or torch.device("cpu") + dtype = dtype or torch.float32 + conv = cast(torch.nn.Conv1d, self.original_component) + channels = conv.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: diff --git a/transformer_lens/model_bridge/generalized_components/lfm2_gated_short_conv.py b/transformer_lens/model_bridge/generalized_components/lfm2_gated_short_conv.py new file mode 100644 index 000000000..b5b10c48b --- /dev/null +++ b/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", + } diff --git a/transformer_lens/model_bridge/supported_architectures/__init__.py b/transformer_lens/model_bridge/supported_architectures/__init__.py index 3a9ad2b74..0eac330e2 100644 --- a/transformer_lens/model_bridge/supported_architectures/__init__.py +++ b/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", diff --git a/transformer_lens/model_bridge/supported_architectures/lfm2.py b/transformer_lens/model_bridge/supported_architectures/lfm2.py new file mode 100644 index 000000000..f915d8822 --- /dev/null +++ b/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) diff --git a/transformer_lens/tools/model_registry/__init__.py b/transformer_lens/tools/model_registry/__init__.py index 0816e2516..e73e1f6ea 100644 --- a/transformer_lens/tools/model_registry/__init__.py +++ b/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"], diff --git a/transformer_lens/tools/model_registry/data/supported_models.json b/transformer_lens/tools/model_registry/data/supported_models.json index 946f8909c..65663df2a 100644 --- a/transformer_lens/tools/model_registry/data/supported_models.json +++ b/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": 1033, "models": [ + { + "architecture_id": "Lfm2ForCausalLM", + "model_id": "LiquidAI/LFM2.5-230M", + "status": 1, + "verified_date": "2026-07-16", + "metadata": null, + "note": "Full verification completed", + "phase1_score": 100.0, + "phase2_score": 100.0, + "phase3_score": 100.0, + "phase4_score": 98.8, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Lfm2ForCausalLM", + "model_id": "LiquidAI/LFM2-350M", + "status": 1, + "verified_date": "2026-07-16", + "metadata": null, + "note": "Full verification completed", + "phase1_score": 100.0, + "phase2_score": 100.0, + "phase3_score": 100.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-16", + "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", diff --git a/transformer_lens/tools/model_registry/data/verification_history.json b/transformer_lens/tools/model_registry/data/verification_history.json index eec6f0a61..f9e577ee1 100644 --- a/transformer_lens/tools/model_registry/data/verification_history.json +++ b/transformer_lens/tools/model_registry/data/verification_history.json @@ -1,5 +1,5 @@ { - "last_updated": "2026-07-08T04:16:00.084004", + "last_updated": "2026-07-16T08:45:00.706089", "records": [ { "model_id": "Macropodus/macbert4mdcspell_v1", @@ -16890,6 +16890,126 @@ "notes": "Full verification completed", "invalidated": false, "invalidation_reason": null + }, + { + "model_id": "LiquidAI/LFM2.5-230M", + "architecture_id": "Lfm2ForCausalLM", + "verified_date": "2026-07-12", + "verified_by": "verify_models", + "transformerlens_version": null, + "notes": "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)", + "invalidated": false, + "invalidation_reason": null + }, + { + "model_id": "LiquidAI/LFM2-350M", + "architecture_id": "Lfm2ForCausalLM", + "verified_date": "2026-07-12", + "verified_by": "verify_models", + "transformerlens_version": null, + "notes": "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)", + "invalidated": false, + "invalidation_reason": null + }, + { + "model_id": "LiquidAI/LFM2-1.2B", + "architecture_id": "Lfm2ForCausalLM", + "verified_date": "2026-07-12", + "verified_by": "verify_models", + "transformerlens_version": null, + "notes": "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)", + "invalidated": false, + "invalidation_reason": null + }, + { + "model_id": "LiquidAI/LFM2-1.2B", + "architecture_id": "Lfm2ForCausalLM", + "verified_date": "2026-07-14", + "verified_by": "verify_models", + "transformerlens_version": null, + "notes": "Below threshold: P1=50.0% < 100.0% (failed: all_components); P3=89.5% but required tests failed: log \u2014 10/114 components failed (10 critical)", + "invalidated": false, + "invalidation_reason": null + }, + { + "model_id": "LiquidAI/LFM2-1.2B", + "architecture_id": "Lfm2ForCausalLM", + "verified_date": "2026-07-14", + "verified_by": "verify_models", + "transformerlens_version": null, + "notes": "Below threshold: P1=50.0% < 100.0% (failed: all_components); P3=89.5% but required tests failed: log \u2014 10/114 components failed (10 critical)", + "invalidated": false, + "invalidation_reason": null + }, + { + "model_id": "LiquidAI/LFM2-1.2B", + "architecture_id": "Lfm2ForCausalLM", + "verified_date": "2026-07-14", + "verified_by": "verify_models", + "transformerlens_version": null, + "notes": "Below threshold: P1=50.0% < 100.0% (failed: all_components); P3=89.5% but required tests failed: log \u2014 10/114 components failed (10 critical)", + "invalidated": false, + "invalidation_reason": null + }, + { + "model_id": "LiquidAI/LFM2-1.2B", + "architecture_id": "Lfm2ForCausalLM", + "verified_date": "2026-07-14", + "verified_by": "verify_models", + "transformerlens_version": null, + "notes": "Below threshold: P1=50.0% < 100.0% (failed: all_components); P3=89.5% but required tests failed: log \u2014 20/114 components failed (20 critical)", + "invalidated": false, + "invalidation_reason": null + }, + { + "model_id": "LiquidAI/LFM2-1.2B", + "architecture_id": "Lfm2ForCausalLM", + "verified_date": "2026-07-14", + "verified_by": "verify_models", + "transformerlens_version": null, + "notes": "Below threshold: P3=89.5% but required tests failed: logits_equivalence, loss_equivalence \u2014 Tensors differ: max_diff=23.349285, mean_rel=0.575189", + "invalidated": false, + "invalidation_reason": null + }, + { + "model_id": "LiquidAI/LFM2-1.2B", + "architecture_id": "Lfm2ForCausalLM", + "verified_date": "2026-07-15", + "verified_by": "verify_models", + "transformerlens_version": null, + "notes": "Full verification completed", + "invalidated": false, + "invalidation_reason": null + }, + { + "model_id": "LiquidAI/LFM2-1.2B", + "architecture_id": "Lfm2ForCausalLM", + "verified_date": "2026-07-16", + "verified_by": "verify_models", + "transformerlens_version": null, + "notes": "Full verification completed", + "invalidated": false, + "invalidation_reason": null + }, + { + "model_id": "LiquidAI/LFM2-350M", + "architecture_id": "Lfm2ForCausalLM", + "verified_date": "2026-07-16", + "verified_by": "verify_models", + "transformerlens_version": null, + "notes": "Full verification completed", + "invalidated": false, + "invalidation_reason": null + }, + { + "model_id": "LiquidAI/LFM2.5-230M", + "architecture_id": "Lfm2ForCausalLM", + "verified_date": "2026-07-16", + "verified_by": "verify_models", + "transformerlens_version": null, + "notes": "Full verification completed", + "invalidated": false, + "invalidation_reason": null } ] } diff --git a/transformer_lens/tools/model_registry/generate_report.py b/transformer_lens/tools/model_registry/generate_report.py index f1d91b0f7..cdbea0570 100644 --- a/transformer_lens/tools/model_registry/generate_report.py +++ b/transformer_lens/tools/model_registry/generate_report.py @@ -32,6 +32,7 @@ "GPTNeoXForCausalLM": "EleutherAI's GPT-NeoX architecture used in Pythia models", "GPTJForCausalLM": "EleutherAI's GPT-J 6B parameter model", "LlamaForCausalLM": "Meta's LLaMA architecture, basis for many open models", + "Lfm2ForCausalLM": "LiquidAI's LFM2 hybrid convolution/attention dense model", "LLaDAModelLM": "GSAI's bidirectional masked-diffusion language model (single-forward support)", "Lfm2MoeForCausalLM": "LiquidAI's LFM2 hybrid convolution/attention Mixture of Experts model", "MistralForCausalLM": "Mistral AI's efficient 7B parameter model with sliding window attention",