diff --git a/docs/api/build_from_gguf.md b/docs/api/build_from_gguf.md index 2fc1d49de..15b424216 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': 88, 'rejected': 3, 'supported': 56}; importable: 54; quantized import: {'rejected': 12, 'supported': 135}; runtime: {'deferred': 144, 'rejected': 3} | +| Architectures | 147 | graph verdicts: {'deferred': 88, 'rejected': 3, 'supported': 56}; importable: 54; quantized import: {'rejected': 11, 'supported': 136}; 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` | @@ -411,7 +411,7 @@ before graph construction or durable output. | `internlm2` | — | model=`internlm2`; tensor=`llama` | not claimed | config=supported; tensor_map=supported; graph=supported; runtime=deferred; quantized_import=rejected | 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. 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. | | `jais` | — | none (fails before config extraction) | audited-direct-loader-conditional-union | config=deferred; tensor_map=deferred; graph=deferred; runtime=deferred; quantized_import=supported | JAIS combines fused biased QKV, causal ALiBi, 1/head_dim attention scaling, parallel SwiGLU, and converter-baked MuP embedding/output scales. Reusing Falcon or Llama would lose required value transforms. | | `jais2` | — | none (fails before config extraction) | audited-direct-loader-conditional-union | config=deferred; tensor_map=deferred; graph=deferred; runtime=deferred; quantized_import=supported | JAIS2 is a distinct RoPE, bias-bearing LayerNorm decoder with split Q/K/V and a non-gated ReLU-squared FFN. It is not the ALiBi/SwiGLU JAIS graph and has no exact Mobius tensor recipe. | -| `jamba` | — | model=`jamba`; tensor=`jamba` | 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. | +| `jamba` | — | model=`jamba`; tensor=`jamba` | not claimed | config=supported; tensor_map=supported; graph=supported; runtime=deferred; quantized_import=supported | Exact mixed attention/Mamba and dense/routed-MoE schedules, strict tensor closure and shapes, GGUF value transforms, compatible projection quantization, value-checked expert ordering, reduced Transformers parity, and multi-token ORT state threading, reorder, and replay are covered. Generic ORT GenAI runtime packaging remains deferred because its released cache schema cannot represent heterogeneous KV, convolution, and recurrent state slots; tracked by #605. | | `jina-bert-v2` | — | none (fails before config extraction) | not claimed | config=deferred; tensor_map=deferred; graph=deferred; runtime=deferred; quantized_import=supported | JinaBERT v2 uses ALiBi, optional full-width Q/K norms, an extra attention norm, and either separate or fused GeGLU inputs. Mobius has no graph with that exact combination. | | `jina-bert-v3` | — | none (fails before config extraction) | not claimed | config=deferred; tensor_map=deferred; graph=deferred; runtime=deferred; quantized_import=supported | JinaBERT v3 uses RoPE and may alternate dense GELU and routed MoE layers. BertModel has absolute positions and no MoE path. | | `kimi-k3` | — | none (fails before config extraction) | not claimed | config=deferred; tensor_map=deferred; graph=deferred; runtime=deferred; quantized_import=supported | Kimi-K3 alternates KDA recurrent and NoPE MLA layers and requires convolution plus matrix state, sigmoid routed latent MoE with SiTU experts, optional shared experts, and cross-layer residual banks. Kimi-Linear is not an alias: its gates, expert activation, latent projections, and residual contract differ. Mobius has no exact graph or mixed-state task for either ABI. | @@ -559,19 +559,26 @@ independent full-logit and generation parity. ### Second hybrid cohort `jamba`, `nemotron_h`, and `granitehybrid` have graph-import support for exact -dense subsets only; runtime packaging remains deferred pending independent +pinned subsets. Jamba includes routed MoE layers; Nemotron-H and GraniteHybrid +remain dense-only. Runtime packaging remains deferred pending independent real-artifact full-logit and stateful-generation parity. - Schedules come from suffix-exact per-layer metadata. Jamba and GraniteHybrid use `attention.head_count_kv` (`0` selects Mamba/Mamba2). Nemotron-H combines that array with per-layer `feed_forward_length` to select exactly one of Mamba2, attention, or dense ReLU² MLP. -- Jamba requires `ssm.inner_size == 2 * embedding_length`. Nemotron-H rejects - MTP and all MoE files. GraniteHybrid rejects routed-MoE files until 3-D expert - fusion, ordering, and quantized preservation have independent value tests. +- Jamba requires `ssm.inner_size == 2 * embedding_length`, SiLU experts, + Mamba-1 with biased depthwise convolution and bias-free projections, and + softmax-first top-k routing without post-top-k renormalization. Routed layers + are inferred exactly from `ffn_gate_inp`; stacked expert tensors are split in + numeric order. There are no shared experts. Nemotron-H rejects MTP and all MoE + files. GraniteHybrid rejects routed-MoE files until 3-D expert fusion, + ordering, and quantized preservation have independent value tests. - Every layer must provide exactly its pinned loader tensor family. Missing, wrong-mixer, partial, auxiliary, scale/input-scale, and out-of-range tensors - are rejected before graph construction. GGUF Mamba decay values are inverted + are rejected before graph construction. Compatible attention, dense-FFN, and + expert MatMul weights may remain quantized; Mamba and other state-sensitive + tensors are dequantized. GGUF Mamba decay values are inverted from `-exp(A_log)`; convolution and grouped Mamba2 tensors are restored to graph shapes. - State inputs and outputs are caller-owned. Mamba1 uses conv diff --git a/src/mobius/_configs/_base.py b/src/mobius/_configs/_base.py index adcefe20b..e222457bc 100644 --- a/src/mobius/_configs/_base.py +++ b/src/mobius/_configs/_base.py @@ -2592,6 +2592,9 @@ class JambaConfig(ArchitectureConfig): attn_layer_offset: int = 4 expert_layer_period: int = 2 expert_layer_offset: int = 1 + # GGUF serializes the resolved schedule through tensor presence rather than + # preserving the source period/offset pair. + expert_layer_indices: list[int] | None = None @classmethod def from_transformers(cls, config, parent_config=None) -> JambaConfig: @@ -2601,18 +2604,33 @@ def from_transformers(cls, config, parent_config=None) -> JambaConfig: n = base.num_hidden_layers attn_period = getattr(config, "attn_layer_period", 8) attn_offset = getattr(config, "attn_layer_offset", 4) + expert_period = getattr(config, "expert_layer_period", 2) + expert_offset = getattr(config, "expert_layer_offset", 1) + if attn_period <= 0 or not 0 <= attn_offset < attn_period: + raise ValueError("Jamba attn_layer_offset must be in [0, attn_layer_period)") + if expert_period <= 0 or not 0 <= expert_offset < expert_period: + raise ValueError("Jamba expert_layer_offset must be in [0, expert_layer_period)") layer_types = [] for i in range(n): - if (i - attn_offset) % attn_period == 0: + if i % attn_period == attn_offset: layer_types.append("full_attention") else: layer_types.append("mamba") num_experts = getattr(config, "num_experts", 16) num_experts_per_tok = getattr(config, "num_experts_per_tok", 2) + dt_rank = getattr(config, "mamba_dt_rank", "auto") + if dt_rank == "auto": + dt_rank = math.ceil(base.hidden_size / 16) # Exclude fields we set explicitly below to avoid duplicate keyword args - _exclude = {"layer_types", "num_local_experts", "num_experts_per_tok"} + _exclude = { + "layer_types", + "num_local_experts", + "num_experts_per_tok", + "norm_topk_prob", + "rope_type", + } base_fields = {k: v for k, v in _shallow_fields(base).items() if k not in _exclude} return cls( **base_fields, @@ -2623,13 +2641,16 @@ def from_transformers(cls, config, parent_config=None) -> JambaConfig: mamba_d_state=getattr(config, "mamba_d_state", 16), mamba_d_conv=getattr(config, "mamba_d_conv", 4), mamba_expand=getattr(config, "mamba_expand", 2), - mamba_dt_rank=getattr(config, "mamba_dt_rank", 256), + mamba_dt_rank=int(dt_rank), mamba_conv_bias=getattr(config, "mamba_conv_bias", True), mamba_proj_bias=getattr(config, "mamba_proj_bias", False), attn_layer_period=attn_period, attn_layer_offset=attn_offset, - expert_layer_period=getattr(config, "expert_layer_period", 2), - expert_layer_offset=getattr(config, "expert_layer_offset", 1), + expert_layer_period=expert_period, + expert_layer_offset=expert_offset, + expert_layer_indices=[i for i in range(n) if i % expert_period == expert_offset], + norm_topk_prob=False, + rope_type=None, ) diff --git a/src/mobius/components/_ssm.py b/src/mobius/components/_ssm.py index e71319cd3..3039ff6c8 100644 --- a/src/mobius/components/_ssm.py +++ b/src/mobius/components/_ssm.py @@ -188,6 +188,60 @@ def _project_ssm_params(self, op, x_db): c_mat = self.c_layernorm(op, c_mat) return dt_raw, b_mat, c_mat + def _repeat_for_channels(self, op: OpBuilder, value: ir.Value) -> ir.Value: + # Jamba Mamba-1 shares token-dependent B/C across all expanded channels. + value = op.Unsqueeze(value, [2]) # (B, T, 1, d_state) + value = op.Tile(value, [1, 1, self.d_inner, 1]) + return op.Reshape(value, [0, 0, self.d_inner * self.d_state]) + + def forward( + self, + op: OpBuilder, + x: ir.Value, + ssm_state: ir.Value, + padding_mask: ir.Value | None = None, + ): + """Run the complete multi-token Jamba selective scan.""" + dt_raw, b_mat, c_mat = self._project_ssm_params(op, self.x_proj(op, x)) + dt = op.Softplus(op.Cast(self.dt_proj(op, dt_raw), to=ir.DataType.FLOAT)) + del padding_mask + + x_f32 = op.Cast(x, to=ir.DataType.FLOAT) + decay = op.Mul( + op.Unsqueeze(dt, [-1]), + op.Unsqueeze( + op.Neg(op.Exp(op.Cast(self.A_log, to=ir.DataType.FLOAT))), + [0, 1], + ), + ) + decay = op.Reshape(decay, [0, 0, self.d_inner * self.d_state]) + value = op.Mul(dt, x_f32) + internal_state = op.Unsqueeze( + op.Cast(ssm_state, to=ir.DataType.FLOAT), + [-1], + ) + output, present_state = op.LinearAttention( + self._repeat_for_channels(op, op.Cast(c_mat, to=ir.DataType.FLOAT)), + self._repeat_for_channels(op, op.Cast(b_mat, to=ir.DataType.FLOAT)), + value, + internal_state, + decay, + scale=1.0, + q_num_heads=self.d_inner, + kv_num_heads=self.d_inner, + update_rule="gated", + _domain="com.microsoft", + _outputs=2, + ) + output = op.Add( + output, + op.Mul(x_f32, op.Cast(self.D, to=ir.DataType.FLOAT)), + ) + return ( + op.CastLike(output, x), + op.CastLike(op.Squeeze(present_state, [-1]), ssm_state), + ) + class Mamba2Scan(nn.Module): """Multi-head selective scan for Mamba2/SSD architecture. diff --git a/src/mobius/integrations/gguf/_arch_registry.py b/src/mobius/integrations/gguf/_arch_registry.py index ceb119e21..e865eac12 100644 --- a/src/mobius/integrations/gguf/_arch_registry.py +++ b/src/mobius/integrations/gguf/_arch_registry.py @@ -261,6 +261,15 @@ "stateful ORT generation. Runtime packaging remains deferred until that evidence exists." ) +_JAMBA_RUNTIME_VALIDATION_PENDING = ( + "Exact mixed attention/Mamba and dense/routed-MoE schedules, strict tensor closure " + "and shapes, GGUF value transforms, compatible projection quantization, value-checked " + "expert ordering, reduced Transformers parity, and multi-token ORT state threading, " + "reorder, and replay are covered. Generic ORT GenAI runtime packaging remains deferred " + "because its released cache schema cannot represent heterogeneous KV, convolution, and " + "recurrent state slots; tracked by #605." +) + _RWKV_GRAPH_REASONS = { "rwkv6": ( "RWKV6 carries two F32 states per layer (two token-shift vectors and a per-head " @@ -1063,8 +1072,7 @@ "ssm.time_step_rank", ), runtime=Support.DEFERRED, - quantized_import=Support.REJECTED, - reason=_RECURRENT_RUNTIME_VALIDATION_PENDING + " " + _NO_QUANTIZED_PROJECTION_REASON, + reason=_JAMBA_RUNTIME_VALIDATION_PENDING, ), GGUFArchitectureSpec( gguf_arch="nemotron_h", diff --git a/src/mobius/integrations/gguf/_arch_registry_test.py b/src/mobius/integrations/gguf/_arch_registry_test.py index e55f79af8..94d4a11f8 100644 --- a/src/mobius/integrations/gguf/_arch_registry_test.py +++ b/src/mobius/integrations/gguf/_arch_registry_test.py @@ -162,6 +162,7 @@ "gpt2", "granitemoe", "hunyuan-dense", + "jamba", "lfm2", "llada", "llada-moe", @@ -270,7 +271,6 @@ def test_every_float_importable_architecture_has_a_quantized_verdict(self) -> No "chatglm", "granitehybrid", "internlm2", - "jamba", "lfm2moe", "mamba", "mamba2", diff --git a/src/mobius/integrations/gguf/_builder.py b/src/mobius/integrations/gguf/_builder.py index 6cbf2a53d..a2e1ebe9f 100644 --- a/src/mobius/integrations/gguf/_builder.py +++ b/src/mobius/integrations/gguf/_builder.py @@ -1699,6 +1699,8 @@ def _raise_for_invalid_encoder_tensor_contract(gguf_model) -> None: def _raise_for_invalid_mamba_hybrid_tensor_contract(gguf_model) -> None: """Require the exact dense tensor family for each audited hybrid layer.""" + import numpy as np + from mobius.integrations.gguf._config_mapping import _derive_hybrid_layout architecture = gguf_model.architecture @@ -1715,13 +1717,21 @@ def _raise_for_invalid_mamba_hybrid_tensor_contract(gguf_model) -> None: optional_by_type: dict[str, set[str]] common: set[str] if architecture == "jamba": - common = { - "attn_norm.weight", - "ffn_norm.weight", - "ffn_gate.weight", - "ffn_up.weight", - "ffn_down.weight", - } + common = {"attn_norm.weight", "ffn_norm.weight"} + num_experts = int(metadata.get("jamba.expert_count", 0)) + top_k = int(metadata.get("jamba.expert_used_count", 0)) + if bool(num_experts) != bool(top_k): + raise ValueError( + "Jamba expert_count and expert_used_count must both be zero or both positive" + ) + if num_experts and not 1 <= top_k <= num_experts: + raise ValueError( + f"Jamba expert_used_count must be in [1, {num_experts}], got {top_k}" + ) + if num_experts == 1: + raise ValueError( + "Jamba expert_count=1 is not a routed-MoE layout; use dense FFN tensors" + ) required_by_type = { "mamba": { "ssm_in.weight", @@ -1865,7 +1875,34 @@ def require_all_or_none(label: str, names: list[str]) -> None: for index, layer_type in enumerate(layer_types): prefix = f"blk.{index}." required = common | required_by_type[layer_type] - optional = optional_by_type[layer_type] + optional = set(optional_by_type[layer_type]) + if architecture == "jamba": + has_router = f"{prefix}ffn_gate_inp.weight" in actual + dense_ffn = {"ffn_gate.weight", "ffn_up.weight", "ffn_down.weight"} + moe_ffn = { + "ffn_gate_inp.weight", + "ffn_gate_exps.weight", + "ffn_up_exps.weight", + "ffn_down_exps.weight", + } + if has_router: + if not num_experts: + raise ValueError( + f"Jamba layer {index} has routed experts without expert metadata" + ) + required |= moe_ffn + optional |= { + suffix + for stem in ("ffn_gate_exps", "ffn_up_exps", "ffn_down_exps") + for suffix in (f"{stem}.scale", f"{stem}.input_scale") + } + else: + required |= dense_ffn + optional |= { + suffix + for stem in ("ffn_gate", "ffn_up", "ffn_down") + for suffix in (f"{stem}.scale", f"{stem}.input_scale") + } expected.update(prefix + suffix for suffix in required) allowed.update(prefix + suffix for suffix in required | optional) @@ -1882,6 +1919,127 @@ def require_all_or_none(label: str, names: list[str]) -> None: f"unexpected={unexpected}, out_of_range={out_of_range}" ) + if architecture != "jamba" or not hasattr(gguf_model, "tensor_items_raw"): + return + + shapes = { + name: tuple(int(dimension) for dimension in shape) + for name, _raw, _qtype, shape in gguf_model.tensor_items_raw() + } + hidden = int(metadata["jamba.embedding_length"]) + intermediate = int(metadata["jamba.feed_forward_length"]) + vocab = int(metadata.get("jamba.vocab_size", 0)) + if not vocab: + vocab = len(metadata.get("tokenizer.ggml.tokens", ())) + raw_head_counts = metadata["jamba.attention.head_count"] + head_counts = ( + [int(value) for value in raw_head_counts] + if isinstance(raw_head_counts, (list, tuple)) + else [int(raw_head_counts)] + ) + positive_head_counts = {value for value in head_counts if value > 0} + if len(positive_head_counts) != 1: + raise ValueError("Jamba GGUF attention layers must use one consistent head count") + heads = positive_head_counts.pop() + if hidden <= 0 or intermediate <= 0 or vocab <= 0 or hidden % heads: + raise ValueError( + "Jamba GGUF has inconsistent embedding, FFN, vocabulary, or head geometry" + ) + head_dim = hidden // heads + state = int(metadata["jamba.ssm.state_size"]) + inner = int(metadata["jamba.ssm.inner_size"]) + rank = int(metadata["jamba.ssm.time_step_rank"]) + conv = int(metadata["jamba.ssm.conv_kernel"]) + if inner != 2 * hidden or min(state, rank, conv) <= 0: + raise ValueError("Jamba GGUF has inconsistent Mamba-1 geometry") + + expected_shapes: dict[str, tuple[int, ...]] = { + "token_embd.weight": (vocab, hidden), + "output_norm.weight": (hidden,), + } + if "output.weight" in actual: + expected_shapes["output.weight"] = (vocab, hidden) + for index, layer_type in enumerate(layer_types): + prefix = f"blk.{index}." + expected_shapes[prefix + "attn_norm.weight"] = (hidden,) + expected_shapes[prefix + "ffn_norm.weight"] = (hidden,) + if layer_type == "mamba": + expected_shapes.update( + { + prefix + "ssm_in.weight": (2 * inner, hidden), + prefix + "ssm_conv1d.weight": (inner, conv), + prefix + "ssm_conv1d.bias": (inner,), + prefix + "ssm_x.weight": (rank + 2 * state, inner), + prefix + "ssm_dt_norm.weight": (rank,), + prefix + "ssm_dt.weight": (inner, rank), + prefix + "ssm_dt.bias": (inner,), + prefix + "ssm_b_norm.weight": (state,), + prefix + "ssm_c_norm.weight": (state,), + prefix + "ssm_a": (inner, state), + prefix + "ssm_d": (inner,), + prefix + "ssm_out.weight": (hidden, inner), + } + ) + else: + kv_heads = int(metadata["jamba.attention.head_count_kv"][index]) + if kv_heads <= 0 or heads % kv_heads: + raise ValueError(f"Jamba attention layer {index} has invalid KV head count") + kv_width = kv_heads * head_dim + expected_shapes.update( + { + prefix + "attn_q.weight": (hidden, hidden), + prefix + "attn_k.weight": (kv_width, hidden), + prefix + "attn_v.weight": (kv_width, hidden), + prefix + "attn_output.weight": (hidden, hidden), + } + ) + if prefix + "ffn_gate_inp.weight" in actual: + expected_shapes.update( + { + prefix + "ffn_gate_inp.weight": (num_experts, hidden), + prefix + "ffn_gate_exps.weight": ( + num_experts, + intermediate, + hidden, + ), + prefix + "ffn_up_exps.weight": ( + num_experts, + intermediate, + hidden, + ), + prefix + "ffn_down_exps.weight": ( + num_experts, + hidden, + intermediate, + ), + } + ) + else: + expected_shapes.update( + { + prefix + "ffn_gate.weight": (intermediate, hidden), + prefix + "ffn_up.weight": (intermediate, hidden), + prefix + "ffn_down.weight": (hidden, intermediate), + } + ) + malformed = sorted( + f"{name}: expected {expected_shape}, got {shapes[name]}" + for name, expected_shape in expected_shapes.items() + if name in shapes and shapes[name] != expected_shape + ) + if malformed: + raise ValueError(f"Invalid Jamba GGUF tensor shape(s): {malformed}") + for index, layer_type in enumerate(layer_types): + if layer_type != "mamba": + continue + decay_name = f"blk.{index}.ssm_a" + decay = np.asarray(gguf_model.get_tensor(decay_name)) + if not np.all(np.isfinite(decay)) or not np.all(decay < 0): + raise ValueError( + f"Malformed Jamba GGUF Mamba decay tensor {decay_name!r}: " + "ssm_a must contain only finite negative -exp(A_log) values" + ) + def _raise_for_invalid_t5_tensor_contract(gguf_model) -> None: """Validate the pinned T5/T5-encoder tensor closure and logical shapes.""" @@ -2716,6 +2874,10 @@ def build_from_gguf( module, config, reuse_candidates=reuse_candidates_by_id, + dequantize_float_linear_types=_float_linear_dequantization_types( + module, + gguf_arch, + ), ) else: state_dict = _load_dequantized_state_dict( @@ -3082,6 +3244,34 @@ def _replace_native_block_linears( ) +def _float_linear_dequantization_types( + module, + gguf_arch: str, +) -> Mapping[str, Collection[str]] | None: + """Return explicitly float projection types for mixed quantized imports.""" + if gguf_arch != "jamba": + return None + + from mobius.integrations.gguf._quant_registry import iter_quant_specs + + quantized_types = frozenset( + spec.name + for spec in iter_quant_specs() + if spec.is_quantized_storage and spec.dequantize is Support.SUPPORTED + ) + mamba_projection_suffixes = ( + ".mamba.in_proj", + ".mamba.out_proj", + ".mamba.ssm.x_proj", + ".mamba.ssm.dt_proj", + ) + return { + name: quantized_types + for name, _child in module.named_modules() + if name.endswith(mamba_projection_suffixes) + } + + #: GGUF architectures whose transformer RMSNorms are zero-centered #: (``output = norm(x) * (1 + weight)``, mobius :class:`OffsetRMSNorm`). Their #: llama.cpp converter bakes the ``+1`` into every ``*norm.weight`` *except* the @@ -4095,6 +4285,13 @@ def _load_quantized_state_dict( target_bits=target_bits, target_block_size=target_block_size, ) + explicitly_dequantized = ( + dequantize_float_linear_types is not None + and module_stem in dequantize_float_linear_types + and quant_spec.name in dequantize_float_linear_types[module_stem] + ) + if explicitly_dequantized and quant_spec.dequantize is Support.SUPPORTED: + route = QuantImportRoute.DEQUANTIZE_FLOAT if route is QuantImportRoute.REJECTED: raise ValueError( f"Cannot import GGUF tensor {gguf_name} mapped to {hf_name} " diff --git a/src/mobius/integrations/gguf/_builder_test.py b/src/mobius/integrations/gguf/_builder_test.py index c565aa8fb..a50613687 100644 --- a/src/mobius/integrations/gguf/_builder_test.py +++ b/src/mobius/integrations/gguf/_builder_test.py @@ -948,6 +948,161 @@ def add_q4(name: str, shape: tuple[int, int]) -> None: writer.close() +def _write_jamba_gguf( + path: Path, + *, + quantized: bool, + omit: str | None = None, + extra: str | None = None, + expert_count: int = 2, + expert_used_count: int = 1, + malformed_shape: str | None = None, + invalid_decay: bool = False, +) -> None: + """Write a tiny mixed Jamba GGUF with one dense and one routed layer.""" + from gguf import GGMLQuantizationType, GGUFWriter + + hidden = 32 + intermediate = 64 + vocab = 64 + heads = 4 + kv_heads = 2 + inner = 64 + state = 4 + rank = 2 + kernel = 4 + rng = np.random.default_rng(612) + + writer = GGUFWriter(str(path), "jamba") + writer.add_context_length(32) + 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_layer_norm_rms_eps(1e-6) + writer.add_vocab_size(vocab) + writer.add_ssm_conv_kernel(kernel) + writer.add_ssm_inner_size(inner) + writer.add_ssm_state_size(state) + writer.add_ssm_time_step_rank(rank) + writer.add_expert_count(expert_count) + writer.add_expert_used_count(expert_used_count) + writer.add_string("jamba.feed_forward.activation", "silu") + + def adjusted_shape(name: str, shape: tuple[int, ...]) -> tuple[int, ...]: + if name == malformed_shape: + return (*shape[:-1], shape[-1] + 1) + return shape + + def add_float( + name: str, + shape: tuple[int, ...], + *, + expert_order: bool = False, + negative: bool = False, + ) -> None: + if name == omit: + return + shape = adjusted_shape(name, shape) + values = rng.normal(0.0, 0.03, size=shape).astype(np.float32) + if negative: + values = -np.exp(values) + if invalid_decay: + values.flat[0] = -np.inf + if expert_order: + for expert in range(shape[0]): + values[expert].fill(expert + 1) + writer.add_tensor(name, values) + + def add_q4(name: str, shape: tuple[int, ...]) -> None: + if name == omit: + return + shape = adjusted_shape(name, shape) + assert shape[-1] % 32 == 0 + byte_shape = (*shape[:-1], shape[-1] // 32 * 18) + raw = np.zeros(byte_shape, dtype=np.uint8) + for index in np.ndindex(shape[:-1]): + for block in range(shape[-1] // 32): + offset = block * 18 + raw[(*index, slice(offset, offset + 2))] = np.array( + [rng.uniform(0.01, 0.05)], dtype=np.float16 + ).view(np.uint8) + raw[(*index, slice(offset + 2, offset + 18))] = rng.integers( + 0, 256, size=16, dtype=np.uint8 + ) + writer.add_tensor(name, raw, raw_dtype=GGMLQuantizationType.Q4_0) + + projection = add_q4 if quantized else add_float + add_float("token_embd.weight", (vocab, hidden)) + add_float("output_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,)) + + projection("blk.0.ffn_gate.weight", (intermediate, hidden)) + projection("blk.0.ffn_up.weight", (intermediate, hidden)) + projection("blk.0.ffn_down.weight", (hidden, intermediate)) + (add_q4 if quantized else add_float)("blk.0.ssm_in.weight", (2 * inner, hidden)) + add_float("blk.0.ssm_conv1d.weight", (inner, kernel)) + add_float("blk.0.ssm_conv1d.bias", (inner,)) + (add_q4 if quantized else add_float)( + "blk.0.ssm_x.weight", + (rank + 2 * state, inner), + ) + add_float("blk.0.ssm_dt_norm.weight", (rank,)) + add_float("blk.0.ssm_dt.weight", (inner, rank)) + add_float("blk.0.ssm_dt.bias", (inner,)) + add_float("blk.0.ssm_b_norm.weight", (state,)) + add_float("blk.0.ssm_c_norm.weight", (state,)) + add_float("blk.0.ssm_a", (inner, state), negative=True) + add_float("blk.0.ssm_d", (inner,)) + (add_q4 if quantized else add_float)("blk.0.ssm_out.weight", (hidden, inner)) + + projection("blk.1.attn_q.weight", (hidden, hidden)) + projection("blk.1.attn_k.weight", (kv_heads * hidden // heads, hidden)) + projection("blk.1.attn_v.weight", (kv_heads * hidden // heads, hidden)) + projection("blk.1.attn_output.weight", (hidden, hidden)) + add_float("blk.1.ffn_gate_inp.weight", (expert_count, hidden)) + if quantized: + projection( + "blk.1.ffn_gate_exps.weight", + (expert_count, intermediate, hidden), + ) + projection( + "blk.1.ffn_up_exps.weight", + (expert_count, intermediate, hidden), + ) + projection( + "blk.1.ffn_down_exps.weight", + (expert_count, hidden, intermediate), + ) + else: + add_float( + "blk.1.ffn_gate_exps.weight", + (expert_count, intermediate, hidden), + expert_order=True, + ) + add_float( + "blk.1.ffn_up_exps.weight", + (expert_count, intermediate, hidden), + expert_order=True, + ) + add_float( + "blk.1.ffn_down_exps.weight", + (expert_count, hidden, intermediate), + expert_order=True, + ) + if extra is not None: + add_float(extra, (1,)) + + writer.write_header_to_file() + writer.write_kv_data_to_file() + writer.write_tensors_to_file() + writer.close() + + def _write_lfm2_gguf(path: Path, *, quantized: bool) -> None: """Write a tiny two-layer LFM2 GGUF with one conv and one attention layer.""" from gguf import GGMLQuantizationType, GGUFWriter @@ -4669,6 +4824,151 @@ def unexpected_graph_build(*args, **kwargs): assert not graph_build_started +class TestJambaGGUFBuild: + """Jamba GGUF import preserves mixed mixers, routed experts, and state.""" + + @staticmethod + def _inputs(tokens: np.ndarray) -> dict[str, np.ndarray]: + batch, sequence = tokens.shape + return { + "input_ids": tokens, + "position_ids": np.broadcast_to( + np.arange(sequence, dtype=np.int64), + (batch, sequence), + ).copy(), + "attention_mask": np.ones((batch, sequence), np.int64), + "past_key_values.0.conv_state": np.zeros((batch, 64, 3), np.float32), + "past_key_values.0.ssm_state": np.zeros((batch, 64, 4), np.float32), + "past_key_values.1.key": np.zeros((batch, 2, 0, 8), np.float32), + "past_key_values.1.value": np.zeros((batch, 2, 0, 8), np.float32), + } + + def test_float_import_preserves_expert_order_and_round_trips(self, tmp_path: Path) -> None: + from mobius._model_package import ModelPackage + from mobius._testing.ort_inference import OnnxModelSession + from mobius.integrations.gguf import build_from_gguf + + path = tmp_path / "jamba-f32.gguf" + _write_jamba_gguf(path, quantized=False) + package = build_from_gguf(path) + model = package["model"] + assert model.metadata_props["mobius.runtime_support"].endswith( + "onnxruntime/mobius#605" + ) + assert [value.name for value in model.graph.outputs] == [ + "logits", + "present.0.conv_state", + "present.0.ssm_state", + "present.1.key", + "present.1.value", + ] + for expert in range(2): + for projection in ("gate_proj", "up_proj", "down_proj"): + value = model.graph.initializers[ + f"model.layers.1.feed_forward.experts.{expert}.{projection}.weight_t" + ].const_value.numpy() + np.testing.assert_array_equal(value, expert + 1) + + output_dir = tmp_path / "saved-jamba" + package.save(output_dir, progress_bar=False) + session = OnnxModelSession(ModelPackage.load(output_dir)["model"]) + outputs = session.run(self._inputs(np.asarray([[1, 2], [3, 4]], np.int64))) + assert outputs["logits"].shape == (2, 2, 64) + assert outputs["present.0.ssm_state"].dtype == np.float32 + + def test_prefill_decode_reorder_and_snapshot_replay(self, tmp_path: Path) -> None: + from mobius._testing.ort_inference import OnnxModelSession + from mobius.integrations.gguf import build_from_gguf + + path = tmp_path / "jamba-state.gguf" + _write_jamba_gguf(path, quantized=False) + session = OnnxModelSession(build_from_gguf(path)["model"]) + histories = np.asarray([[1, 2], [3, 4]], np.int64) + prefill = session.run(self._inputs(histories)) + + order = np.asarray([1, 0], np.int64) + next_tokens = np.asarray([[5], [6]], np.int64) + decode_inputs = { + "input_ids": next_tokens, + "position_ids": np.full((2, 1), 2, np.int64), + "attention_mask": np.ones((2, 3), np.int64), + "past_key_values.0.conv_state": prefill["present.0.conv_state"][order], + "past_key_values.0.ssm_state": prefill["present.0.ssm_state"][order], + "past_key_values.1.key": prefill["present.1.key"][order], + "past_key_values.1.value": prefill["present.1.value"][order], + } + decoded = session.run(decode_inputs) + replayed = session.run(decode_inputs) + for name in decoded: + np.testing.assert_array_equal(decoded[name], replayed[name]) + + full_tokens = np.concatenate([histories[order], next_tokens], axis=1) + full = session.run(self._inputs(full_tokens)) + np.testing.assert_allclose( + decoded["logits"][:, -1], + full["logits"][:, -1], + atol=2e-5, + rtol=2e-5, + ) + + def test_quantized_source_keeps_only_compatible_matmul_roles(self, tmp_path: Path) -> None: + from mobius.integrations.gguf import build_from_gguf + + path = tmp_path / "jamba-q4.gguf" + _write_jamba_gguf(path, quantized=True) + model = build_from_gguf(path, keep_quantized=True)["model"] + # 4 attention + 3 dense FFN + 2 experts * 3 expert projections. + assert sum(node.op_type == "MatMulNBits" for node in model.graph) == 13 + for stem in ( + "model.layers.0.mamba.in_proj.weight_t", + "model.layers.0.mamba.conv1d.weight", + "model.layers.0.mamba.ssm.x_proj.weight_t", + "model.layers.0.mamba.ssm.dt_proj.weight_t", + "model.layers.0.mamba.out_proj.weight_t", + ): + assert model.graph.initializers[stem].const_value.dtype == ir.DataType.FLOAT + + @pytest.mark.parametrize( + ("kwargs", "match"), + [ + ({"omit": "blk.1.ffn_up_exps.weight"}, "tensor closure"), + ({"extra": "blk.0.ffn_gate_inp.weight"}, "tensor closure"), + ( + {"malformed_shape": "blk.1.ffn_gate_exps.weight"}, + "tensor shape", + ), + ({"expert_count": 0, "expert_used_count": 1}, "expert_count"), + ({"expert_count": 1, "expert_used_count": 1}, "not a routed-MoE"), + ({"expert_count": 2, "expert_used_count": 3}, "expert_used_count"), + ({"invalid_decay": True}, "finite negative"), + ({"extra": "blk.0.ssm_in.scale"}, "auxiliary|tensor closure"), + ], + ) + def test_malformed_sources_fail_before_graph( + self, + tmp_path: Path, + monkeypatch, + kwargs: dict[str, object], + match: str, + ) -> None: + from mobius import _builder as core_builder + from mobius.integrations.gguf import build_from_gguf + + path = tmp_path / "jamba-invalid.gguf" + _write_jamba_gguf(path, quantized=False, **kwargs) + graph_build_started = False + + def unexpected_graph_build(*args, **kwargs): + nonlocal graph_build_started + graph_build_started = True + raise AssertionError("graph construction must not start") + + monkeypatch.setattr(core_builder, "build_from_module", unexpected_graph_build) + with pytest.raises(ValueError, match=match): + build_from_gguf(path) + assert not graph_build_started + + class TestBuildGgufStaticCache: """Tests for build_from_gguf(static_cache=True). diff --git a/src/mobius/integrations/gguf/_config_mapping.py b/src/mobius/integrations/gguf/_config_mapping.py index 2dc8e7580..b3adb6e31 100644 --- a/src/mobius/integrations/gguf/_config_mapping.py +++ b/src/mobius/integrations/gguf/_config_mapping.py @@ -24,6 +24,7 @@ import dataclasses import logging +import re from types import MappingProxyType from typing import TYPE_CHECKING, Any @@ -1441,26 +1442,58 @@ def _jamba_postprocess( metadata: dict[str, Any], model: Any, ) -> JambaConfig: - """Build the dense Jamba subset from the serialized per-layer schedule.""" + """Build exact Jamba mixer and routed-FFN schedules from serialized tensors.""" inner_size = int(metadata["jamba.ssm.inner_size"]) if inner_size != 2 * config.hidden_size: raise ValueError( "jamba.ssm.inner_size must equal 2 * embedding_length for the pinned loader" ) - if int(metadata.get("jamba.expert_count", 0)) or any( - ".ffn_gate_inp." in name for name in model.tensor_names - ): + if config.hidden_act not in {"silu", "swish"}: + raise ValueError("Jamba GGUF requires the pinned SiLU feed-forward activation") + num_experts = int(metadata.get("jamba.expert_count", 0)) + top_k = int(metadata.get("jamba.expert_used_count", 0)) + expert_layers = sorted( + { + int(match.group(1)) + for name in model.tensor_names + if (match := re.fullmatch(r"blk\.(\d+)\.ffn_gate_inp\.weight", name)) + } + ) + if num_experts: + if num_experts == 1: + raise ValueError( + "Jamba expert_count=1 is not a routed-MoE layout; use dense FFN tensors" + ) + if not 1 <= top_k <= num_experts: + raise ValueError( + f"jamba.expert_used_count must be in [1, {num_experts}], got {top_k}" + ) + if not expert_layers: + raise ValueError("Jamba expert metadata requires at least one routed MoE layer") + elif top_k or expert_layers: raise ValueError( - "Jamba GGUF MoE layers are deferred until stacked-expert parity is established" + "Jamba routed tensors require positive expert_count and expert_used_count" ) + output_present = "output.weight" in set(model.tensor_names) + fields = _shallow_fields(config) + fields.update( + num_local_experts=num_experts or None, + num_experts_per_tok=top_k or None, + moe_intermediate_size=config.intermediate_size if num_experts else None, + norm_topk_prob=False, + routed_scaling_factor=1.0, + tie_word_embeddings=not output_present, + rope_type=None, + ) return JambaConfig( - **_shallow_fields(config), + **fields, mamba_d_state=int(metadata["jamba.ssm.state_size"]), mamba_d_conv=int(metadata["jamba.ssm.conv_kernel"]), mamba_expand=2, mamba_dt_rank=int(metadata["jamba.ssm.time_step_rank"]), mamba_conv_bias=any(".ssm_conv1d.bias" in name for name in model.tensor_names), mamba_proj_bias=False, + expert_layer_indices=expert_layers, ) diff --git a/src/mobius/integrations/gguf/_config_mapping_test.py b/src/mobius/integrations/gguf/_config_mapping_test.py index 3f37c967f..26b2e761b 100644 --- a/src/mobius/integrations/gguf/_config_mapping_test.py +++ b/src/mobius/integrations/gguf/_config_mapping_test.py @@ -107,7 +107,11 @@ def test_jamba_nope_config_builds_hybrid_graph(self) -> None: _FakeDenseGGUF( "jamba", self._metadata("jamba"), - ["token_embd.weight", "output.weight"], + [ + "token_embd.weight", + "output.weight", + "blk.0.ssm_conv1d.bias", + ], ) ) assert config.rope_type is None @@ -141,7 +145,7 @@ def test_wrong_schedule_length_rejects(self, architecture: str) -> None: with pytest.raises(ValueError, match=r"exactly 3|each contain exactly 3"): gguf_to_config(_FakeDenseGGUF(architecture, metadata, ["token_embd.weight"])) - @pytest.mark.parametrize("architecture", ["jamba", "nemotron_h", "granitehybrid"]) + @pytest.mark.parametrize("architecture", ["nemotron_h", "granitehybrid"]) def test_moe_modes_fail_closed(self, architecture: str) -> None: from mobius.integrations.gguf._config_mapping import gguf_to_config @@ -151,6 +155,30 @@ def test_moe_modes_fail_closed(self, architecture: str) -> None: with pytest.raises(ValueError, match="MoE"): gguf_to_config(_FakeDenseGGUF(architecture, metadata, ["token_embd.weight"])) + def test_jamba_derives_exact_routed_layer_schedule(self) -> None: + from mobius.integrations.gguf._config_mapping import gguf_to_config + + metadata = self._metadata("jamba") + metadata["jamba.expert_count"] = 4 + metadata["jamba.expert_used_count"] = 2 + config = gguf_to_config( + _FakeDenseGGUF( + "jamba", + metadata, + [ + "token_embd.weight", + "blk.2.ffn_gate_inp.weight", + "blk.2.ffn_gate_exps.weight", + "blk.2.ffn_up_exps.weight", + "blk.2.ffn_down_exps.weight", + ], + ) + ) + assert config.expert_layer_indices == [2] + assert config.num_local_experts == 4 + assert config.num_experts_per_tok == 2 + assert config.norm_topk_prob is False + def _diffusion_names(architecture: str, *, output: bool = True) -> list[str]: names = ["token_embd.weight", "output_norm.weight"] diff --git a/src/mobius/integrations/gguf/_tensor_mapping.py b/src/mobius/integrations/gguf/_tensor_mapping.py index 0e2db5136..5a2bd3cdd 100644 --- a/src/mobius/integrations/gguf/_tensor_mapping.py +++ b/src/mobius/integrations/gguf/_tensor_mapping.py @@ -387,6 +387,10 @@ "blk.{bid}.ffn_gate": "model.layers.{bid}.feed_forward.gate_proj", "blk.{bid}.ffn_up": "model.layers.{bid}.feed_forward.up_proj", "blk.{bid}.ffn_down": "model.layers.{bid}.feed_forward.down_proj", + "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", } _NEMOTRON_H_MAPPING: dict[str, str] = { diff --git a/src/mobius/models/jamba.py b/src/mobius/models/jamba.py index 07fab48b1..a308f6a9a 100644 --- a/src/mobius/models/jamba.py +++ b/src/mobius/models/jamba.py @@ -7,9 +7,9 @@ Some layers use Mixture-of-Experts (MoE) MLPs instead of dense MLPs. Layer type selection (per HuggingFace JambaConfig): - - Attention if ``(i - attn_layer_offset) % attn_layer_period == 0`` + - Attention if ``i % attn_layer_period == attn_layer_offset`` - Mamba otherwise - - MoE MLP if ``(i - expert_layer_offset) % expert_layer_period == 0`` + - MoE MLP if ``i % expert_layer_period == expert_layer_offset`` - Dense MLP otherwise State per layer: @@ -22,8 +22,8 @@ from __future__ import annotations import math -from typing import TYPE_CHECKING +import onnx_ir as ir import torch from onnxscript import OpBuilder, nn @@ -31,25 +31,44 @@ from mobius.components import ( MLP, Attention, - Embedding, Linear, - MambaBlock, MoELayer, RMSNorm, - TopKGate, + TiedQuantizedLMHead, create_attention_bias, - initialize_rope, ) from mobius.components._ssm import JambaSelectiveScan - -if TYPE_CHECKING: - import onnx_ir as ir +from mobius.models.base import ( + effective_tie_word_embeddings, + embedding_for_config, + linear_class_for_config, +) # --------------------------------------------------------------------------- # Decoder layers # --------------------------------------------------------------------------- +class _JambaTopKGate(nn.Module): + """Jamba router with float32 full-softmax and unnormalized top-k weights.""" + + def __init__(self, hidden_size: int, num_experts: int, top_k: int): + super().__init__() + self.weight = nn.Parameter([num_experts, hidden_size]) + self.top_k = top_k + + def forward(self, op: OpBuilder, hidden_states: ir.Value): + logits = op.MatMul(hidden_states, op.Transpose(self.weight, perm=[1, 0])) + probabilities = op.Softmax(op.Cast(logits, to=ir.DataType.FLOAT), axis=-1) + weights, experts = op.TopK( + probabilities, + op.Constant(value_ints=[self.top_k]), + axis=-1, + _outputs=2, + ) + return op.CastLike(weights, hidden_states), experts + + class JambaMambaDecoderLayer(nn.Module): """Jamba Mamba layer: RMSNorm → MambaBlock → residual + optional MoE MLP. @@ -80,14 +99,21 @@ def __init__(self, config: JambaConfig, *, use_moe: bool = False): # MLP: MoE or dense if use_moe: - gate = TopKGate( + gate = _JambaTopKGate( config.hidden_size, config.num_local_experts, config.num_experts_per_tok, ) - self.feed_forward = MoELayer(config, gate=gate) + self.feed_forward = MoELayer( + config, + gate=gate, + linear_class=linear_class_for_config(config), + ) else: - self.feed_forward = MLP(config) + self.feed_forward = MLP( + config, + linear_class=linear_class_for_config(config), + ) self.pre_moe_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) def forward( @@ -95,8 +121,9 @@ def forward( op: OpBuilder, hidden_states: ir.Value, attention_bias: ir.Value, - position_embeddings: tuple, + position_embeddings: tuple | None, past_key_value: tuple | None, + padding_mask: ir.Value | None = None, ): """Forward pass. Returns (hidden_states, (conv_state, ssm_state)). @@ -111,7 +138,7 @@ def forward( conv_state, ssm_state = past_key_value if past_key_value is not None else (None, None) mamba_out, new_conv_state, new_ssm_state = self.mamba( - op, hidden_states, conv_state, ssm_state + op, hidden_states, conv_state, ssm_state, padding_mask ) hidden_states = op.Add(residual, mamba_out) @@ -136,19 +163,29 @@ class JambaAttentionDecoderLayer(nn.Module): def __init__(self, config: JambaConfig, *, use_moe: bool = False): super().__init__() - self.self_attn = Attention(config) + self.self_attn = Attention( + config, + linear_class=linear_class_for_config(config), + ) self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) # MLP: MoE or dense if use_moe: - gate = TopKGate( + gate = _JambaTopKGate( config.hidden_size, config.num_local_experts, config.num_experts_per_tok, ) - self.feed_forward = MoELayer(config, gate=gate) + self.feed_forward = MoELayer( + config, + gate=gate, + linear_class=linear_class_for_config(config), + ) else: - self.feed_forward = MLP(config) + self.feed_forward = MLP( + config, + linear_class=linear_class_for_config(config), + ) self.pre_moe_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) def forward( @@ -156,10 +193,12 @@ def forward( op: OpBuilder, hidden_states: ir.Value, attention_bias: ir.Value, - position_embeddings: tuple, + position_embeddings: tuple | None, past_key_value: tuple | None, + padding_mask: ir.Value | None = None, ): """Forward pass. Returns (hidden_states, (key, value)).""" + del padding_mask residual = hidden_states hidden_states = self.input_layernorm(op, hidden_states) @@ -196,10 +235,21 @@ class _JambaTextModel(nn.Module): def __init__(self, config: JambaConfig): super().__init__() + if config.hidden_act not in {"silu", "swish"}: + raise ValueError("Jamba requires SiLU expert and dense FFN activation") + if config.mamba_expand != 2: + raise ValueError("Jamba requires mamba_expand=2") + if not config.mamba_conv_bias or config.mamba_proj_bias: + raise ValueError("Jamba requires mamba_conv_bias=True and mamba_proj_bias=False") + if config.head_dim * config.num_attention_heads != config.hidden_size: + raise ValueError("Jamba attention head dimensions must reconstruct hidden_size") + if config.num_local_experts is not None: + if config.num_local_experts <= 0: + raise ValueError("Jamba num_experts must be positive when MoE is enabled") + if not 1 <= config.num_experts_per_tok <= config.num_local_experts: + raise ValueError("Jamba num_experts_per_tok must be in [1, num_experts]") self._dtype = config.dtype - self.embed_tokens = Embedding( - config.vocab_size, config.hidden_size, config.pad_token_id - ) + self.embed_tokens = embedding_for_config(config) layer_types = config.layer_types or [] if len(layer_types) != config.num_hidden_layers: @@ -209,14 +259,28 @@ def __init__(self, config: JambaConfig): if any(layer_type not in {"mamba", "full_attention"} for layer_type in layer_types): raise ValueError(f"Unknown Jamba layer type in {layer_types!r}") self.layers = nn.ModuleList([]) - expert_period = getattr(config, "expert_layer_period", 1) - expert_offset = getattr(config, "expert_layer_offset", 0) + if config.expert_layer_indices is None: + expert_period = config.expert_layer_period + expert_offset = config.expert_layer_offset + if expert_period <= 0 or not 0 <= expert_offset < expert_period: + raise ValueError( + "Jamba expert_layer_offset must be in [0, expert_layer_period)" + ) + expert_layers = { + i + for i in range(config.num_hidden_layers) + if i % expert_period == expert_offset + } + else: + expert_layers = set(config.expert_layer_indices) + if any(i < 0 or i >= config.num_hidden_layers for i in expert_layers): + raise ValueError("Jamba expert_layer_indices contains an out-of-range layer") for i in range(config.num_hidden_layers): ltype = layer_types[i] use_moe = ( config.num_local_experts is not None and config.num_local_experts > 1 - and (i - expert_offset) % expert_period == 0 + and i in expert_layers ) if ltype == "mamba": self.layers.append(JambaMambaDecoderLayer(config, use_moe=use_moe)) @@ -224,7 +288,6 @@ def __init__(self, config: JambaConfig): self.layers.append(JambaAttentionDecoderLayer(config, use_moe=use_moe)) self.final_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.rotary_emb = initialize_rope(config) def forward( self, @@ -235,9 +298,8 @@ def forward( past_key_values: list | None = None, ): hidden_states = self.embed_tokens(op, input_ids) - position_embeddings = ( - self.rotary_emb(op, position_ids) if self.rotary_emb is not None else None - ) + del position_ids + position_embeddings = None attention_bias = create_attention_bias( op, @@ -245,6 +307,16 @@ def forward( attention_mask=attention_mask, dtype=self._dtype, ) + current_length = op.Shape(input_ids, start=1, end=2) + padding_mask = op.Unsqueeze( + op.Slice( + attention_mask, + op.Neg(current_length), + [9223372036854775807], + [1], + ), + [-1], + ) present_key_values = [] past_kvs = past_key_values or [None] * len(self.layers) @@ -255,6 +327,7 @@ def forward( attention_bias=attention_bias, position_embeddings=position_embeddings, past_key_value=past_kv, + padding_mask=padding_mask, ) present_key_values.append(present_kv) @@ -279,8 +352,33 @@ def __init__(self, config: JambaConfig): super().__init__() self.config = config self.model = _JambaTextModel(config) - self.lm_head = Linear(config.hidden_size, config.vocab_size, bias=False) - if config.tie_word_embeddings: + quantization = getattr(config, "quantization", None) + quantized_embedding = quantization is not None and bool( + getattr(quantization, "quantize_embeddings", False) + ) + quantized_head = quantization is not None and bool( + getattr(quantization, "quantize_lm_head", False) + ) + tie = effective_tie_word_embeddings(config) + if tie and quantized_embedding != quantized_head: + raise ValueError( + "Jamba tied embeddings require token embedding and LM head " + "quantization to be enabled together" + ) + if tie and quantized_embedding: + self.lm_head = TiedQuantizedLMHead( + self.model.embed_tokens, + config.hidden_size, + config.vocab_size, + ) + else: + head_class = linear_class_for_config(config) if quantized_head else None + self.lm_head = (head_class or Linear)( + config.hidden_size, + config.vocab_size, + bias=False, + ) + if tie and not quantized_embedding: self.lm_head.weight = self.model.embed_tokens.weight def forward( @@ -313,7 +411,7 @@ def preprocess_weights( 4. SSM params nested under mamba.ssm 5. Attribute name mapping (mamba_mixer → mamba) """ - if self.config.tie_word_embeddings: + if effective_tie_word_embeddings(self.config): if "model.embed_tokens.weight" not in state_dict: state_dict["model.embed_tokens.weight"] = state_dict["lm_head.weight"] state_dict.pop("lm_head.weight", None) @@ -332,8 +430,33 @@ def preprocess_weights( # --------------------------------------------------------------------------- -class _JambaMambaBlock(MambaBlock): - """MambaBlock that uses JambaSelectiveScan (with dt/B/C layernorms).""" +class _JambaCausalConv(nn.Module): + """Jamba depthwise convolution with a fixed K-1 carry state.""" + + def __init__(self, channels: int, kernel_size: int): + super().__init__() + self.weight = nn.Parameter([channels, 1, kernel_size]) + self.bias = nn.Parameter([channels]) + + def forward( + self, + op: OpBuilder, + hidden_states: ir.Value, + conv_state: ir.Value, + ) -> tuple[ir.Value, ir.Value]: + return op.CausalConvWithState( + hidden_states, + self.weight, + self.bias, + conv_state, + activation="silu", + _domain="com.microsoft", + _outputs=2, + ) + + +class _JambaMambaBlock(nn.Module): + """Jamba Mamba-1 block with multi-token convolution and selective scan.""" def __init__( self, @@ -344,11 +467,45 @@ def __init__( conv_kernel: int = 4, rms_norm_eps: float = 1e-6, ): - super().__init__(d_model, d_inner, d_state, dt_rank, conv_kernel) - # Replace the SSM with the Jamba variant (layernormed dt/B/C) + super().__init__() + self.d_inner = d_inner + self.in_proj = Linear(d_model, 2 * d_inner, bias=False) + self.conv1d = _JambaCausalConv(d_inner, conv_kernel) self.ssm = JambaSelectiveScan( - d_inner, d_state, self.dt_rank, layer_norm_epsilon=rms_norm_eps + d_inner, + d_state, + dt_rank if dt_rank is not None else math.ceil(d_model / 16), + layer_norm_epsilon=rms_norm_eps, + ) + self.out_proj = Linear(d_inner, d_model, bias=False) + + def forward( + self, + op: OpBuilder, + hidden_states: ir.Value, + conv_state: ir.Value, + ssm_state: ir.Value, + padding_mask: ir.Value | None, + ) -> tuple[ir.Value, ir.Value, ir.Value]: + x, gate = op.Split( + self.in_proj(op, hidden_states), + [self.d_inner, self.d_inner], + axis=-1, + _outputs=2, + ) + if padding_mask is not None: + x = op.Mul(x, op.CastLike(padding_mask, x)) + x, present_conv = self.conv1d( + op, + op.Transpose(x, perm=[0, 2, 1]), + conv_state, ) + x = op.Transpose(x, perm=[0, 2, 1]) + if padding_mask is not None: + x = op.Mul(x, op.CastLike(padding_mask, x)) + output, present_state = self.ssm(op, x, ssm_state, padding_mask) + output = op.Mul(output, op.Swish(gate)) + return self.out_proj(op, output), present_conv, present_state # --------------------------------------------------------------------------- @@ -404,7 +561,10 @@ def _rename_jamba_weight( return None # handled inline # Fused down_proj: [num_experts, hidden, intermediate] # → split into per-expert down_proj - if key.endswith(".experts.down_proj") and value.dim() == 3: + if ( + key.endswith((".experts.down_proj", ".experts.down_proj.weight")) + and value.dim() == 3 + ): _split_fused_expert_down(key, value, out) return None # handled inline # w1 → gate_proj, w2 → down_proj, w3 → up_proj @@ -439,7 +599,11 @@ def _split_fused_expert_gate_up( intermediate = value.shape[1] // 2 # Base path: e.g. "layers.0.feed_forward.experts" - base = new_key.replace(".gate_up_proj", "") + base = ( + new_key.removesuffix(".gate_up_proj.weight") + if new_key.endswith(".gate_up_proj.weight") + else new_key.removesuffix(".gate_up_proj") + ) for e in range(num_experts): expert_w = value[e] # [2*intermediate, hidden] @@ -463,6 +627,10 @@ def _split_fused_expert_down( ``layers.{i}.feed_forward.experts.{e}.down_proj.weight`` """ num_experts = value.shape[0] - base = key.replace(".down_proj", "") + base = ( + key.removesuffix(".down_proj.weight") + if key.endswith(".down_proj.weight") + else key.removesuffix(".down_proj") + ) for e in range(num_experts): out[f"{base}.{e}.down_proj.weight"] = value[e] diff --git a/src/mobius/tasks/_cache_utils.py b/src/mobius/tasks/_cache_utils.py index 22a82ffde..d8d1bda57 100644 --- a/src/mobius/tasks/_cache_utils.py +++ b/src/mobius/tasks/_cache_utils.py @@ -414,10 +414,17 @@ def _register_linear_attention_functions( layer_types = getattr(config, "layer_types", None) or [] has_deltanet = "linear_attention" in layer_types has_lightning = "lightning_attention" in layer_types + has_mamba = "mamba" in layer_types has_mamba2 = "mamba2" in layer_types or isinstance(config, FalconH1Config) has_short_conv = "conv" in layer_types - if not has_deltanet and not has_lightning and not has_mamba2 and not has_short_conv: + if ( + not has_deltanet + and not has_lightning + and not has_mamba + and not has_mamba2 + and not has_short_conv + ): return from mobius.functions import ( @@ -454,6 +461,24 @@ def _register_linear_attention_functions( ) model.functions[attn_func_gated.identifier()] = attn_func_gated + if has_mamba: + d_inner = config.hidden_size * getattr(config, "mamba_expand", 2) + conv_func = causal_conv_nd_with_state( + kernel_size=getattr(config, "mamba_d_conv", 4), + channels=d_inner, + ndim=1, + activation="silu", + ) + attn_func = linear_attention( + q_num_heads=d_inner, + kv_num_heads=d_inner, + update_rule="gated", + scale=1.0, + stash_type=ir.DataType.FLOAT, + ) + model.functions[conv_func.identifier()] = conv_func + model.functions[attn_func.identifier()] = attn_func + if has_mamba2: mamba2_n_heads = getattr(config, "mamba_n_heads", 0) mamba2_d_head = getattr(config, "mamba_d_head", 0) diff --git a/src/mobius/tasks/_causal_lm.py b/src/mobius/tasks/_causal_lm.py index 72eb6a9b7..486736fcb 100644 --- a/src/mobius/tasks/_causal_lm.py +++ b/src/mobius/tasks/_causal_lm.py @@ -404,6 +404,11 @@ def build( model = _make_model(graph) _register_linear_attention_functions(model, config) + if config.model_type == "jamba": + model.metadata_props["mobius.runtime_support"] = ( + "Deferred: heterogeneous attention KV and Mamba recurrent state " + "discovery is tracked by https://github.com/onnxruntime/mobius#605" + ) return ModelPackage({"model": model}, config=config) diff --git a/tests/build_graph_test.py b/tests/build_graph_test.py index e634fbe74..08cf8a3ec 100644 --- a/tests/build_graph_test.py +++ b/tests/build_graph_test.py @@ -5801,10 +5801,8 @@ def _jamba_config(self): expert_layer_offset=1, num_local_experts=2, num_experts_per_tok=1, - # Jamba's attention layers use standard RoPE; enable it - # explicitly since ArchitectureConfig defaults ``rope_type`` to - # ``None`` (NoPE) to express "no RoPE" structurally. - rope_type="default", + # Jamba attention is positional-encoding-free. + rope_type=None, ) def test_jamba_builds(self): @@ -5865,6 +5863,85 @@ def test_jamba_registry_lookup(self): model_cls = registry.get("jamba") assert model_cls.__name__ == "JambaCausalLMModel" + def test_jamba_transformers_config_uses_exact_schedules(self): + """Transformers periods resolve without duplicate inherited fields.""" + from transformers import JambaConfig as HFJambaConfig + + from mobius._configs import JambaConfig + + config = JambaConfig.from_transformers( + HFJambaConfig( + vocab_size=TINY_VOCAB, + hidden_size=TINY_HIDDEN, + intermediate_size=TINY_INTERMEDIATE, + num_hidden_layers=4, + num_attention_heads=TINY_HEADS, + num_key_value_heads=TINY_KV_HEADS, + attn_layer_period=2, + attn_layer_offset=1, + expert_layer_period=2, + expert_layer_offset=1, + num_experts=2, + num_experts_per_tok=1, + mamba_d_state=8, + mamba_d_conv=4, + mamba_expand=2, + mamba_dt_rank="auto", + ) + ) + assert config.layer_types == [ + "mamba", + "full_attention", + "mamba", + "full_attention", + ] + assert config.expert_layer_indices == [1, 3] + assert config.mamba_dt_rank == (TINY_HIDDEN + 15) // 16 + assert config.rope_type is None + assert config.norm_topk_prob is False + + def test_jamba_router_softmaxes_in_float32(self): + """Jamba routes with a full-expert float32 softmax before top-k.""" + import dataclasses + + import onnx_ir as ir + + from mobius._builder import build_from_module + from mobius.models.jamba import JambaCausalLMModel + from mobius.tasks import HybridCausalLMTask + + config = dataclasses.replace(self._jamba_config(), dtype=ir.DataType.FLOAT16) + model = build_from_module( + JambaCausalLMModel(config), + config, + task=HybridCausalLMTask(), + )["model"] + softmaxes = [node for node in model.graph if node.op_type == "Softmax"] + assert softmaxes + for softmax in softmaxes: + assert softmax.inputs[0].producer().op_type == "Cast" + + def test_jamba_low_precision_state_matches_model_dtype(self): + """The public recurrent ABI matches Transformers cache storage dtype.""" + import dataclasses + + import onnx_ir as ir + + from mobius._builder import build_from_module + from mobius.models.jamba import JambaCausalLMModel + from mobius.tasks import HybridCausalLMTask + + config = dataclasses.replace(self._jamba_config(), dtype=ir.DataType.FLOAT16) + model = build_from_module( + JambaCausalLMModel(config), + config, + task=HybridCausalLMTask(), + )["model"] + inputs = {value.name: value for value in model.graph.inputs} + outputs = {value.name: value for value in model.graph.outputs} + assert inputs["past_key_values.0.ssm_state"].dtype == ir.DataType.FLOAT16 + assert outputs["present.0.ssm_state"].dtype == ir.DataType.FLOAT16 + def test_jamba_preprocess_weights_moe_renames(self): """Verify MoE expert weight renames and SSM nesting.""" import torch @@ -5896,6 +5973,42 @@ def test_jamba_preprocess_weights_moe_renames(self): # Non-SSM stays flat assert "model.layers.0.mamba.in_proj.weight" in result + def test_jamba_preprocesses_fused_experts_in_numeric_order(self): + """Current Transformers stacks every expert in one fused parameter.""" + import torch + + from mobius.models.jamba import JambaCausalLMModel + + module = JambaCausalLMModel(self._jamba_config()) + gate_up = torch.stack( + [ + torch.full((2 * TINY_INTERMEDIATE, TINY_HIDDEN), expert + 1.0) + for expert in range(2) + ] + ) + down = torch.stack( + [torch.full((TINY_HIDDEN, TINY_INTERMEDIATE), expert + 3.0) for expert in range(2)] + ) + result = module.preprocess_weights( + { + "model.layers.1.feed_forward.experts.gate_up_proj.weight": gate_up, + "model.layers.1.feed_forward.experts.down_proj.weight": down, + } + ) + for expert in range(2): + torch.testing.assert_close( + result[f"model.layers.1.feed_forward.experts.{expert}.gate_proj.weight"], + torch.full((TINY_INTERMEDIATE, TINY_HIDDEN), expert + 1.0), + ) + torch.testing.assert_close( + result[f"model.layers.1.feed_forward.experts.{expert}.up_proj.weight"], + torch.full((TINY_INTERMEDIATE, TINY_HIDDEN), expert + 1.0), + ) + torch.testing.assert_close( + result[f"model.layers.1.feed_forward.experts.{expert}.down_proj.weight"], + torch.full((TINY_HIDDEN, TINY_INTERMEDIATE), expert + 3.0), + ) + # =========================================================================== # Registry completeness diff --git a/tests/synthetic_parity_test.py b/tests/synthetic_parity_test.py index d7c7725ee..a698408ce 100644 --- a/tests/synthetic_parity_test.py +++ b/tests/synthetic_parity_test.py @@ -1185,12 +1185,8 @@ def test_synthetic_parity(model_type: str, config_overrides: dict): _fill_random_weights(onnx_model, rng) # 5. Prepare inputs - # Mamba1 (layer_type="mamba") only supports single-token decode (seq_len=1) - # because SelectiveScan uses a sequential recurrence that squeezes the seq - # dimension. Mamba2 and attention layers handle arbitrary seq_len. - layer_types = getattr(config, "layer_types", None) or [] - has_mamba1 = "mamba" in layer_types - prefill_seq_len = 1 if has_mamba1 else 3 + # Exercise multi-token prefill for both recurrent and attention architectures. + prefill_seq_len = 3 input_ids = rng.integers(1, config.vocab_size, size=(1, prefill_seq_len)).astype(np.int64) attention_mask = np.ones_like(input_ids) position_ids = np.arange(input_ids.shape[1], dtype=np.int64)[np.newaxis, :]