diff --git a/docs/api/build_from_gguf.md b/docs/api/build_from_gguf.md index 85a5bed5e..6466102bf 100644 --- a/docs/api/build_from_gguf.md +++ b/docs/api/build_from_gguf.md @@ -15,7 +15,7 @@ from mobius import build_from_gguf | Census | Total | Closure | |---|---:|---| -| Architectures | 147 | graph verdicts: {'deferred': 91, 'rejected': 3, 'supported': 53}; importable: 51; quantized import: {'rejected': 11, 'supported': 136}; runtime: {'deferred': 144, 'rejected': 3} | +| Architectures | 147 | graph verdicts: {'deferred': 90, 'rejected': 3, 'supported': 54}; importable: 52; quantized import: {'rejected': 12, 'supported': 135}; runtime: {'deferred': 144, 'rejected': 3} | | Active stored qtypes | 25 | 24 have an import route; 1 are explicitly deferred with no route | | Serialized projector strings | 60 | {'graph-importable': 2, 'runtime-supported': 0} | | Tokenizer pre identifiers | 87 | 56 semantic groups; all default to deferred and become exact-copy only with a validated embedded `tokenizer.huggingface.json` | @@ -418,7 +418,7 @@ before graph construction or durable output. | `kimi-linear` | — | none (fails before config extraction) | not claimed | config=deferred; tensor_map=deferred; graph=deferred; runtime=deferred; quantized_import=supported | Kimi-Linear alternates KDA recurrent and NoPE MLA layers, carrying three rolling convolution histories and a per-head matrix state in addition to attention cache. Its two-stage decay/output gates and sigmoid correction-bias MoE routing are not represented by any Mobius graph or state task; aliasing it to Kimi-K3, Mamba, or ordinary attention would change the model. | | `laguna` | — | none (fails before config extraction) | audited-direct-loader-conditional-union | config=deferred; tensor_map=deferred; graph=deferred; runtime=deferred; quantized_import=supported | Laguna combines per-head-or-element softplus attention gates, dual-RoPE interleaved sliding-window attention, a dense prefix, and sigmoid correction-biased routed/shared experts. Mobius has no exact graph or iSWA cache contract. | | `lfm2` | — | model=`lfm2`; tensor=`lfm2` | not claimed | config=supported; tensor_map=supported; graph=supported; runtime=deferred; quantized_import=supported | Config extraction, exact pinned tensor-name closure, GGUF value transforms, and synthetic recurrent-state execution are covered, but no representative real-weight GGUF has yet passed independent full-logit parity and deterministic multi-token stateful ORT generation. Runtime packaging remains deferred until that evidence exists. | -| `lfm2moe` | — | none (fails before config extraction) | not claimed | config=deferred; tensor_map=deferred; graph=deferred; runtime=deferred; quantized_import=supported | LFM2MoE has an arbitrary per-layer attention/short-convolution schedule plus a dense-to-sigmoid-routed-MoE transition. Recurrent layers use F32 rolling convolution state with copy-on-write sequence reorder and bounded rollback snapshots, while attention layers use KV cache. The dense LFM2 graph and ordinary KV/static-cache tasks do not own that mixed state or the expert correction-bias semantics. | +| `lfm2moe` | — | model=`lfm2_moe`; tensor=`lfm2`+`lfm2_moe_extras` | not claimed | config=supported; tensor_map=supported; graph=supported; runtime=deferred; quantized_import=rejected | Config extraction, exact pinned tensor-name closure, GGUF value transforms, and synthetic recurrent-state execution are covered, but no representative real-weight GGUF has yet passed independent full-logit parity and deterministic multi-token stateful ORT generation. Runtime packaging remains deferred until that evidence exists. The mobius graph uses floating Linear modules for this architecture, so no MatMulNBits or BlockQuantizedMatMul target can consume preserved GGUF projection weights. Use keep_quantized=False for explicit float import. | | `llada` | — | model=`llada`; tensor=`llama` | not claimed | config=supported; tensor_map=supported; graph=supported; runtime=deferred; quantized_import=supported | Config extraction, suffix-exact tensor closure, masked-diffusion task dispatch, and synthetic full-sequence execution are covered, but no pinned real GGUF has passed independent Hugging Face/llama.cpp masked-step logit parity and deterministic multi-step generation parity. Runtime packaging remains deferred until both exist. | | `llada-moe` | — | model=`llada`; module=`llada_moe`; tensor=`llama`+`diffusion_fused_qkv`+`moe_qk_norm_extras`+`moe_extras` | not claimed | config=supported; tensor_map=supported; graph=supported; runtime=deferred; quantized_import=supported | Config extraction, suffix-exact tensor closure, masked-diffusion task dispatch, and synthetic full-sequence execution are covered, but no pinned real GGUF has passed independent Hugging Face/llama.cpp masked-step logit parity and deterministic multi-step generation parity. Runtime packaging remains deferred until both exist. | | `llama` | `mistral` | model=`llama`; tensor=`llama` | not claimed | config=supported; tensor_map=supported; graph=supported; runtime=deferred; quantized_import=supported | Config extraction, exact tensor-name closure, and a full synthetic GGUF graph build are covered, but no representative real-weight GGUF has yet passed ORT parity or generation validation. Runtime packaging remains deferred until that evidence exists. | diff --git a/src/mobius/_configs/__init__.py b/src/mobius/_configs/__init__.py index 2fbc36c8f..14a3695f0 100644 --- a/src/mobius/_configs/__init__.py +++ b/src/mobius/_configs/__init__.py @@ -41,6 +41,7 @@ JambaConfig, JetMoeConfig, Lfm2Config, + Lfm2MoeConfig, Lfm2VlConfig, LongcatFlashConfig, Mamba2Config, @@ -117,6 +118,7 @@ "JambaConfig", "JetMoeConfig", "Lfm2Config", + "Lfm2MoeConfig", "Lfm2VlConfig", "LongcatFlashConfig", "Mamba2Config", diff --git a/src/mobius/_configs/_base.py b/src/mobius/_configs/_base.py index 92c31576f..2abdebaca 100644 --- a/src/mobius/_configs/_base.py +++ b/src/mobius/_configs/_base.py @@ -1604,6 +1604,23 @@ def from_transformers(cls, config, parent_config=None) -> Lfm2Config: ) +@dataclasses.dataclass +class Lfm2MoeConfig(CausalLMConfig): + """Configuration for LFM2MoE's dense-prefix and routed-expert feed-forwards.""" + + num_dense_layers: int = 2 + use_expert_bias: bool = True + + @classmethod + def from_transformers(cls, config, parent_config=None) -> Lfm2MoeConfig: + base = ArchitectureConfig.from_transformers(config, parent_config) + return cls( + **_shallow_fields(base), + num_dense_layers=getattr(config, "num_dense_layers", 2), + use_expert_bias=getattr(config, "use_expert_bias", True), + ) + + @dataclasses.dataclass class Lfm2VlConfig(Lfm2Config): """Configuration for LiquidAI LFM2-VL (SigLIP2 NaFlex + LFM2 decoder). diff --git a/src/mobius/_registry.py b/src/mobius/_registry.py index 420a08792..06ff17fd0 100644 --- a/src/mobius/_registry.py +++ b/src/mobius/_registry.py @@ -29,6 +29,7 @@ Gemma4AssistantConfig, Gemma4Config, Lfm2Config, + Lfm2MoeConfig, Lfm2VlConfig, MMSConfig, MoonshineConfig, @@ -80,6 +81,7 @@ InternLM2CausalLMModel, LayerNormCausalLMModel, Lfm2CausalLMModel, + Lfm2MoECausalLMModel, Lfm2VlForConditionalGeneration, LLaDAModel, LLaDAMoEModel, @@ -499,6 +501,7 @@ def _detect_fallback_registration(hf_config) -> ModelRegistration | None: "internlm2": ModelRegistration(InternLM2CausalLMModel), "llama4_text": ModelRegistration(Llama4CausalLMModel), "lfm2": ModelRegistration(Lfm2CausalLMModel, config_class=Lfm2Config), + "lfm2_moe": ModelRegistration(Lfm2MoECausalLMModel, config_class=Lfm2MoeConfig), "dream": ModelRegistration(DreamModel, task="masked-diffusion"), "Dream": ModelRegistration(DreamModel, task="masked-diffusion"), "llada": ModelRegistration(LLaDAModel, task="masked-diffusion"), diff --git a/src/mobius/integrations/gguf/_arch_registry.py b/src/mobius/integrations/gguf/_arch_registry.py index e7702b052..0f0c914a9 100644 --- a/src/mobius/integrations/gguf/_arch_registry.py +++ b/src/mobius/integrations/gguf/_arch_registry.py @@ -322,14 +322,6 @@ "ordinary attention would change the model." ) -_LFM2MOE_GRAPH_REASON = ( - "LFM2MoE has an arbitrary per-layer attention/short-convolution schedule plus a " - "dense-to-sigmoid-routed-MoE transition. Recurrent layers use F32 rolling convolution " - "state with copy-on-write sequence reorder and bounded rollback snapshots, while " - "attention layers use KV cache. The dense LFM2 graph and ordinary KV/static-cache " - "tasks do not own that mixed state or the expert correction-bias semantics." -) - _ENCODER_RUNTIME_VALIDATION_PENDING = ( "Config extraction, exact pinned tensor closure, encoder-only task dispatch, and " "synthetic ORT execution are covered, but no pinned real GGUF artifact has passed " @@ -1138,11 +1130,22 @@ ), GGUFArchitectureSpec( gguf_arch="lfm2moe", - config=Support.DEFERRED, - tensor_map=Support.DEFERRED, - graph=Support.DEFERRED, + model_type="lfm2_moe", + tensor_map_recipe=("lfm2", "lfm2_moe_extras"), + config_postprocessor="lfm2moe", + required_metadata=( + "attention.head_count_kv", + "attention.layer_norm_rms_epsilon", + "rope.freq_base", + "shortconv.l_cache", + "expert_count", + "expert_used_count", + "expert_feed_forward_length", + "expert_gating_func", + ), runtime=Support.DEFERRED, - reason=_LFM2MOE_GRAPH_REASON, + quantized_import=Support.REJECTED, + reason=_RECURRENT_RUNTIME_VALIDATION_PENDING + " " + _NO_QUANTIZED_PROJECTION_REASON, ), GGUFArchitectureSpec( gguf_arch="minimax-01", diff --git a/src/mobius/integrations/gguf/_arch_registry_test.py b/src/mobius/integrations/gguf/_arch_registry_test.py index 5e6d0a487..0907fbbce 100644 --- a/src/mobius/integrations/gguf/_arch_registry_test.py +++ b/src/mobius/integrations/gguf/_arch_registry_test.py @@ -57,7 +57,7 @@ #: Number of importable architectures. Pinned so that adding support is a #: deliberate act that also updates the documented support matrix, and so that #: accidentally losing an architecture is a failure rather than a silence. -_EXPECTED_SUPPORTED_COUNT = 51 +_EXPECTED_SUPPORTED_COUNT = 52 _FINAL_CENSUS_CLOSURE = frozenset( { "afmoe", @@ -268,6 +268,7 @@ def test_every_float_importable_architecture_has_a_quantized_verdict(self) -> No "granitehybrid", "internlm2", "jamba", + "lfm2moe", "mamba", "mamba2", "nemotron_h", @@ -899,14 +900,12 @@ class TestPinnedRemainingHybridCohort: "deepseek4", "kimi-k3", "kimi-linear", - "lfm2moe", ) _EXPECTED_TENSOR_COUNTS: ClassVar[dict[str, int]] = { "bailingmoe3": 41, "deepseek4": 44, "kimi-k3": 46, "kimi-linear": 40, - "lfm2moe": 24, } @pytest.mark.parametrize("architecture", _ARCHITECTURES) @@ -954,7 +953,6 @@ def test_no_unpinned_alias_or_config_mutation_is_reachable( ("deepseek4", ("compressed-cache", "rollback", "ordinary KV")), ("kimi-k3", ("matrix state", "latent MoE", "residual banks")), ("kimi-linear", ("convolution histories", "matrix state", "correction-bias")), - ("lfm2moe", ("rolling convolution", "reorder", "rollback")), ], ) def test_state_and_schedule_mismatch_is_explicit( @@ -965,6 +963,19 @@ def test_state_and_schedule_mismatch_is_explicit( for term in state_terms: assert term in reason + def test_lfm2moe_graph_and_import_advance_but_runtime_stays_deferred(self) -> None: + spec = try_get_arch_spec("lfm2moe") + assert spec is not None + assert spec.model_type == "lfm2_moe" + assert spec.config is Support.SUPPORTED + assert spec.tensor_map is Support.SUPPORTED + assert spec.graph is Support.SUPPORTED + assert spec.runtime is Support.DEFERRED + assert spec.quantized_import is Support.REJECTED + assert spec.reason is not None + assert "representative real-weight GGUF" in spec.reason + assert "keep_quantized=False" in spec.reason + def test_hugging_face_deepseek_v4_registration_remains_valid(self) -> None: assert "deepseek_v4" in _REGISTRATIONS assert try_get_arch_spec("deepseek4").model_type is None diff --git a/src/mobius/integrations/gguf/_builder.py b/src/mobius/integrations/gguf/_builder.py index 1a27d41aa..5249f8be7 100644 --- a/src/mobius/integrations/gguf/_builder.py +++ b/src/mobius/integrations/gguf/_builder.py @@ -906,7 +906,7 @@ def _raise_for_invalid_hybrid_tensor_contract(gguf_model) -> None: if architecture in {"jamba", "nemotron_h", "granitehybrid"}: _raise_for_invalid_mamba_hybrid_tensor_contract(gguf_model) return - if architecture not in {"lfm2", "qwen35", "qwen35moe", "qwen3next"}: + if architecture not in {"lfm2", "lfm2moe", "qwen35", "qwen35moe", "qwen3next"}: return metadata = gguf_model.metadata @@ -934,7 +934,7 @@ def _raise_for_invalid_hybrid_tensor_contract(gguf_model) -> None: f"{match.group(1)} (block_count={total_layers})" ) - if architecture == "lfm2": + if architecture in {"lfm2", "lfm2moe"}: required_global = {"token_embd.weight", "token_embd_norm.weight"} auxiliary = sorted( name @@ -946,13 +946,7 @@ def _raise_for_invalid_hybrid_tensor_contract(gguf_model) -> None: "lfm2 causal-LM import does not support embedding/ColBERT head " f"tensor(s): {auxiliary}" ) - common_suffixes = { - "attn_norm.weight", - "ffn_norm.weight", - "ffn_gate.weight", - "ffn_up.weight", - "ffn_down.weight", - } + common_suffixes = {"attn_norm.weight", "ffn_norm.weight"} full_suffixes = { "attn_q.weight", "attn_k.weight", @@ -1019,6 +1013,30 @@ def _raise_for_invalid_hybrid_tensor_contract(gguf_model) -> None: required.update( full_suffixes if layer_type == "full_attention" else recurrent_suffixes ) + if architecture == "lfm2": + required.update({"ffn_gate.weight", "ffn_up.weight", "ffn_down.weight"}) + elif architecture == "lfm2moe": + dense_ffn = {"ffn_gate.weight", "ffn_up.weight", "ffn_down.weight"} + routed_ffn = { + "ffn_gate_inp.weight", + "ffn_gate_exps.weight", + "ffn_up_exps.weight", + "ffn_down_exps.weight", + "exp_probs_b.bias", + } + dense_layers = int(metadata.get("lfm2moe.leading_dense_block_count", 0)) + if layer < dense_layers: + required.update(dense_ffn) + wrong_ffn = sorted(layer_names & routed_ffn) + else: + required.update(routed_ffn) + wrong_ffn = sorted(layer_names & dense_ffn) + if wrong_ffn: + raise ValueError( + f"lfm2moe layer {layer} contains tensor(s) from the wrong " + f"{'routed' if layer < dense_layers else 'dense'} FFN family: " + f"{wrong_ffn}" + ) if architecture in {"qwen35moe", "qwen3next"}: fused = "ffn_gate_up_exps.weight" in layer_names @@ -1062,7 +1080,7 @@ def _raise_for_invalid_hybrid_tensor_contract(gguf_model) -> None: required.update( {"attn_qkv.weight", "attn_gate.weight"} if modern else {"ssm_in.weight"} ) - elif architecture != "lfm2": + elif architecture not in {"lfm2", "lfm2moe"}: required.update({"attn_qkv.weight", "attn_gate.weight"}) missing = sorted(required - layer_names) @@ -2070,6 +2088,7 @@ def build_from_gguf( ) if gguf_arch in { "lfm2", + "lfm2moe", "qwen35", "qwen35moe", "qwen3next", @@ -2753,28 +2772,42 @@ def _normalize_gguf_weights( # Fused stacked gate/up experts [num_experts, 2*out, ...] are split # before ordinary stacked-expert unpacking. This handles float weights # and packed MatMulNBits companions without dropping either half. - fused_marker = ".mlp.experts.gate_up_proj." - if fused_marker in key and value.dim() >= 3: + fused_marker = next( + ( + marker + for marker in ( + ".mlp.experts.gate_up_proj.", + ".feed_forward.experts.gate_up_proj.", + ) + if marker in key + ), + None, + ) + if fused_marker is not None and value.dim() >= 3: prefix, suffix = key.rsplit(fused_marker, 1) + container = fused_marker.removesuffix(".gate_up_proj.") if value.shape[1] % 2: raise ValueError( f"Fused expert tensor {key!r} has odd gate/up width {value.shape[1]}" ) gate, up = value.chunk(2, dim=1) for i in range(value.shape[0]): - result[f"{prefix}.mlp.experts.{i}.gate_proj.{suffix}"] = gate[i] - result[f"{prefix}.mlp.experts.{i}.up_proj.{suffix}"] = up[i] + result[f"{prefix}{container}.{i}.gate_proj.{suffix}"] = gate[i] + result[f"{prefix}{container}.{i}.up_proj.{suffix}"] = up[i] continue # Stacked expert weights [num_experts, out, ...] → per-expert. unpacked = False for proj in ("gate_proj", "up_proj", "down_proj"): - marker = f".mlp.experts.{proj}." - if marker in key and value.dim() >= 3: - prefix, suffix = key.rsplit(marker, 1) - for i in range(value.shape[0]): - result[f"{prefix}.mlp.experts.{i}.{proj}.{suffix}"] = value[i] - unpacked = True + for container in (".mlp.experts", ".feed_forward.experts"): + marker = f"{container}.{proj}." + if marker in key and value.dim() >= 3: + prefix, suffix = key.rsplit(marker, 1) + for i in range(value.shape[0]): + result[f"{prefix}{container}.{i}.{proj}.{suffix}"] = value[i] + unpacked = True + break + if unpacked: break if unpacked: continue @@ -3960,8 +3993,12 @@ def _validate_moe_weight_shape( if num_experts is None: return expert_size = getattr(config, "moe_intermediate_size", None) or config.intermediate_size - if ".mlp.experts." in name: - projection = name.rsplit(".mlp.experts.", 1)[1].split(".", 1)[0] + expert_marker = next( + (marker for marker in (".mlp.experts.", ".feed_forward.experts.") if marker in name), + None, + ) + if expert_marker is not None: + projection = name.rsplit(expert_marker, 1)[1].split(".", 1)[0] if projection not in {"gate_proj", "up_proj", "down_proj"}: return expected = ( @@ -3973,7 +4010,7 @@ def _validate_moe_weight_shape( raise ValueError( f"Invalid stacked expert shape for {name}: expected {expected}, got {shape}" ) - elif name.endswith(".mlp.gate.weight"): + elif name.endswith((".mlp.gate.weight", ".feed_forward.gate.weight")): expected = (num_experts, config.hidden_size) if shape != expected: raise ValueError( diff --git a/src/mobius/integrations/gguf/_builder_test.py b/src/mobius/integrations/gguf/_builder_test.py index d7b5a3ee6..5a480a79b 100644 --- a/src/mobius/integrations/gguf/_builder_test.py +++ b/src/mobius/integrations/gguf/_builder_test.py @@ -788,6 +788,105 @@ def add_q4(name: str, shape: tuple[int, int]) -> None: writer.close() +def _write_lfm2moe_gguf(path: Path, *, quantized: bool) -> None: + """Write a tiny LFM2MoE GGUF with dense-conv and routed-attention layers.""" + from gguf import GGMLQuantizationType, GGUFWriter + + hidden = 32 + intermediate = 64 + expert_intermediate = 32 + experts = 4 + vocab = 64 + heads = 4 + kv_heads = 2 + head_dim = hidden // heads + kernel = 3 + rng = np.random.default_rng(31) + + writer = GGUFWriter(str(path), "lfm2moe") + writer.add_context_length(64) + writer.add_embedding_length(hidden) + writer.add_feed_forward_length(intermediate) + writer.add_block_count(2) + writer.add_head_count(heads) + writer.add_head_count_kv([0, kv_heads]) + writer.add_rope_freq_base(10_000.0) + writer.add_rope_dimension_count(head_dim) + writer.add_layer_norm_rms_eps(1e-5) + writer.add_vocab_size(vocab) + writer.add_expert_count(experts) + writer.add_expert_used_count(2) + writer.add_expert_feed_forward_length(expert_intermediate) + writer.add_uint32("lfm2moe.shortconv.l_cache", kernel) + writer.add_uint32("lfm2moe.leading_dense_block_count", 1) + writer.add_uint32("lfm2moe.expert_gating_func", 2) + + def add_float(name: str, shape: tuple[int, ...]) -> None: + writer.add_tensor(name, rng.normal(0, 0.03, shape).astype(np.float32)) + + def add_q4(name: str, shape: tuple[int, ...]) -> None: + rows = int(np.prod(shape[:-1])) + columns = shape[-1] + assert columns % 32 == 0 + byte_shape = (*shape[:-1], columns // 32 * 18) + raw = np.zeros((rows, byte_shape[-1]), dtype=np.uint8) + for row in range(rows): + for block in range(columns // 32): + offset = block * 18 + raw[row, offset : offset + 2] = np.array( + [rng.uniform(0.01, 0.05)], dtype=np.float16 + ).view(np.uint8) + raw[row, offset + 2 : offset + 18] = rng.integers( + 0, 256, size=16, dtype=np.uint8 + ) + writer.add_tensor( + name, + raw.reshape(byte_shape), + raw_dtype=GGMLQuantizationType.Q4_0, + ) + + add_projection = add_q4 if quantized else add_float + add_float("token_embd.weight", (vocab, hidden)) + add_float("token_embd_norm.weight", (hidden,)) + for layer in range(2): + prefix = f"blk.{layer}." + add_float(prefix + "attn_norm.weight", (hidden,)) + add_float(prefix + "ffn_norm.weight", (hidden,)) + + add_projection("blk.0.ffn_gate.weight", (intermediate, hidden)) + add_projection("blk.0.ffn_up.weight", (intermediate, hidden)) + add_projection("blk.0.ffn_down.weight", (hidden, intermediate)) + add_float("blk.0.shortconv.conv.weight", (hidden, kernel)) + add_projection("blk.0.shortconv.in_proj.weight", (3 * hidden, hidden)) + add_projection("blk.0.shortconv.out_proj.weight", (hidden, hidden)) + + add_float("blk.1.ffn_gate_inp.weight", (experts, hidden)) + add_projection( + "blk.1.ffn_gate_exps.weight", + (experts, expert_intermediate, hidden), + ) + add_projection( + "blk.1.ffn_up_exps.weight", + (experts, expert_intermediate, hidden), + ) + add_projection( + "blk.1.ffn_down_exps.weight", + (experts, hidden, expert_intermediate), + ) + add_float("blk.1.exp_probs_b.bias", (experts,)) + add_projection("blk.1.attn_q.weight", (heads * head_dim, hidden)) + add_projection("blk.1.attn_k.weight", (kv_heads * head_dim, hidden)) + add_projection("blk.1.attn_v.weight", (kv_heads * head_dim, hidden)) + add_projection("blk.1.attn_output.weight", (hidden, heads * head_dim)) + add_float("blk.1.attn_q_norm.weight", (head_dim,)) + add_float("blk.1.attn_k_norm.weight", (head_dim,)) + + writer.write_header_to_file() + writer.write_kv_data_to_file() + writer.write_tensors_to_file() + writer.close() + + def _write_qwen35_gguf( path: Path, *, @@ -3662,6 +3761,95 @@ def test_lfm2_quantized_preservation_fails_closed_and_float_import_executes( assert outputs[1]["present.1.key"].shape == (1, 2, 4, 8) assert all(np.isfinite(output["logits"]).all() for output in outputs) + def test_lfm2moe_float_prefill_decode_threads_mixed_state(self, tmp_path: Path) -> None: + from mobius._model_package import ModelPackage + from mobius.integrations.gguf import build_from_gguf + + path = tmp_path / "lfm2moe-f32.gguf" + _write_lfm2moe_gguf(path, quantized=False) + package = build_from_gguf(path) + model = package["model"] + + assert [value.name for value in model.graph.inputs if "past_" in value.name] == [ + "past_key_values.0.conv_state", + "past_key_values.1.key", + "past_key_values.1.value", + ] + outputs = self._run_lfm2(model) + assert outputs[0]["present.0.conv_state"].shape == (1, 32, 2) + assert outputs[0]["present.1.key"].shape == (1, 2, 3, 8) + assert outputs[1]["present.1.key"].shape == (1, 2, 4, 8) + assert all(np.isfinite(output["logits"]).all() for output in outputs) + + saved = tmp_path / "saved-lfm2moe" + package.save(str(saved), progress_bar=False, check_weights=True) + reloaded = ModelPackage.load(str(saved))["model"] + reloaded_outputs = self._run_lfm2(reloaded) + for actual, expected in zip(reloaded_outputs, outputs): + for name in actual: + np.testing.assert_allclose(actual[name], expected[name], rtol=0, atol=0) + + def test_lfm2moe_quantized_source_requires_explicit_dequantization( + self, tmp_path: Path + ) -> None: + from mobius.integrations.gguf import build_from_gguf + + path = tmp_path / "lfm2moe-q4.gguf" + _write_lfm2moe_gguf(path, quantized=True) + with pytest.raises(ValueError, match="keep_quantized=False"): + build_from_gguf(path) + + explicit_float = build_from_gguf(path, keep_quantized=False)["model"] + outputs = self._run_lfm2(explicit_float) + assert all(np.isfinite(output["logits"]).all() for output in outputs) + + def test_lfm2moe_state_rollback_and_batch_reorder(self, tmp_path: Path) -> None: + from mobius._testing.ort_inference import OnnxModelSession + from mobius.integrations.gguf import build_from_gguf + + path = tmp_path / "lfm2moe-state.gguf" + _write_lfm2moe_gguf(path, quantized=False) + session = OnnxModelSession(build_from_gguf(path)["model"]) + try: + prefill = session.run( + { + "input_ids": np.asarray([[1, 2], [3, 4]], dtype=np.int64), + "attention_mask": np.ones((2, 2), dtype=np.int64), + "position_ids": np.asarray([[0, 1], [0, 1]], dtype=np.int64), + "past_key_values.0.conv_state": np.zeros((2, 32, 2), dtype=np.float32), + "past_key_values.1.key": np.zeros((2, 2, 0, 8), dtype=np.float32), + "past_key_values.1.value": np.zeros((2, 2, 0, 8), dtype=np.float32), + } + ) + snapshot = { + "past_key_values.0.conv_state": prefill["present.0.conv_state"], + "past_key_values.1.key": prefill["present.1.key"], + "past_key_values.1.value": prefill["present.1.value"], + } + + def decode(tokens: list[list[int]], states: dict[str, np.ndarray]): + return session.run( + { + "input_ids": np.asarray(tokens, dtype=np.int64), + "attention_mask": np.ones((2, 3), dtype=np.int64), + "position_ids": np.asarray([[2], [2]], dtype=np.int64), + **states, + } + ) + + first = decode([[5], [6]], snapshot) + replayed = decode([[5], [6]], snapshot) + reordered = decode( + [[6], [5]], + {name: value[[1, 0]] for name, value in snapshot.items()}, + ) + finally: + session.close() + + for name in first: + np.testing.assert_allclose(replayed[name], first[name], rtol=0, atol=0) + np.testing.assert_allclose(reordered[name], first[name][[1, 0]], rtol=0, atol=0) + def test_qwen35_float_and_quantized_prefill_decode_thread_mixed_state( self, tmp_path: Path ) -> None: @@ -3742,6 +3930,23 @@ def test_lfm2_cache_task_misdispatch_is_rejected( with pytest.raises(ValueError, match=re.escape(message)): build_from_gguf(path, **kwargs) + @pytest.mark.parametrize( + ("kwargs", "message"), + [ + ({"static_cache": True}, "architecture-specific"), + ({"task": "text-generation"}, "hybrid-text-generation"), + ], + ) + def test_lfm2moe_cache_task_misdispatch_is_rejected( + self, tmp_path: Path, kwargs: dict, message: str + ) -> None: + from mobius.integrations.gguf import build_from_gguf + + path = tmp_path / "lfm2moe-misdispatch.gguf" + _write_lfm2moe_gguf(path, quantized=False) + with pytest.raises(ValueError, match=re.escape(message)): + build_from_gguf(path, **kwargs) + class TestRecurrentGGUFBuild: """Mamba GGUF imports preserve recurrent state and tensor-role semantics.""" @@ -4214,7 +4419,6 @@ def test_deferred_audio_architectures_fail_before_all_downstream_stages( "deepseek4", "kimi-k3", "kimi-linear", - "lfm2moe", "arctic", "dbrx", "gpt-oss", @@ -4805,6 +5009,27 @@ def _lfm2_names() -> list[str]: names.extend(f"blk.{layer}.{suffix}" for suffix in [*common, *conditional]) return names + @classmethod + def _lfm2moe_names(cls) -> list[str]: + names = cls._lfm2_names() + dense_layer_1 = { + "blk.1.ffn_gate.weight", + "blk.1.ffn_up.weight", + "blk.1.ffn_down.weight", + } + names = [name for name in names if name not in dense_layer_1] + names.extend( + f"blk.1.{suffix}" + for suffix in ( + "ffn_gate_inp.weight", + "ffn_gate_exps.weight", + "ffn_up_exps.weight", + "ffn_down_exps.weight", + "exp_probs_b.bias", + ) + ) + return names + def test_lfm2_exact_mixer_closure_passes(self) -> None: from mobius.integrations.gguf._builder import ( _raise_for_invalid_hybrid_tensor_contract, @@ -4820,6 +5045,45 @@ def test_lfm2_exact_mixer_closure_passes(self) -> None: ) _raise_for_invalid_hybrid_tensor_contract(model) + def test_lfm2moe_exact_dense_and_routed_closure_passes(self) -> None: + from mobius.integrations.gguf._builder import ( + _raise_for_invalid_hybrid_tensor_contract, + ) + + _raise_for_invalid_hybrid_tensor_contract( + self._FakeGGUF( + "lfm2moe", + { + "lfm2moe.block_count": 2, + "lfm2moe.attention.head_count_kv": [0, 2], + "lfm2moe.leading_dense_block_count": 1, + }, + self._lfm2moe_names(), + ) + ) + + @pytest.mark.parametrize( + "wrong_tensor", + ["blk.0.ffn_gate_inp.weight", "blk.1.ffn_gate.weight"], + ) + def test_lfm2moe_wrong_ffn_family_is_rejected(self, wrong_tensor: str) -> None: + from mobius.integrations.gguf._builder import ( + _raise_for_invalid_hybrid_tensor_contract, + ) + + with pytest.raises(ValueError, match=r"wrong .* FFN family"): + _raise_for_invalid_hybrid_tensor_contract( + self._FakeGGUF( + "lfm2moe", + { + "lfm2moe.block_count": 2, + "lfm2moe.attention.head_count_kv": [0, 2], + "lfm2moe.leading_dense_block_count": 1, + }, + [*self._lfm2moe_names(), wrong_tensor], + ) + ) + def test_wrong_mixer_tensor_is_rejected(self) -> None: from mobius.integrations.gguf._builder import ( _raise_for_invalid_hybrid_tensor_contract, diff --git a/src/mobius/integrations/gguf/_config_mapping.py b/src/mobius/integrations/gguf/_config_mapping.py index f9c8899ce..71e3a1786 100644 --- a/src/mobius/integrations/gguf/_config_mapping.py +++ b/src/mobius/integrations/gguf/_config_mapping.py @@ -35,6 +35,7 @@ Gemma4Config, GraniteMoeHybridConfig, JambaConfig, + Lfm2MoeConfig, Mamba2Config, MambaConfig, MuseGlimmerConfig, @@ -338,7 +339,7 @@ def _derive_hybrid_layout( f"nextn predict layers ({mtp_count}) for architecture {gguf_arch}." ) - if gguf_arch in {"lfm2", "jamba", "granitehybrid"}: + if gguf_arch in {"lfm2", "lfm2moe", "jamba", "granitehybrid"}: raw_kv_heads = metadata.get(f"{gguf_arch}.attention.head_count_kv") if not isinstance(raw_kv_heads, (list, tuple, np.ndarray)): raise ValueError( @@ -357,6 +358,7 @@ def _derive_hybrid_layout( ) recurrent_type = { "lfm2": "conv", + "lfm2moe": "conv", "jamba": "mamba", "granitehybrid": "mamba2", }[gguf_arch] @@ -565,7 +567,13 @@ def gguf_to_config( num_kv_heads = hf_fields.get("num_key_value_heads", num_attention_heads) if isinstance(num_kv_heads, (list, np.ndarray)): values = [int(value) for value in num_kv_heads] - if canonical_arch in {"lfm2", "jamba", "nemotron_h", "granitehybrid"}: + if canonical_arch in { + "lfm2", + "lfm2moe", + "jamba", + "nemotron_h", + "granitehybrid", + }: nonzero = {value for value in values if value} if len(nonzero) != 1: raise ValueError( @@ -600,6 +608,7 @@ def gguf_to_config( canonical_arch in { "lfm2", + "lfm2moe", "jamba", "nemotron_h", "granitehybrid", @@ -891,6 +900,69 @@ def gguf_to_config( return config +def _lfm2moe_postprocess( + config: ArchitectureConfig, + metadata: dict[str, Any], + model: Any = None, +) -> Lfm2MoeConfig: + """Restore LFM2MoE fields serialized by the pinned llama.cpp converter. + + The pinned loader defaults a missing dense-prefix length to zero and + requires the SIGMOID gating enum. Its graph always normalizes selected + probabilities, and the architecture loader does not read a scaling + override, so metadata that conflicts with those invariants is rejected. + """ + del model + arch = "lfm2moe" + gating = int(metadata[f"{arch}.expert_gating_func"]) + if gating != 2: + raise ValueError(f"{arch}.expert_gating_func must be SIGMOID (2), got {gating}") + num_dense_layers = int(metadata.get(f"{arch}.leading_dense_block_count", 0)) + if not 0 <= num_dense_layers <= config.num_hidden_layers: + raise ValueError( + f"{arch}.leading_dense_block_count must be in [0, " + f"{config.num_hidden_layers}], got {num_dense_layers}" + ) + if ( + config.num_local_experts is None + or config.num_experts_per_tok is None + or config.moe_intermediate_size is None + ): + raise ValueError("lfm2moe requires expert count, top-k, and expert FFN width") + if not 0 < config.num_experts_per_tok <= config.num_local_experts: + raise ValueError( + "lfm2moe expert_used_count must be positive and no greater than expert_count" + ) + + if metadata.get(f"{arch}.expert_weights_norm", True) is not True: + raise ValueError( + "lfm2moe.expert_weights_norm=False is incompatible with the pinned " + "llama.cpp graph, which always normalizes selected expert weights" + ) + expert_scale = float(metadata.get(f"{arch}.expert_weights_scale", 1.0)) + if expert_scale != 1.0: # noqa: RUF069 + raise ValueError( + "lfm2moe.expert_weights_scale must be 1.0 because the pinned loader " + "does not read an architecture-specific override" + ) + + fields = _shallow_fields(config) + fields.update( + hidden_act="silu", + attn_qk_norm=True, + short_conv_bias=False, + scoring_func="sigmoid", + norm_topk_prob=True, + routed_scaling_factor=1.0, + ) + return Lfm2MoeConfig( + **fields, + num_dense_layers=num_dense_layers, + # The pinned loader requires exp_probs_b for every routed layer. + use_expert_bias=True, + ) + + def _gemma2_postprocess( config: ArchitectureConfig, metadata: dict[str, Any], @@ -2230,6 +2302,7 @@ def _eagle3_postprocess( "mamba": _mamba_postprocess, "mamba2": _mamba2_postprocess, "jamba": _jamba_postprocess, + "lfm2moe": _lfm2moe_postprocess, "nemotron_h": _nemotron_h_postprocess, "granitehybrid": _granitehybrid_postprocess, "bert_encoder": _bert_encoder_postprocess, diff --git a/src/mobius/integrations/gguf/_config_mapping_test.py b/src/mobius/integrations/gguf/_config_mapping_test.py index 9cb025eb8..3f37c967f 100644 --- a/src/mobius/integrations/gguf/_config_mapping_test.py +++ b/src/mobius/integrations/gguf/_config_mapping_test.py @@ -1180,6 +1180,95 @@ def test_lfm2_schedule_comes_only_from_kv_head_array(self) -> None: _derive_hybrid_layout("lfm2", md) +class TestLfm2MoePostprocess: + """Pinned LFM2MoE routing defaults remain overrideable and fail closed.""" + + @staticmethod + def _base_config(): + from mobius._configs import ArchitectureConfig + + return ArchitectureConfig( + hidden_size=64, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + vocab_size=256, + intermediate_size=128, + num_local_experts=4, + num_experts_per_tok=2, + moe_intermediate_size=32, + ) + + @staticmethod + def _metadata() -> dict[str, int]: + return { + "lfm2moe.expert_gating_func": 2, + "lfm2moe.leading_dense_block_count": 1, + } + + def test_pinned_invariant_routing_defaults(self) -> None: + from mobius.integrations.gguf._config_mapping import _lfm2moe_postprocess + + result = _lfm2moe_postprocess(self._base_config(), self._metadata()) + + assert result.num_dense_layers == 1 + assert result.norm_topk_prob is True + assert result.routed_scaling_factor == pytest.approx(1.0) + assert result.use_expert_bias is True + + def test_explicit_pinned_routing_values_are_accepted(self) -> None: + from mobius.integrations.gguf._config_mapping import _lfm2moe_postprocess + + metadata: dict[str, object] = { + **self._metadata(), + "lfm2moe.expert_weights_norm": True, + "lfm2moe.expert_weights_scale": 1.0, + } + result = _lfm2moe_postprocess(self._base_config(), metadata) + + assert result.norm_topk_prob is True + assert result.routed_scaling_factor == pytest.approx(1.0) + + @pytest.mark.parametrize( + ("key", "value"), + [ + ("lfm2moe.expert_weights_norm", False), + ("lfm2moe.expert_weights_scale", 0.75), + ], + ) + def test_non_pinned_routing_overrides_are_rejected(self, key: str, value: object) -> None: + from mobius.integrations.gguf._config_mapping import _lfm2moe_postprocess + + metadata: dict[str, object] = {**self._metadata(), key: value} + with pytest.raises(ValueError, match=key.split(".")[-1]): + _lfm2moe_postprocess(self._base_config(), metadata) + + def test_missing_dense_prefix_uses_pinned_loader_default(self) -> None: + from mobius.integrations.gguf._config_mapping import _lfm2moe_postprocess + + metadata = self._metadata() + del metadata["lfm2moe.leading_dense_block_count"] + + assert _lfm2moe_postprocess(self._base_config(), metadata).num_dense_layers == 0 + + def test_missing_gating_function_remains_fail_closed(self) -> None: + from mobius.integrations.gguf._config_mapping import _lfm2moe_postprocess + + metadata = self._metadata() + del metadata["lfm2moe.expert_gating_func"] + + with pytest.raises(KeyError, match="expert_gating_func"): + _lfm2moe_postprocess(self._base_config(), metadata) + + def test_all_dense_prefix_is_valid(self) -> None: + from mobius.integrations.gguf._config_mapping import _lfm2moe_postprocess + + metadata = self._metadata() + metadata["lfm2moe.leading_dense_block_count"] = 2 + + assert _lfm2moe_postprocess(self._base_config(), metadata).num_dense_layers == 2 + + class TestMuseGlimmerPostprocess: """Muse Glimmer config postprocessing. diff --git a/src/mobius/integrations/gguf/_tensor_mapping.py b/src/mobius/integrations/gguf/_tensor_mapping.py index 834cb15fd..80a50f71a 100644 --- a/src/mobius/integrations/gguf/_tensor_mapping.py +++ b/src/mobius/integrations/gguf/_tensor_mapping.py @@ -624,6 +624,16 @@ "blk.{bid}.shortconv.out_proj": "model.layers.{bid}.conv.out_proj", } +_LFM2_MOE_EXTRAS: dict[str, str] = { + "blk.{bid}.ffn_gate_inp": "model.layers.{bid}.feed_forward.gate", + "blk.{bid}.ffn_gate_exps": "model.layers.{bid}.feed_forward.experts.gate_proj", + "blk.{bid}.ffn_up_exps": "model.layers.{bid}.feed_forward.experts.up_proj", + "blk.{bid}.ffn_down_exps": "model.layers.{bid}.feed_forward.experts.down_proj", + # GGUF names this parameter as a bias sidecar, while the reference model + # stores it as a bare fp32 tensor on the routed feed-forward block. + "blk.{bid}.exp_probs_b": "model.layers.{bid}.feed_forward.expert_bias@", +} + # Architectures sharing the llama HF naming convention are declared in # ``_arch_registry`` via ``tensor_map_recipe=("llama", ...)`` rather than by a # frozenset here, so the "which architectures does this cover?" question has one @@ -674,6 +684,7 @@ { "llama": _LLAMA_MAPPING, "lfm2": _LFM2_MAPPING, + "lfm2_moe_extras": _LFM2_MOE_EXTRAS, "dflash": _DFLASH_MAPPING, "eagle3": _EAGLE3_MAPPING, "olmo": _OLMO_MAPPING, diff --git a/src/mobius/models/__init__.py b/src/mobius/models/__init__.py index 77e8505f1..f0a384fa4 100644 --- a/src/mobius/models/__init__.py +++ b/src/mobius/models/__init__.py @@ -87,6 +87,7 @@ "LLaVAModel", "LayerNormCausalLMModel", "Lfm2CausalLMModel", + "Lfm2MoECausalLMModel", "Lfm2VlForConditionalGeneration", "LongcatFlashCausalLMModel", "MPTCausalLMModel", @@ -255,7 +256,7 @@ from mobius.models.internvl import InternVL2Model from mobius.models.jamba import JambaCausalLMModel from mobius.models.jetmoe import JetMoeCausalLMModel -from mobius.models.lfm2 import Lfm2CausalLMModel +from mobius.models.lfm2 import Lfm2CausalLMModel, Lfm2MoECausalLMModel from mobius.models.lfm2_vl import Lfm2VlForConditionalGeneration from mobius.models.llada import DreamModel, LLaDAModel, LLaDAMoEModel, RND1Model from mobius.models.llama4 import Llama4CausalLMModel diff --git a/src/mobius/models/lfm2.py b/src/mobius/models/lfm2.py index 4ee17e681..c0a7a535f 100644 --- a/src/mobius/models/lfm2.py +++ b/src/mobius/models/lfm2.py @@ -12,17 +12,19 @@ import torch from onnxscript import OpBuilder, nn -from mobius._configs import ArchitectureConfig, Lfm2Config +from mobius._configs import ArchitectureConfig, Lfm2Config, Lfm2MoeConfig from mobius.components import ( MLP, Attention, Embedding, GatedShortConv, + MoELayer, RMSNorm, create_padding_mask, initialize_rope, ) from mobius.models.base import CausalLMModel +from mobius.models.moe import _rename_moe_expert_weights _ConfigT = TypeVar("_ConfigT", bound=ArchitectureConfig) @@ -140,6 +142,96 @@ def forward( return hidden_states, present_key_value +class Lfm2MoETopKGate(nn.Module): + """LFM2MoE sigmoid router with selection-only correction bias.""" + + def __init__( + self, + hidden_size: int, + num_experts: int, + top_k: int, + *, + norm_topk_prob: bool, + routed_scaling_factor: float, + ): + super().__init__() + self.weight = nn.Parameter([num_experts, hidden_size]) + self._top_k = top_k + self._norm_topk_prob = norm_topk_prob + self._routed_scaling_factor = routed_scaling_factor + + def forward( + self, + op: OpBuilder, + hidden_states: ir.Value, + expert_bias: ir.Value | None, + ) -> tuple[ir.Value, ir.Value]: + router_logits = op.MatMul(hidden_states, op.Transpose(self.weight, perm=[1, 0])) + routing_probs = op.Sigmoid(router_logits) + selection_scores = routing_probs + if expert_bias is not None: + # The checkpoint keeps the learned correction bias in fp32. It changes + # expert selection but never the probability used to combine experts. + selection_scores = op.Add( + op.Cast(routing_probs, to=ir.DataType.FLOAT), + expert_bias, + ) + _, selected_experts = op.TopK( + selection_scores, + op.Constant(value_ints=[self._top_k]), + axis=-1, + _outputs=2, + ) + routing_weights = op.GatherElements(routing_probs, selected_experts, axis=-1) + if self._norm_topk_prob: + weight_sum = op.ReduceSum(routing_weights, [-1], keepdims=True) + routing_weights = op.Div( + routing_weights, + op.Add(weight_sum, op.CastLike(1e-6, weight_sum)), + ) + if self._routed_scaling_factor != 1.0: # noqa: RUF069 + routing_weights = op.Mul( + routing_weights, + op.CastLike(self._routed_scaling_factor, routing_weights), + ) + return routing_weights, selected_experts + + +class Lfm2MoEFeedForward(MoELayer): + """LFM2MoE routed SwiGLU experts with optional selection correction bias.""" + + def __init__(self, config: Lfm2MoeConfig): + assert config.num_local_experts is not None + assert config.num_experts_per_tok is not None + gate = Lfm2MoETopKGate( + config.hidden_size, + config.num_local_experts, + config.num_experts_per_tok, + norm_topk_prob=config.norm_topk_prob, + routed_scaling_factor=config.routed_scaling_factor, + ) + super().__init__(config, gate=gate) + self.expert_bias = ( + nn.Parameter([config.num_local_experts], dtype=ir.DataType.FLOAT) + if config.use_expert_bias + else None + ) + if self.expert_bias is not None: + self.expert_bias._keep_float32 = True + + def forward(self, op: OpBuilder, hidden_states: ir.Value) -> ir.Value: + return super().forward(op, hidden_states, self.expert_bias) + + +class Lfm2MoEDecoderLayer(Lfm2DecoderLayer): + """LFM2MoE hybrid operator layer with dense-prefix or routed feed-forward.""" + + def __init__(self, config: Lfm2MoeConfig, layer_idx: int): + super().__init__(config, layer_idx) + if layer_idx >= config.num_dense_layers: + self.feed_forward = Lfm2MoEFeedForward(config) + + class Lfm2TextModel(nn.Module): """LFM2 decoder backbone with mixed convolution and full-attention layers.""" @@ -197,6 +289,16 @@ def forward( return hidden_states, present_key_values +class Lfm2MoETextModel(Lfm2TextModel): + """LFM2MoE decoder with the serialized dense-to-expert layer transition.""" + + def __init__(self, config: Lfm2MoeConfig): + super().__init__(config) + self.layers = nn.ModuleList( + [Lfm2MoEDecoderLayer(config, i) for i in range(config.num_hidden_layers)] + ) + + class Lfm2CausalLMModel(CausalLMModel): """LiquidAI LFM2 causal LM with double-gated short convolutions and QK-norm GQA.""" @@ -218,3 +320,27 @@ def preprocess_weights( ) -> dict[str, torch.Tensor]: """Map upstream LFM2 projection names to shared mobius components.""" return super().preprocess_weights(rename_lfm2_weights(state_dict)) + + +class Lfm2MoECausalLMModel(CausalLMModel): + """LiquidAI LFM2MoE hybrid causal LM with correction-biased sigmoid routing.""" + + default_task: str = "hybrid-text-generation" + category: str = "Hybrid Convolution+Attention MoE" + config_class: type = Lfm2MoeConfig + + def __init__(self, config: Lfm2MoeConfig): + config = apply_lfm2_config_defaults(config) + super().__init__(config) + self.model = Lfm2MoETextModel(config) + if config.tie_word_embeddings: + self.lm_head.weight = self.model.embed_tokens.weight + + def preprocess_weights( + self, + state_dict: dict[str, torch.Tensor], + ) -> dict[str, torch.Tensor]: + """Map dense and routed LFM2MoE weights onto the dedicated graph.""" + state_dict = rename_lfm2_weights(state_dict) + state_dict = _rename_moe_expert_weights(state_dict) + return super().preprocess_weights(state_dict) diff --git a/src/mobius/models/lfm2_test.py b/src/mobius/models/lfm2_test.py index 419a68ece..a87cef6fb 100644 --- a/src/mobius/models/lfm2_test.py +++ b/src/mobius/models/lfm2_test.py @@ -9,9 +9,10 @@ import onnx_ir as ir from mobius import build_from_module -from mobius._configs import Lfm2Config +from mobius._configs import Lfm2Config, Lfm2MoeConfig +from mobius._optimizations import SymbolicShapeInferencePass from mobius._registry import registry -from mobius.models.lfm2 import Lfm2CausalLMModel +from mobius.models.lfm2 import Lfm2CausalLMModel, Lfm2MoECausalLMModel def _hf_config(**overrides): @@ -108,3 +109,43 @@ def test_cuda_graph_uses_lfm2_fusions(): assert all( node.attributes.get_int("stash_type") == ir.DataType.FLOAT for node in remaining_norms ) + + +def test_lfm2moe_fp16_keeps_selection_bias_and_scores_in_float32(): + config = Lfm2MoeConfig( + hidden_size=64, + intermediate_size=128, + moe_intermediate_size=32, + num_attention_heads=4, + num_key_value_heads=2, + num_hidden_layers=2, + num_local_experts=4, + num_experts_per_tok=2, + num_dense_layers=1, + vocab_size=256, + max_position_embeddings=128, + head_dim=16, + layer_types=["conv", "full_attention"], + short_conv_kernel=3, + rope_type="default", + rope_theta=10_000.0, + dtype=ir.DataType.FLOAT16, + ) + model = build_from_module( + Lfm2MoECausalLMModel(config), + config, + task="hybrid-text-generation", + )["model"] + SymbolicShapeInferencePass()(model) + + bias_name = "model.layers.1.feed_forward.expert_bias" + assert model.graph.initializers[bias_name].dtype == ir.DataType.FLOAT + selection_add = next( + node + for node in model.graph + if node.op_type == "Add" + and any(value is not None and value.name == bias_name for value in node.inputs) + ) + assert all( + value is None or value.dtype == ir.DataType.FLOAT for value in selection_add.inputs + ) diff --git a/tests/_test_configs.py b/tests/_test_configs.py index 9b88a4880..01a85713a 100644 --- a/tests/_test_configs.py +++ b/tests/_test_configs.py @@ -34,6 +34,7 @@ JambaConfig, JetMoeConfig, Lfm2Config, + Lfm2MoeConfig, Lfm2VlConfig, LongcatFlashConfig, Mamba2Config, @@ -148,6 +149,25 @@ def _base_config(config_cls=None, **overrides) -> ArchitectureConfig: }, True, ), + ( + "lfm2_moe", + { + "_config_cls": Lfm2MoeConfig, + "layer_types": ["conv", "full_attention"], + "attn_qk_norm": True, + "short_conv_kernel": 3, + "short_conv_bias": False, + "num_dense_layers": 1, + "num_local_experts": 4, + "num_experts_per_tok": 2, + "moe_intermediate_size": 32, + "norm_topk_prob": True, + "routed_scaling_factor": 1.0, + "use_expert_bias": True, + "tie_word_embeddings": True, + }, + True, + ), ("mistral", {}, False), ("qwen2", {}, True), ("muse_glimmer_text", dict(_TINY_MUSE_GLIMMER_TEXT_OVERRIDES), True), diff --git a/tests/synthetic_parity_test.py b/tests/synthetic_parity_test.py index 80a266a77..6eadb8877 100644 --- a/tests/synthetic_parity_test.py +++ b/tests/synthetic_parity_test.py @@ -548,7 +548,7 @@ def _create_hf_config(model_type: str, config_overrides: dict): i for i, lt in enumerate(layer_types) if lt in ("full_attention", "attention") ] - if hf_model_type == "lfm2": + if hf_model_type in {"lfm2", "lfm2_moe"}: hf_kwargs["conv_L_cache"] = hf_kwargs.pop("short_conv_kernel", 3) hf_kwargs["conv_bias"] = hf_kwargs.pop("short_conv_bias", False) hf_kwargs["norm_eps"] = hf_kwargs.pop("rms_norm_eps") @@ -556,6 +556,8 @@ def _create_hf_config(model_type: str, config_overrides: dict): "rope_type": "default", "rope_theta": 10_000.0, } + if hf_model_type == "lfm2_moe": + hf_kwargs["num_experts"] = hf_kwargs.pop("num_local_experts") # Jamba uses attn_layer_offset/attn_layer_period if hf_model_type in ("jamba",) and "layer_types" in hf_kwargs: diff --git a/tests/weight_alignment_test.py b/tests/weight_alignment_test.py index c970a41c8..7bc0d5188 100644 --- a/tests/weight_alignment_test.py +++ b/tests/weight_alignment_test.py @@ -295,6 +295,60 @@ def test_lfm2_vl_huggingface_weight_alignment() -> None: assert not missing, f"LFM2-VL preprocessing missed: {sorted(missing)}" +def _lfm2moe_module(): + overrides = next( + overrides for mt, overrides, _ in ALL_CAUSAL_LM_CONFIGS if mt == "lfm2_moe" + ) + return registry.get("lfm2_moe")( + _base_config(**{**overrides, "tie_word_embeddings": False}) + ) + + +def test_lfm2moe_individual_expert_weight_alignment() -> None: + """The published checkpoint's w1/w2/w3 experts map to the dedicated MoE graph.""" + module = _lfm2moe_module() + state_dict = { + "model.layers.1.feed_forward.experts.0.w1.weight": torch.ones(32, 64), + "model.layers.1.feed_forward.experts.0.w2.weight": torch.ones(64, 32), + "model.layers.1.feed_forward.experts.0.w3.weight": torch.ones(32, 64), + } + + aligned = module.preprocess_weights(state_dict) + + assert set(aligned) == { + "model.layers.1.feed_forward.experts.0.gate_proj.weight", + "model.layers.1.feed_forward.experts.0.down_proj.weight", + "model.layers.1.feed_forward.experts.0.up_proj.weight", + } + + +def test_lfm2moe_fused_expert_weight_alignment() -> None: + """Current Transformers fused expert tensors split into per-expert projections.""" + module = _lfm2moe_module() + gate_up = torch.arange(4 * 64 * 64).reshape(4, 64, 64) + down = torch.arange(4 * 64 * 32).reshape(4, 64, 32) + state_dict = { + "model.layers.1.feed_forward.experts.gate_up_proj.weight": gate_up, + "model.layers.1.feed_forward.experts.down_proj.weight": down, + } + + aligned = module.preprocess_weights(state_dict) + + assert len(aligned) == 12 + assert torch.equal( + aligned["model.layers.1.feed_forward.experts.3.gate_proj.weight"], + gate_up[3, :32], + ) + assert torch.equal( + aligned["model.layers.1.feed_forward.experts.3.up_proj.weight"], + gate_up[3, 32:], + ) + assert torch.equal( + aligned["model.layers.1.feed_forward.experts.3.down_proj.weight"], + down[3], + ) + + # --------------------------------------------------------------------------- # Detection model weight alignment # ---------------------------------------------------------------------------