diff --git a/transformer_lens/model_bridge/supported_architectures/lfm2_moe.py b/transformer_lens/model_bridge/supported_architectures/lfm2_moe.py index 496b1527a..ac2ef8d42 100644 --- a/transformer_lens/model_bridge/supported_architectures/lfm2_moe.py +++ b/transformer_lens/model_bridge/supported_architectures/lfm2_moe.py @@ -1,76 +1,164 @@ """LiquidAI LFM2 MoE architecture adapter.""" -from typing import Any +from typing import Any, Dict, Optional + +import torch from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter from transformer_lens.model_bridge.generalized_components import ( BlockBridge, + DepthwiseConv1DBridge, EmbeddingBridge, + Lfm2ShortConvBridge, + LinearBridge, + MoEBridge, + MoERouterBridge, + PositionEmbeddingsAttentionBridge, RMSNormalizationBridge, + RotaryEmbeddingBridge, UnembeddingBridge, ) -class Lfm2MoeBlockBridge(BlockBridge): - """Whole-layer LFM2 bridge exposing only residual stream hooks. - - LFM2 MoE interleaves short-convolution and full-attention operator layers. - Wrapping the HF layer as a whole preserves correct execution while avoiding - unresolved standard attention/MLP aliases on layers that do not have them. - """ - - hook_aliases = { - "hook_resid_pre": "hook_in", - "hook_resid_post": "hook_out", - } +class Lfm2MoeGateBridge(MoERouterBridge): + def get_random_inputs( + self, + batch_size: int = 2, + seq_len: int = 8, + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, + ) -> Dict[str, Any]: + """Random inputs for router component testing. + + The router runs on the reshaped [N, d_model] hidden states and takes a + second `expert_bias` arg (use_expert_bias=True); its top-k gather is + hardcoded to dim=1, so the input must be 2D or the gather indexes the + sequence axis out of bounds. + + Args: + batch_size: Batch size for generated inputs + seq_len: Sequence length for generated inputs + device: Device to place tensors on + dtype: Dtype for generated tensors (defaults to float32) + + Returns: + Dictionary of input tensors matching the component's expected input signature + """ + if device is None: + device = torch.device("cpu") + if dtype is None: + dtype = torch.float32 + d_model = self.config.d_model if self.config and hasattr(self.config, "d_model") else 768 + num_experts = ( + self.config.num_experts if self.config and hasattr(self.config, "num_experts") else 0 + ) + hidden_states = torch.randn(batch_size * seq_len, d_model, device=device, dtype=dtype) + expert_bias = torch.zeros(num_experts, device=device) + return {"args": (hidden_states, expert_bias)} class Lfm2MoeArchitectureAdapter(ArchitectureAdapter): - """Architecture adapter for LiquidAI LFM2 MoE models. - - LFM2 MoE is a hybrid decoder with both short-convolution and full-attention - layers. The adapter delegates each decoder layer to HF and exposes residual - hooks around the whole layer rather than pretending every layer has a - homogeneous attention/MLP substructure. - """ - - # Phases 1-3 compare standard attention/MLP components, which this hybrid - # adapter intentionally doesn't expose (whole-layer residual hooks only). - # Phase 4 (generation + text-quality) needs no component comparison, so it applies. - applicable_phases: list[int] = [4] + """Architecture adapter for LiquidAI Lfm2 MoE models.""" def __init__(self, cfg: Any) -> None: - """Initialize the LFM2 MoE architecture adapter.""" + """Initialize the Lfm2 MoE architecture adapter.""" super().__init__(cfg) self._set_rms_rotary_defaults() - # Hookable attention needs eager; the base prepare hooks force it through - # from_pretrained and onto the loaded config. + + self.cfg.act_fn = "silu" self.cfg.attn_implementation = "eager" self.cfg.default_prepend_bos = False - if hasattr(cfg, "num_experts"): - self.cfg.num_experts = cfg.num_experts - if hasattr(cfg, "experts_per_token"): - self.cfg.experts_per_token = cfg.experts_per_token - if hasattr(cfg, "moe_intermediate_size"): - setattr(self.cfg, "moe_intermediate_size", cfg.moe_intermediate_size) - if hasattr(cfg, "layer_types"): - setattr(self.cfg, "layer_types", cfg.layer_types) - - norm_eps = getattr(cfg, "norm_eps", None) - if norm_eps is not None: - self.cfg.eps = norm_eps - rope_parameters = getattr(cfg, "rope_parameters", None) or {} rope_theta = rope_parameters.get("rope_theta") or getattr(cfg, "rope_theta", None) if rope_theta is not None: self.cfg.rotary_base = rope_theta + self.weight_processing_conversions = { + **self._qkvo_weight_conversions(), + } + self.component_mapping = { "embed": EmbeddingBridge(name="model.embed_tokens"), - "blocks": Lfm2MoeBlockBridge(name="model.layers", config=self.cfg), - # LFM2 stores the decoder-final norm at embedding_norm, not model.norm. + "rotary_emb": RotaryEmbeddingBridge(name="model.pos_emb"), + "blocks": BlockBridge( + name="model.layers", + config=self.cfg, + 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": MoEBridge( + name="feed_forward", + config=self.cfg, + sparse_required=("gate",), + submodules={ + "gate": Lfm2MoeGateBridge(name="gate", config=self.cfg, optional=True), + "dense_gate": LinearBridge(name="w1", optional=True), + "dense_in": LinearBridge(name="w3", optional=True), + "dense_out": LinearBridge(name="w2", optional=True), + }, + ), + }, + ), "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.""" + rotary_emb = hf_model.model.pos_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/data/supported_models.json b/transformer_lens/tools/model_registry/data/supported_models.json index c20df5eb9..6a3dd2564 100644 --- a/transformer_lens/tools/model_registry/data/supported_models.json +++ b/transformer_lens/tools/model_registry/data/supported_models.json @@ -9,7 +9,7 @@ "total_architectures": 143, "total_models": 15670, "total_provisional": 7, - "total_verified": 1203, + "total_verified": 1202, "models": [ { "architecture_id": "FalconH1ForCausalLM", @@ -30390,16 +30390,17 @@ { "architecture_id": "Lfm2MoeForCausalLM", "model_id": "LiquidAI/LFM2.5-8B-A1B", - "status": 1, - "verified_date": "2026-06-26", + "status": 3, + "verified_date": "2026-08-15", "metadata": null, - "note": "Full verification completed with issues, low text quality", - "phase1_score": null, - "phase2_score": null, - "phase3_score": null, - "phase4_score": 23.6, + "note": "Below threshold: P3=85.0% but required tests failed: logits_equivalence, loss_equivalence \u2014 Text quality score: 24.7/100 (avg perplexity: 19.6) \u2014 generated text may be incoherent", + "phase1_score": 100.0, + "phase2_score": 100.0, + "phase3_score": 85.0, + "phase4_score": 24.7, "phase7_score": null, - "phase8_score": null + "phase8_score": null, + "phase9_score": null }, { "architecture_id": "LlamaForCausalLM", diff --git a/transformer_lens/tools/model_registry/data/verification_history.json b/transformer_lens/tools/model_registry/data/verification_history.json index a80a94f25..0ae85e7f1 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-31T10:36:39.286088", + "last_updated": "2026-08-15T07:23:13.344517", "records": [ { "model_id": "Macropodus/macbert4mdcspell_v1", @@ -21060,6 +21060,26 @@ "notes": "Full verification completed", "invalidated": false, "invalidation_reason": null + }, + { + "model_id": "LiquidAI/LFM2.5-8B-A1B", + "architecture_id": "Lfm2MoeForCausalLM", + "verified_date": "2026-08-14", + "verified_by": "verify_models", + "transformerlens_version": null, + "notes": "Below threshold: P1=50.0% < 100.0% (failed: all_components); P3=38.9% < 75.0% (failed: layer_norm_fo \u2014 22/152 components failed (22 critical)", + "invalidated": false, + "invalidation_reason": null + }, + { + "model_id": "LiquidAI/LFM2.5-8B-A1B", + "architecture_id": "Lfm2MoeForCausalLM", + "verified_date": "2026-08-15", + "verified_by": "verify_models", + "transformerlens_version": null, + "notes": "Below threshold: P3=85.0% but required tests failed: logits_equivalence, loss_equivalence \u2014 Text quality score: 24.7/100 (avg perplexity: 19.6) \u2014 generated text may be incoherent", + "invalidated": false, + "invalidation_reason": null } ] }