Coverage for transformer_lens/model_bridge/supported_architectures/lfm2_moe.py: 76%
44 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-01 16:23 +0000
1"""LiquidAI LFM2 MoE architecture adapter."""
3from typing import Any, Dict, Optional
5import torch
7from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter
8from transformer_lens.model_bridge.generalized_components import (
9 BlockBridge,
10 DepthwiseConv1DBridge,
11 EmbeddingBridge,
12 Lfm2ShortConvBridge,
13 LinearBridge,
14 MoEBridge,
15 MoERouterBridge,
16 PositionEmbeddingsAttentionBridge,
17 RMSNormalizationBridge,
18 RotaryEmbeddingBridge,
19 UnembeddingBridge,
20)
23class Lfm2MoeGateBridge(MoERouterBridge):
24 def get_random_inputs(
25 self,
26 batch_size: int = 2,
27 seq_len: int = 8,
28 device: Optional[torch.device] = None,
29 dtype: Optional[torch.dtype] = None,
30 ) -> Dict[str, Any]:
31 """Random inputs for router component testing.
33 The router runs on the reshaped [N, d_model] hidden states and takes a
34 second `expert_bias` arg (use_expert_bias=True); its top-k gather is
35 hardcoded to dim=1, so the input must be 2D or the gather indexes the
36 sequence axis out of bounds.
38 Args:
39 batch_size: Batch size for generated inputs
40 seq_len: Sequence length for generated inputs
41 device: Device to place tensors on
42 dtype: Dtype for generated tensors (defaults to float32)
44 Returns:
45 Dictionary of input tensors matching the component's expected input signature
46 """
47 if device is None:
48 device = torch.device("cpu")
49 if dtype is None:
50 dtype = torch.float32
51 d_model = self.config.d_model if self.config and hasattr(self.config, "d_model") else 768
52 num_experts = (
53 self.config.num_experts if self.config and hasattr(self.config, "num_experts") else 0
54 )
55 hidden_states = torch.randn(batch_size * seq_len, d_model, device=device, dtype=dtype)
56 expert_bias = torch.zeros(num_experts, device=device)
57 return {"args": (hidden_states, expert_bias)}
60class Lfm2MoeArchitectureAdapter(ArchitectureAdapter):
61 """Architecture adapter for LiquidAI Lfm2 MoE models."""
63 def __init__(self, cfg: Any) -> None:
64 """Initialize the Lfm2 MoE architecture adapter."""
65 super().__init__(cfg)
67 self._set_rms_rotary_defaults()
69 self.cfg.act_fn = "silu"
70 self.cfg.attn_implementation = "eager"
72 rope_parameters = getattr(cfg, "rope_parameters", None) or {}
73 rope_theta = rope_parameters.get("rope_theta") or getattr(cfg, "rope_theta", None)
74 if rope_theta is not None:
75 self.cfg.rotary_base = rope_theta
77 self.weight_processing_conversions = {
78 **self._qkvo_weight_conversions(),
79 }
81 self.component_mapping = {
82 "embed": EmbeddingBridge(name="model.embed_tokens"),
83 "rotary_emb": RotaryEmbeddingBridge(name="model.pos_emb"),
84 "blocks": BlockBridge(
85 name="model.layers",
86 config=self.cfg,
87 submodules={
88 "ln1": RMSNormalizationBridge(
89 name="operator_norm",
90 config=self.cfg,
91 ),
92 "ln2": RMSNormalizationBridge(
93 name="ffn_norm",
94 config=self.cfg,
95 ),
96 "attn": PositionEmbeddingsAttentionBridge(
97 name="self_attn",
98 config=self.cfg,
99 optional=True,
100 submodules={
101 "q": LinearBridge(name="q_proj"),
102 "k": LinearBridge(name="k_proj"),
103 "v": LinearBridge(name="v_proj"),
104 "o": LinearBridge(name="out_proj"),
105 "q_norm": RMSNormalizationBridge(name="q_layernorm", config=self.cfg),
106 "k_norm": RMSNormalizationBridge(name="k_layernorm", config=self.cfg),
107 },
108 requires_attention_mask=True,
109 requires_position_embeddings=True,
110 ),
111 "conv": Lfm2ShortConvBridge(
112 name="conv",
113 config=self.cfg,
114 optional=True,
115 submodules={
116 "in": LinearBridge(name="in_proj"),
117 "conv": DepthwiseConv1DBridge(name="conv"),
118 "out": LinearBridge(name="out_proj"),
119 },
120 ),
121 "mlp": MoEBridge(
122 name="feed_forward",
123 config=self.cfg,
124 sparse_required=("gate",),
125 submodules={
126 "gate": Lfm2MoeGateBridge(name="gate", config=self.cfg, optional=True),
127 "dense_gate": LinearBridge(name="w1", optional=True),
128 "dense_in": LinearBridge(name="w3", optional=True),
129 "dense_out": LinearBridge(name="w2", optional=True),
130 },
131 ),
132 },
133 ),
134 "ln_final": RMSNormalizationBridge(name="model.embedding_norm", config=self.cfg),
135 "unembed": UnembeddingBridge(name="lm_head", config=self.cfg),
136 }
138 def setup_component_testing(self, hf_model: Any, bridge_model: Any = None) -> None:
139 """Set up model-specific references for component testing."""
140 rotary_emb = hf_model.model.pos_emb
142 # Set attention implementation on HF model to eager (vs sdpa default)
143 if hasattr(hf_model, "config") and hasattr(hf_model.config, "_attn_implementation"): 143 ↛ 146line 143 didn't jump to line 146 because the condition on line 143 was always true
144 hf_model.config._attn_implementation = "eager"
146 if hasattr(hf_model, "model") and hasattr(hf_model.model, "layers"): 146 ↛ 152line 146 didn't jump to line 152 because the condition on line 146 was always true
147 for layer in hf_model.model.layers:
148 if hasattr(layer, "self_attn") and hasattr(layer.self_attn, "config"): 148 ↛ 147line 148 didn't jump to line 147 because the condition on line 148 was always true
149 layer.self_attn.config._attn_implementation = "eager"
151 # Set rotary_emb on actual bridge instances
152 if bridge_model is not None and hasattr(bridge_model, "blocks"):
153 for block in bridge_model.blocks:
154 if hasattr(block, "attn"):
155 block.attn.set_rotary_emb(rotary_emb)
157 # Set on template for get_generalized_component() calls
158 # Find the first attention layer (LFM2 layer 0 is conv, not attn)
159 layer_types = getattr(self.cfg, "layer_types", None)
160 if layer_types is not None and "full_attention" in layer_types: 160 ↛ exitline 160 didn't return from function 'setup_component_testing' because the condition on line 160 was always true
161 first_attn_idx = layer_types.index("full_attention")
162 attn_bridge = self.get_generalized_component(f"blocks.{first_attn_idx}.attn")
163 attn_bridge.set_rotary_emb(rotary_emb)