diff --git a/docs/api/build_from_gguf.md b/docs/api/build_from_gguf.md index 15b424216..bd008dae6 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': 11, 'supported': 136}; runtime: {'deferred': 144, 'rejected': 3} | +| Architectures | 147 | graph verdicts: {'deferred': 88, 'rejected': 2, 'supported': 57}; importable: 55; quantized import: {'rejected': 11, 'supported': 136}; runtime: {'deferred': 145, 'rejected': 2} | | 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` | @@ -442,7 +442,7 @@ before graph construction or durable output. | `nanbeige` | — | none (fails before config extraction) | audited-direct-loader-conditional-union | config=deferred; tensor_map=deferred; graph=deferred; runtime=deferred; quantized_import=supported | Nanbeige reuses physical layer weights across a configurable logical loop count, optionally normalizes between loops, and allocates a distinct KV slot for every logical occurrence. A Llama alias would build the wrong layer count and cache ABI. | | `nemotron` | — | model=`nemotron`; 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. | | `nemotron_h` | — | model=`nemotron_h`; tensor=`nemotron_h` | 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. | -| `nemotron_h_moe` | — | none (fails before config extraction) | not claimed | config=rejected; tensor_map=rejected; graph=rejected; runtime=rejected; quantized_import=rejected | Direct GGUF conversion is intentionally disabled. GGUF block_count includes a combined attention+MoE MTP auxiliary block, so aliasing it to the 52-layer 'nemotron_h' backbone would build the wrong graph. The Nemotron-H Mamba2 path also lacks passing full-logit/generation parity, and common GGUF presets contain Q5_0/Q5_1 expert tensors that cannot be preserved by MatMulNBits. Use llama.cpp/Unsloth to run the GGUF without changing its quantization, or start from the official pinned BF16 Hugging Face checkpoint and quantize the validated ONNX export with Olive only after L4/L5 semantic generation passes. See docs/api/build_from_gguf.md for the pinned recipe and waiver. | +| `nemotron_h_moe` | — | model=`nemotron_h`; tensor=`nemotron_h_moe` | not claimed | config=supported; tensor_map=supported; graph=supported; runtime=deferred; quantized_import=rejected | Exact mixed attention/Mamba2/dense/MoE scheduling, sigmoid correction-bias routing, shared experts, optional latent projections, and strict GGUF tensor closure are supported. Generic ORT GenAI runtime packaging remains deferred because its released cache schema cannot represent heterogeneous KV, convolution, and recurrent state slots; tracked by onnxruntime/mobius#605. Quantization preservation is unsupported because mixed Mamba2 recurrent parameters must remain dequantized and correction-biased sigmoid experts cannot use the fused MoE ABI. Use keep_quantized=False for explicit float import. | | `neo-bert` | — | none (fails before config extraction) | not claimed | config=deferred; tensor_map=deferred; graph=deferred; runtime=deferred; quantized_import=supported | NeoBERT uses pre-norm RMSNorm, RoPE, fused QKV, and fused SwiGLU. The existing encoder graphs differ in normalization and projection layout. | | `nomic-bert` | — | none (fails before config extraction) | not claimed | config=deferred; tensor_map=deferred; graph=deferred; runtime=deferred; quantized_import=supported | NomicBERT uses RoPE and a parallel gated FFN with BERT-style post norms. BertModel uses absolute positions and a non-gated sequential GELU FFN. | | `nomic-bert-moe` | — | none (fails before config extraction) | not claimed | config=deferred; tensor_map=deferred; graph=deferred; runtime=deferred; quantized_import=supported | NomicBERT-MoE alternates dense and routed-expert FFNs according to moe_every_n_layers. Mobius has no encoder MoE graph with that schedule. | @@ -604,8 +604,10 @@ deferred pending the heterogeneous-state schema tracked by [`onnxruntime/mobius#605`](https://github.com/onnxruntime/mobius/issues/605) and real full-logit plus deterministic stateful-generation evidence. -`nemotron_h_moe` remains rejected because its folded MTP attention+MoE head has -no equivalent package contract. +`nemotron_h_moe` backbone import is supported with exact mixed Mamba2, +attention, dense-MLP, and routed-MoE scheduling. Files with a folded MTP block +still fail before graph construction because no released package contract +represents that auxiliary attention+MoE head. ### Audio/TTS/codec cohort @@ -1393,163 +1395,42 @@ reported only as graph/import evidence. as documented above. It does not establish runtime support for an architecture. -## NVIDIA Nemotron 3.5 Lightning waiver - -Direct conversion of GGUF architecture `nemotron_h_moe` is intentionally -disabled. The following evidence is pinned: - -- GGUF repository: - `unsloth/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-GGUF` at - `f2d3fe3694501008786e81e5f20360cbf715496a`. -- Official BF16 comparison: - `nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16` at - `d468880b6ad3c6e0d21377ce7242adaea4cc884d`. -- The official backbone has exactly 52 layers: 23 Mamba, 23 MoE, and - 6 attention layers. GGUF block 52 is a separate combined attention+MoE MTP - auxiliary block, so `block_count=53` cannot be aliased to the backbone. - -### Quantization findings - -| GGUF file | Relevant tensor inventory | Direct preservation | -|---|---|---| -| `...-Q8_0.gguf` | 32.904B parameters in `Q8_0` | Qtype-compatible, but blocked by architecture and semantic validation | -| `...-MXFP4_MOE.gguf` | 14.687B `MXFP4`, 12.772B `Q5_1`, 5.445B `Q8_0` | No; the 5-bit expert weights require a quantization-changing float round-trip | -| `...-UD-Q4_K_M.gguf` | 15.326B `Q5_0`, 12.772B `Q5_1`, 4.806B `Q8_0` | No; the preset name does not describe its actual per-tensor types | -| `BF16/...-0000*-of-00002.gguf` | 329 tensors in shard 1 and 88 in shard 2 | No; Mobius does not assemble GGUF shards | - -The GGUF embeds GPT-2/Pixtral BPE metadata with BOS 1 and EOS 11, but declares -padding ID 999 (``). The pinned official tokenizer declares -`<|im_end|>` (ID 11) as padding. The GGUF also names the BF16 base repository -without recording its immutable source commit. Both discrepancies must be -resolved before a self-contained runtime package can be accepted. - -The guard also reflects missing semantic evidence: Nemotron-H Mamba2 synthetic -full-logit parity is not passing, and no real-weight ORT or ORT GenAI generation -has passed. Graph creation, config emission, or session creation is not a -substitute for generation. - -### Reproduce the guard with a pinned download - -The `Q8_0` file is the only practical candidate whose large quantized tensors -all use a currently repackable type. Download it explicitly so the source does -not move: - -```powershell -$repo = "unsloth/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-GGUF" -$revision = "f2d3fe3694501008786e81e5f20360cbf715496a" -$file = "NVIDIA-Nemotron-3.5-Lightning-30B-A3B-Q8_0.gguf" - -python -m pip install ` - --index-url https://packagefeedproxy.microsoft.io/pypi/simple ` - huggingface_hub -hf download $repo $file --revision $revision --local-dir .\nemotron-gguf - -# Expected: fail-fast NotImplementedError; no ONNX package is emitted. -python -m mobius build-gguf ".\nemotron-gguf\$file" ` - --ep cpu ` - --external-data safetensors --output .\nemotron-gguf-onnx -``` - -`mobius build-gguf --runtime ort-genai` is rejected separately. The GGUF CLI -does not emit `genai_config.json` until a selected architecture's cache and -tokenizer contracts have passed real ORT GenAI generation. - -To execute the pinned GGUF without changing its quantization, use current -llama.cpp instead: - -```powershell -.\llama-cli.exe ` - --model ".\nemotron-gguf\$file" ` - --temp 0.6 --top-p 0.95 --min-p 0.01 -``` - -### Option A: official BF16, then Olive - -Option A is the ONNX route because it preserves authoritative config, -tokenizer, and weight provenance. It is still a candidate until Nemotron-H -semantic tests pass, and currently targets direct ONNX Runtime rather than -ORT GenAI: - -```powershell -$repo = "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16" -$revision = "d468880b6ad3c6e0d21377ce7242adaea4cc884d" - -hf download $repo --revision $revision --local-dir .\nemotron-bf16 -python -m mobius build ` - --config .\nemotron-bf16 ` - --dtype bf16 --ep cuda ` - --external-data safetensors --max-shard-size 5GB ` - --output .\nemotron-bf16-onnx -``` - -After the BF16 package passes full-logit and generation parity, quantize its -decoder with an initialized Olive environment: - -```json -{ - "input_model": { - "type": "OnnxModel", - "model_path": "nemotron-bf16-onnx/model.onnx" - }, - "passes": { - "int4": { - "type": "OnnxKQuantQuantization", - "bits": 4, - "block_size": 32 - } - }, - "output_dir": "nemotron-int4-onnx" -} -``` - -```powershell -python -m pip install ` - --index-url https://packagefeedproxy.microsoft.io/pypi/simple ` - olive-ai onnxruntime -olive run --config .\olive-int4.json -Copy-Item .\nemotron-bf16\tokenizer* .\nemotron-int4-onnx\ -Copy-Item .\nemotron-bf16\special_tokens_map.json .\nemotron-int4-onnx\ -Copy-Item .\nemotron-bf16\chat_template.jinja .\nemotron-int4-onnx\ -``` - -The candidate can be checked for direct ORT session loading: - -```python -import onnxruntime as ort - -session = ort.InferenceSession( - r".\nemotron-int4-onnx\model.onnx", - providers=["CUDAExecutionProvider", "CPUExecutionProvider"], -) -print(session.get_providers()) -print([(value.name, value.shape, value.type) for value in session.get_inputs()]) -``` - -Session loading is not generation evidence. There is intentionally no ORT -GenAI generation command for this model at the pinned revisions: - -- ORT GenAI 0.15.2 does not register model type `nemotron_h`. -- The generated generic decoder config does not bind the graph's Mamba - `conv_state` and `recurrent_state` cache inputs. -- The official `generation_config.json` uses EOS IDs `[2, 11]`, while the - architecture config alone supplies EOS 2. - -Do not publish the package unless BF16 full logits match the pinned reference, -direct-ORT greedy generation is coherent and deterministic through an -independently validated hybrid-cache loop, the quantized package remains -non-degenerate, and ORT GenAI model/cache/token support is implemented before -claiming ORT GenAI compatibility. - -### Prerequisites for revisiting direct GGUF conversion - -1. Map the 52-layer schedule exactly and model block 52 as MTP, or explicitly - exclude it with generation evidence. -2. Fix Nemotron-H Mamba2 full-logit parity before testing quantized output. -3. Preserve every large source qtype. For Q5 variants this requires a validated - 5-bit runtime kernel and repacker; dequantize/requantize is not direct - preservation. -4. Resolve the GGUF padding-token mismatch and record an immutable upstream - BF16 source revision. -5. Pass real-weight prefill, cached decode, deterministic multi-token - generation, ORT load/inference, and ORT GenAI package generation on each - claimed EP. +## NVIDIA Nemotron-H MoE support boundary + +`nemotron_h_moe` backbone conversion is supported against pinned llama.cpp +`8d9af256337d1a501250f9bbf4c0859a654bddd6`. Import reconstructs the exact +per-layer Mamba2/attention/dense/MoE schedule and validates every required +tensor before graph construction. Routed experts use sigmoid probabilities, +correction bias only for top-k selection, optional probability normalization, +and the serialized routed scaling factor. Experts are non-gated ReLU-squared +MLPs; one shared expert and optional latent down/up projections are represented +explicitly. + +The importer accepts only the canonical stacked expert tensors and canonical +`blk.N.exp_probs_b.bias`. Missing expert tensors, separate or fused variants, +partial latent projections, unsupported scale/input-scale sidecars, and +inconsistent metadata or logical shapes fail before graph construction. +Quantized-source import is available only with `keep_quantized=False` +(`mobius build-gguf --dequantize`): mixed recurrent roles and the custom +sigmoid/ReLU-squared expert graph do not have an exact current quantized ABI. + +Synthetic evidence covers: + +- full-logit parity with Transformers for dense, routed/shared, and latent-MoE + schedules, including correction-bias-driven expert selection; +- float and dequantized Q4 GGUF import, expert-order value checks, ORT execution + with convolution/recurrent/KV state threading, and package round-trip; +- strict malformed-family, sidecar, shape, and MTP rejection. + +The smallest public full-MoE source is still the 30B-A3B model, so no practical +small real-weight checkpoint exists for CI parity. The pinned public pair is +`nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16` and +`unsloth/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-GGUF`. + +Those released GGUFs append a combined attention+MoE MTP block to the +52-layer backbone (23 Mamba2, 23 MoE, 6 attention). Mobius rejects that MTP +sidecar before graph construction rather than aliasing `block_count=53` to a +52-layer decoder. ORT GenAI packaging also remains deferred because released +runtime schemas do not represent the heterogeneous KV, convolution, and +recurrent state slots; see +[`onnxruntime/mobius#605`](https://github.com/onnxruntime/mobius/issues/605). diff --git a/docs/cli_reference.md b/docs/cli_reference.md index f66dcbbe4..7561c4bbc 100644 --- a/docs/cli_reference.md +++ b/docs/cli_reference.md @@ -389,10 +389,12 @@ embeddings dequantize because these graphs do not yet implement quantized fused QKV is also rejected, while float fused QKV is split losslessly. Sharded GGUF inputs are rejected because a single shard has an incomplete -tensor table. `nemotron_h_moe` is also rejected until its MTP block, Mamba2 -parity, mixed expert quantization, tokenizer provenance, and real ORT/ORT GenAI -generation are validated. See -[`build_from_gguf()`](api/build_from_gguf.md#nvidia-nemotron-35-lightning-waiver). +tensor table. MTP-free `nemotron_h_moe` backbones are supported with exact +hybrid scheduling and routed/shared/latent expert semantics; quantized sources +require `--dequantize`. Files with the released combined attention+MoE MTP +sidecar fail before graph construction, and ORT GenAI packaging remains +deferred. See +[`build_from_gguf()`](api/build_from_gguf.md#nvidia-nemotron-h-moe-support-boundary). Runtime packaging requires a validated embedded `tokenizer.huggingface.json`; opaque tokenizer pre-types are never reconstructed. Deferred/rejected diff --git a/src/mobius/integrations/gguf/_arch_registry.py b/src/mobius/integrations/gguf/_arch_registry.py index e865eac12..38601a1b4 100644 --- a/src/mobius/integrations/gguf/_arch_registry.py +++ b/src/mobius/integrations/gguf/_arch_registry.py @@ -74,18 +74,6 @@ "checkpoint with `mobius build` instead." ) -_NEMOTRON_H_MOE_REASON = ( - "Direct GGUF conversion is intentionally disabled. GGUF block_count includes a " - "combined attention+MoE MTP auxiliary block, so aliasing it to the 52-layer " - "'nemotron_h' backbone would build the wrong graph. The Nemotron-H Mamba2 path " - "also lacks passing full-logit/generation parity, and common GGUF presets " - "contain Q5_0/Q5_1 expert tensors that cannot be preserved by MatMulNBits. Use " - "llama.cpp/Unsloth to run the GGUF without changing its quantization, or start " - "from the official pinned BF16 Hugging Face checkpoint and quantize the " - "validated ONNX export with Olive only after L4/L5 semantic generation passes. " - "See docs/api/build_from_gguf.md for the pinned recipe and waiver." -) - _MMPROJ_REASON = ( "This is a multimodal projector sidecar, not a language model. Upstream it is a " "quant-only stub whose runtime lives outside libllama. Pass it to " @@ -1093,6 +1081,38 @@ quantized_import=Support.REJECTED, reason=_RECURRENT_RUNTIME_VALIDATION_PENDING + " " + _NO_QUANTIZED_PROJECTION_REASON, ), + GGUFArchitectureSpec( + gguf_arch="nemotron_h_moe", + model_type="nemotron_h", + tensor_map_recipe=("nemotron_h_moe",), + config_key_map="nemotron_h", + config_postprocessor="nemotron_h_moe", + tensor_processor="mamba", + required_metadata=( + "attention.layer_norm_rms_epsilon", + "ssm.conv_kernel", + "ssm.group_count", + "ssm.inner_size", + "ssm.state_size", + "ssm.time_step_rank", + "expert_count", + "expert_used_count", + "expert_feed_forward_length", + "expert_shared_feed_forward_length", + ), + runtime=Support.DEFERRED, + quantized_import=Support.REJECTED, + reason=( + "Exact mixed attention/Mamba2/dense/MoE scheduling, sigmoid correction-bias " + "routing, shared experts, optional latent projections, and strict GGUF tensor " + "closure are supported. Generic ORT GenAI runtime packaging remains deferred " + "because its released cache schema cannot represent heterogeneous KV, " + "convolution, and recurrent state slots; tracked by onnxruntime/mobius#605. " + "Quantization preservation is unsupported because mixed Mamba2 recurrent " + "parameters must remain dequantized and correction-biased sigmoid experts " + "cannot use the fused MoE ABI. Use keep_quantized=False for explicit float import." + ), + ), GGUFArchitectureSpec( gguf_arch="granitehybrid", model_type="granitemoehybrid", @@ -1746,15 +1766,6 @@ runtime=Support.DEFERRED, reason=_RWKV_GRAPH_REASONS["rwkv7"], ), - GGUFArchitectureSpec( - gguf_arch="nemotron_h_moe", - config=Support.REJECTED, - tensor_map=Support.REJECTED, - graph=Support.REJECTED, - runtime=Support.REJECTED, - quantized_import=Support.REJECTED, - reason=_NEMOTRON_H_MOE_REASON, - ), GGUFArchitectureSpec( gguf_arch=MMPROJ_ARCHITECTURE, config=Support.REJECTED, diff --git a/src/mobius/integrations/gguf/_arch_registry_test.py b/src/mobius/integrations/gguf/_arch_registry_test.py index 94d4a11f8..eca74635b 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 = 54 +_EXPECTED_SUPPORTED_COUNT = 55 _FINAL_CENSUS_CLOSURE = frozenset( { "afmoe", @@ -275,6 +275,7 @@ def test_every_float_importable_architecture_has_a_quantized_verdict(self) -> No "mamba", "mamba2", "nemotron_h", + "nemotron_h_moe", "phi2", } assert all(actual[arch] is Support.REJECTED for arch in rejected) @@ -1496,7 +1497,7 @@ def test_configurable_but_unmappable_architectures_are_refused( with pytest.raises(UnsupportedGGUFArchitectureError, match="Unsupported GGUF"): get_arch_spec(architecture) - @pytest.mark.parametrize("architecture", ["nemotron_h_moe", MMPROJ_ARCHITECTURE]) + @pytest.mark.parametrize("architecture", [MMPROJ_ARCHITECTURE]) def test_deliberately_disabled_architectures_raise_the_disabled_error( self, architecture: str ) -> None: @@ -1523,7 +1524,18 @@ def test_legacy_exception_types_still_catch_everything(self) -> None: with pytest.raises(ValueError): get_arch_spec("definitely-not-real") with pytest.raises(NotImplementedError): - get_arch_spec("nemotron_h_moe") + get_arch_spec(MMPROJ_ARCHITECTURE) + + def test_nemotron_h_moe_graph_is_supported_but_runtime_is_deferred(self) -> None: + spec = get_arch_spec("nemotron_h_moe") + assert spec.model_type == "nemotron_h" + 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 "onnxruntime/mobius#605" in spec.reason class TestDocumentedSupportMatrix: diff --git a/src/mobius/integrations/gguf/_builder.py b/src/mobius/integrations/gguf/_builder.py index a2e1ebe9f..8c4e78fc8 100644 --- a/src/mobius/integrations/gguf/_builder.py +++ b/src/mobius/integrations/gguf/_builder.py @@ -28,6 +28,7 @@ from typing import TYPE_CHECKING, Any from urllib.parse import urlparse +import numpy as np import tqdm from huggingface_hub import ( HfApi, @@ -1281,7 +1282,7 @@ def _raise_for_invalid_hybrid_tensor_contract(gguf_model) -> None: from mobius.integrations.gguf._config_mapping import _derive_hybrid_layout architecture = gguf_model.architecture - if architecture in {"jamba", "nemotron_h", "granitehybrid"}: + if architecture in {"jamba", "nemotron_h", "nemotron_h_moe", "granitehybrid"}: _raise_for_invalid_mamba_hybrid_tensor_contract(gguf_model) return if architecture not in {"lfm2", "lfm2moe", "qwen35", "qwen35moe", "qwen3next"}: @@ -1705,7 +1706,9 @@ def _raise_for_invalid_mamba_hybrid_tensor_contract(gguf_model) -> None: architecture = gguf_model.architecture metadata = gguf_model.metadata - layer_count, layer_types, mtp_count = _derive_hybrid_layout(architecture, metadata) + layer_count, layer_types, mtp_count = _derive_hybrid_layout( + architecture, metadata, gguf_model.tensor_names + ) assert layer_types is not None if mtp_count: raise ValueError(f"{architecture} GGUF auxiliary/MTP blocks are not supported") @@ -1755,7 +1758,7 @@ def _raise_for_invalid_mamba_hybrid_tensor_contract(gguf_model) -> None: }, } optional_by_type = {"mamba": set(), "full_attention": set()} - elif architecture == "nemotron_h": + elif architecture in {"nemotron_h", "nemotron_h_moe"}: common = {"attn_norm.weight"} required_by_type = { "mamba2": { @@ -1774,11 +1777,23 @@ def _raise_for_invalid_mamba_hybrid_tensor_contract(gguf_model) -> None: "attn_output.weight", }, "mlp": {"ffn_up.weight", "ffn_down.weight"}, + "moe": { + "ffn_gate_inp.weight", + "exp_probs_b.bias", + "ffn_up_exps.weight", + "ffn_down_exps.weight", + "ffn_up_shexp.weight", + "ffn_down_shexp.weight", + }, } optional_by_type = { "mamba2": {"ssm_conv1d.bias"}, "full_attention": {"attn_output.bias"}, "mlp": {"ffn_up.bias", "ffn_down.bias"}, + "moe": { + "ffn_latent_down.weight", + "ffn_latent_up.weight", + }, } else: common = { @@ -1818,11 +1833,6 @@ def _raise_for_invalid_mamba_hybrid_tensor_contract(gguf_model) -> None: expected = set(required_global) allowed = required_global | optional_global - if architecture == "nemotron_h" and "moe" in layer_types: - raise ValueError( - "Nemotron-H MoE GGUF import is deferred; fused softmax MoE is incompatible " - "with its sigmoid correction-bias routing" - ) def require_all_or_none(label: str, names: list[str]) -> None: present = sorted(set(names) & actual) @@ -1854,12 +1864,12 @@ def require_all_or_none(label: str, names: list[str]) -> None: for projection in ("q", "k", "v") ], ) - if architecture in {"nemotron_h", "granitehybrid"}: + if architecture in {"nemotron_h", "nemotron_h_moe", "granitehybrid"}: require_all_or_none( "attention output projection", [f"blk.{index}.attn_output.bias" for index in attention_layers], ) - if architecture == "nemotron_h": + if architecture in {"nemotron_h", "nemotron_h_moe"}: mlp_layers = [ index for index, layer_type in enumerate(layer_types) if layer_type == "mlp" ] @@ -1871,6 +1881,26 @@ def require_all_or_none(label: str, names: list[str]) -> None: for projection in ("up", "down") ], ) + moe_layers = [ + index for index, layer_type in enumerate(layer_types) if layer_type == "moe" + ] + require_all_or_none( + "MoE latent projection", + [ + f"blk.{index}.ffn_latent_{direction}.weight" + for index in moe_layers + for direction in ("down", "up") + ], + ) + has_latent_metadata = f"{architecture}.moe_latent_size" in metadata + has_latent_tensors = any( + f"blk.{index}.ffn_latent_down.weight" in actual for index in moe_layers + ) + if has_latent_metadata != has_latent_tensors: + raise ValueError( + f"{architecture} moe_latent_size metadata and latent projection tensors " + "must either both be present or both be absent" + ) for index, layer_type in enumerate(layer_types): prefix = f"blk.{index}." @@ -1919,13 +1949,19 @@ 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"): + if architecture not in {"jamba", "nemotron_h", "nemotron_h_moe"} 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() } + if architecture in {"nemotron_h", "nemotron_h_moe"}: + _validate_nemotron_h_tensor_shapes(gguf_model, layer_types, actual) + return + hidden = int(metadata["jamba.embedding_length"]) intermediate = int(metadata["jamba.feed_forward_length"]) vocab = int(metadata.get("jamba.vocab_size", 0)) @@ -2041,6 +2077,127 @@ def require_all_or_none(label: str, names: list[str]) -> None: ) +def _validate_nemotron_h_tensor_shapes(gguf_model, layer_types, actual: set[str]) -> None: + """Validate logical Nemotron-H tensor shapes before graph construction.""" + metadata = gguf_model.metadata + arch = gguf_model.architecture + shapes = { + name: tuple(int(dimension) for dimension in shape) + for name, _raw, _qtype, shape in gguf_model.tensor_items_raw() + } + hidden = int(metadata[f"{arch}.embedding_length"]) + vocab = int(metadata.get(f"{arch}.vocab_size", 0)) + if not vocab: + vocab = len(metadata.get("tokenizer.ggml.tokens", ())) + heads_raw = metadata[f"{arch}.attention.head_count"] + heads_by_layer = ( + [int(value) for value in heads_raw] + if isinstance(heads_raw, (list, tuple, np.ndarray)) + else [int(heads_raw)] * len(layer_types) + ) + kv_raw = metadata[f"{arch}.attention.head_count_kv"] + kv_by_layer = [int(value) for value in kv_raw] + head_dim = int( + metadata.get( + f"{arch}.attention.key_length", + hidden // next(value for value in heads_by_layer if value > 0), + ) + ) + state = int(metadata[f"{arch}.ssm.state_size"]) + inner = int(metadata[f"{arch}.ssm.inner_size"]) + groups = int(metadata[f"{arch}.ssm.group_count"]) + ssm_heads = int(metadata[f"{arch}.ssm.time_step_rank"]) + conv = int(metadata[f"{arch}.ssm.conv_kernel"]) + + expected: dict[str, tuple[int, ...]] = { + "token_embd.weight": (vocab, hidden), + "output_norm.weight": (hidden,), + } + if "output.weight" in actual: + expected["output.weight"] = (vocab, hidden) + for layer, layer_type in enumerate(layer_types): + prefix = f"blk.{layer}." + expected[prefix + "attn_norm.weight"] = (hidden,) + if layer_type == "mamba2": + conv_width = inner + 2 * groups * state + expected.update( + { + prefix + "ssm_in.weight": ( + 2 * inner + 2 * groups * state + ssm_heads, + hidden, + ), + prefix + "ssm_conv1d.weight": (conv_width, conv), + prefix + "ssm_dt.bias": (ssm_heads,), + prefix + "ssm_a": (ssm_heads, 1), + prefix + "ssm_d": (ssm_heads, 1), + prefix + "ssm_norm.weight": (groups, inner // groups), + prefix + "ssm_out.weight": (hidden, inner), + } + ) + if prefix + "ssm_conv1d.bias" in actual: + expected[prefix + "ssm_conv1d.bias"] = (conv_width,) + elif layer_type == "full_attention": + heads = heads_by_layer[layer] + kv_heads = kv_by_layer[layer] + expected.update( + { + prefix + "attn_q.weight": (heads * head_dim, hidden), + prefix + "attn_k.weight": (kv_heads * head_dim, hidden), + prefix + "attn_v.weight": (kv_heads * head_dim, hidden), + prefix + "attn_output.weight": (hidden, heads * head_dim), + } + ) + if prefix + "attn_output.bias" in actual: + expected[prefix + "attn_output.bias"] = (hidden,) + elif layer_type == "mlp": + width = int(metadata[f"{arch}.feed_forward_length"][layer]) + expected.update( + { + prefix + "ffn_up.weight": (width, hidden), + prefix + "ffn_down.weight": (hidden, width), + } + ) + for projection, size in (("up", width), ("down", hidden)): + name = prefix + f"ffn_{projection}.bias" + if name in actual: + expected[name] = (size,) + else: + experts = int(metadata[f"{arch}.expert_count"]) + expert_width = int(metadata[f"{arch}.expert_feed_forward_length"]) + shared_width = int(metadata[f"{arch}.expert_shared_feed_forward_length"]) + latent = metadata.get(f"{arch}.moe_latent_size") + expert_input = int(latent) if latent is not None else hidden + expected.update( + { + prefix + "ffn_gate_inp.weight": (experts, hidden), + prefix + "exp_probs_b.bias": (experts,), + prefix + "ffn_up_exps.weight": ( + experts, + expert_width, + expert_input, + ), + prefix + "ffn_down_exps.weight": ( + experts, + expert_input, + expert_width, + ), + prefix + "ffn_up_shexp.weight": (shared_width, hidden), + prefix + "ffn_down_shexp.weight": (hidden, shared_width), + } + ) + if latent is not None: + expected[prefix + "ffn_latent_down.weight"] = (int(latent), hidden) + expected[prefix + "ffn_latent_up.weight"] = (hidden, int(latent)) + + malformed = sorted( + f"{name}: expected {expected_shape}, got {shapes[name]}" + for name, expected_shape in expected.items() + if name in shapes and shapes[name] != expected_shape + ) + if malformed: + raise ValueError(f"Invalid Nemotron-H GGUF tensor shape(s): {malformed}") + + def _raise_for_invalid_t5_tensor_contract(gguf_model) -> None: """Validate the pinned T5/T5-encoder tensor closure and logical shapes.""" from mobius.integrations.gguf._quant_registry import float_storage_type_ids @@ -2669,6 +2826,7 @@ def build_from_gguf( "qwen3next", "jamba", "nemotron_h", + "nemotron_h_moe", "granitehybrid", }: from mobius.tasks import HybridCausalLMTask @@ -3645,11 +3803,7 @@ def _has_quantized_weights(gguf_model, gguf_arch: str) -> bool: for name, _raw, qtype, _shape in gguf_model.tensor_items_raw(): hf_name = map_gguf_to_hf_names(name, gguf_arch) type_id = getattr(qtype, "value", qtype) - if ( - hf_name is not None - and hf_name.endswith(".weight") - and type_id not in float_type_ids - ): + if hf_name is not None and name.endswith(".weight") and type_id not in float_type_ids: return True return False diff --git a/src/mobius/integrations/gguf/_builder_test.py b/src/mobius/integrations/gguf/_builder_test.py index a50613687..964fd640d 100644 --- a/src/mobius/integrations/gguf/_builder_test.py +++ b/src/mobius/integrations/gguf/_builder_test.py @@ -1103,6 +1103,172 @@ def add_q4(name: str, shape: tuple[int, ...]) -> None: writer.close() +def _write_nemotron_h_moe_gguf( + path: Path, + *, + quantized: bool, + latent: bool = False, + omit: str | None = None, + extra: str | None = None, + malformed_shape: str | None = None, + mtp: bool = False, + quantized_only: str | None = None, +) -> None: + """Write a tiny exact Nemotron-H backbone covering all four layer kinds.""" + from gguf import GGMLQuantizationType, GGUFWriter + + hidden = 32 + dense_width = 64 + expert_width = 64 + shared_width = 64 + latent_width = 32 + vocab = 64 + heads = 2 + kv_heads = 1 + ssm_heads = 4 + head_dim = hidden // heads + inner = 64 + groups = 1 + state = 4 + kernel = 4 + experts = 2 + rng = np.random.default_rng(613) + + writer = GGUFWriter(str(path), "nemotron_h_moe") + writer.add_context_length(32) + writer.add_embedding_length(hidden) + writer.add_feed_forward_length( + [0, dense_width, 0, dense_width] + ([dense_width] if mtp else []) + ) + writer.add_block_count(5 if mtp else 4) + writer.add_head_count([0, 0, heads, 0] + ([heads] if mtp else [])) + writer.add_head_count_kv([0, 0, kv_heads, 0] + ([kv_heads] if mtp else [])) + 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(ssm_heads) + writer.add_ssm_group_count(groups) + writer.add_expert_count(experts) + writer.add_expert_used_count(1) + writer.add_expert_feed_forward_length(expert_width) + writer.add_expert_shared_count(1) + writer.add_expert_shared_feed_forward_length(shared_width) + writer.add_expert_weights_norm(True) + writer.add_expert_weights_scale(2.5) + if latent: + writer.add_uint32("nemotron_h_moe.moe_latent_size", latent_width) + if mtp: + writer.add_uint32("nemotron_h_moe.nextn_predict_layers", 1) + + 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 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) + + def projection(name: str, shape: tuple[int, ...]) -> None: + if quantized and (quantized_only is None or name == quantized_only): + add_q4(name, shape) + else: + add_float(name, shape) + + add_float("token_embd.weight", (vocab, hidden)) + add_float("output_norm.weight", (hidden,)) + projection("output.weight", (vocab, hidden)) + for layer in range(4): + add_float(f"blk.{layer}.attn_norm.weight", (hidden,)) + + conv_width = inner + 2 * groups * state + projection( + "blk.0.ssm_in.weight", + (2 * inner + 2 * groups * state + ssm_heads, hidden), + ) + add_float("blk.0.ssm_conv1d.weight", (conv_width, kernel)) + add_float("blk.0.ssm_conv1d.bias", (conv_width,)) + add_float("blk.0.ssm_dt.bias", (ssm_heads,)) + add_float("blk.0.ssm_a", (ssm_heads, 1), negative=True) + add_float("blk.0.ssm_d", (ssm_heads, 1)) + add_float("blk.0.ssm_norm.weight", (groups, inner // groups)) + projection("blk.0.ssm_out.weight", (hidden, inner)) + + expert_input = latent_width if latent else hidden + add_float("blk.1.ffn_gate_inp.weight", (experts, hidden)) + add_float("blk.1.exp_probs_b.bias", (experts,)) + if quantized: + projection("blk.1.ffn_up_exps.weight", (experts, expert_width, expert_input)) + projection("blk.1.ffn_down_exps.weight", (experts, expert_input, expert_width)) + else: + add_float( + "blk.1.ffn_up_exps.weight", + (experts, expert_width, expert_input), + expert_order=True, + ) + add_float( + "blk.1.ffn_down_exps.weight", + (experts, expert_input, expert_width), + expert_order=True, + ) + projection("blk.1.ffn_up_shexp.weight", (shared_width, hidden)) + projection("blk.1.ffn_down_shexp.weight", (hidden, shared_width)) + if latent: + projection("blk.1.ffn_latent_down.weight", (latent_width, hidden)) + projection("blk.1.ffn_latent_up.weight", (hidden, latent_width)) + + projection("blk.2.attn_q.weight", (heads * head_dim, hidden)) + projection("blk.2.attn_k.weight", (kv_heads * head_dim, hidden)) + projection("blk.2.attn_v.weight", (kv_heads * head_dim, hidden)) + projection("blk.2.attn_output.weight", (hidden, heads * head_dim)) + + projection("blk.3.ffn_up.weight", (dense_width, hidden)) + projection("blk.3.ffn_down.weight", (hidden, dense_width)) + if mtp: + projection("blk.4.nextn.eh_proj.weight", (hidden, 2 * hidden)) + 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 @@ -4969,6 +5135,208 @@ def unexpected_graph_build(*args, **kwargs): assert not graph_build_started +class TestNemotronHMoEGGUFBuild: + """Nemotron-H GGUF import preserves hybrid state and exact expert semantics.""" + + @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, 72, 3), np.float32), + "past_key_values.0.ssm_state": np.zeros((batch, 4, 4, 16), np.float32), + "past_key_values.2.key": np.zeros((batch, 1, 0, 16), np.float32), + "past_key_values.2.value": np.zeros((batch, 1, 0, 16), np.float32), + } + + def test_float_import_preserves_expert_order_state_and_roundtrip( + 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 / "nemotron-h-moe-f32.gguf" + _write_nemotron_h_moe_gguf(path, quantized=False) + package = build_from_gguf(path) + model = package["model"] + assert [value.name for value in model.graph.outputs] == [ + "logits", + "present.0.conv_state", + "present.0.ssm_state", + "present.2.key", + "present.2.value", + ] + for expert in range(2): + for projection in ("up_proj", "down_proj"): + value = model.graph.initializers[ + f"model.layers.1.moe.experts.{expert}.{projection}.weight_t" + ].const_value.numpy() + np.testing.assert_array_equal(value, expert + 1) + + output_dir = tmp_path / "saved-nemotron-h-moe" + 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.conv_state"].shape == (2, 72, 3) + assert outputs["present.0.ssm_state"].shape == (2, 4, 4, 16) + session.close() + + def test_latent_projection_imports_and_executes(self, tmp_path: Path) -> None: + from mobius._testing.ort_inference import OnnxModelSession + from mobius.integrations.gguf import build_from_gguf + + path = tmp_path / "nemotron-h-moe-latent.gguf" + _write_nemotron_h_moe_gguf(path, quantized=False, latent=True) + model = build_from_gguf(path)["model"] + assert model.graph.initializers[ + "model.layers.1.moe.fc1_latent_proj.weight_t" + ].shape == [32, 32] + session = OnnxModelSession(model) + outputs = session.run(self._inputs(np.asarray([[1, 2]], np.int64))) + assert outputs["logits"].shape == (1, 2, 64) + session.close() + + def test_prefill_decode_threads_all_state_families(self, tmp_path: Path) -> None: + from mobius._testing.ort_inference import OnnxModelSession + from mobius.integrations.gguf import build_from_gguf + + path = tmp_path / "nemotron-h-moe-state.gguf" + _write_nemotron_h_moe_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.2.key": prefill["present.2.key"][order], + "past_key_values.2.value": prefill["present.2.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, + ) + session.close() + + def test_quantized_source_requires_explicit_dequantization(self, tmp_path: Path) -> None: + from mobius.integrations.gguf import build_from_gguf + + path = tmp_path / "nemotron-h-moe-q4.gguf" + _write_nemotron_h_moe_gguf(path, quantized=True) + with pytest.raises(ValueError, match=r"keep_quantized=False"): + build_from_gguf(path, keep_quantized=True) + + model = build_from_gguf(path, keep_quantized=False)["model"] + assert all(node.op_type != "MatMulNBits" for node in model.graph) + assert ( + model.graph.initializers[ + "model.layers.1.moe.experts.0.up_proj.weight_t" + ].const_value.dtype + == ir.DataType.FLOAT + ) + + def test_quantized_expert_only_requires_explicit_dequantization( + self, tmp_path: Path + ) -> None: + from mobius.integrations.gguf import build_from_gguf + + path = tmp_path / "nemotron-h-moe-expert-q4.gguf" + _write_nemotron_h_moe_gguf( + path, + quantized=True, + quantized_only="blk.1.ffn_up_exps.weight", + ) + with pytest.raises(ValueError, match=r"keep_quantized=False"): + build_from_gguf(path, keep_quantized=True) + + @pytest.mark.parametrize( + ("kwargs", "match"), + [ + ({"omit": "blk.1.ffn_up_exps.weight"}, "tensor closure"), + ({"omit": "blk.1.exp_probs_b.bias"}, "tensor closure"), + ({"extra": "blk.1.ffn_exp_probs_b.bias"}, "tensor closure"), + ({"extra": "blk.1.ffn_up_exps.scale"}, "auxiliary quantization"), + ( + {"latent": True, "omit": "blk.1.ffn_latent_up.weight"}, + "latent projection", + ), + ( + {"malformed_shape": "blk.1.ffn_down_exps.weight"}, + "tensor shape", + ), + ], + ) + 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 / "nemotron-h-moe-invalid.gguf" + _write_nemotron_h_moe_gguf(path, quantized=False, **kwargs) + graph_build = mock.Mock(side_effect=AssertionError("graph construction reached")) + monkeypatch.setattr(core_builder, "build_from_module", graph_build) + with pytest.raises(ValueError, match=match): + build_from_gguf(path) + graph_build.assert_not_called() + + def test_mtp_sidecar_fails_before_graph(self, tmp_path: Path, monkeypatch) -> None: + from mobius import _builder as core_builder + from mobius.integrations.gguf import build_from_gguf + + path = tmp_path / "nemotron-h-moe-mtp.gguf" + _write_nemotron_h_moe_gguf(path, quantized=False, mtp=True) + graph_build = mock.Mock(side_effect=AssertionError("graph construction reached")) + monkeypatch.setattr(core_builder, "build_from_module", graph_build) + with pytest.raises(NotImplementedError, match="routed experts"): + build_from_gguf(path) + graph_build.assert_not_called() + + @pytest.mark.parametrize( + ("options", "match"), + [ + ({"static_cache": True}, "static_cache=True"), + ({"task": "text-generation"}, "hybrid-text-generation"), + ], + ) + def test_incompatible_task_options_fail_closed( + self, + tmp_path: Path, + options: dict[str, object], + match: str, + ) -> None: + from mobius.integrations.gguf import build_from_gguf + + path = tmp_path / "nemotron-h-moe-options.gguf" + _write_nemotron_h_moe_gguf(path, quantized=False) + with pytest.raises(ValueError, match=match): + build_from_gguf(path, **options) + + class TestBuildGgufStaticCache: """Tests for build_from_gguf(static_cache=True). @@ -5184,22 +5552,6 @@ def test_nemotron_layout_excludes_combined_mtp_block(self): assert mtp_blocks == (52,) assert mtp_kinds == {52: frozenset({"attention", "moe"})} - def test_local_nemotron_h_moe_fails_before_graph_build(self, tmp_path: Path): - from mobius.integrations.gguf import build_from_gguf - - path = tmp_path / "nemotron-h-moe-q8.gguf" - _write_quantized_gguf(path, architecture="nemotron_h_moe") - - with pytest.raises(NotImplementedError) as exc_info: - build_from_gguf(path, keep_quantized=True) - - message = str(exc_info.value) - assert "intentionally disabled" in message - assert "MTP auxiliary block" in message - assert "Q5_0/Q5_1" in message - assert "llama.cpp/Unsloth" in message - assert "Olive" in message - @pytest.mark.parametrize( ("architecture", "projection_quantization"), [ @@ -5380,23 +5732,6 @@ def test_standalone_clip_fails_before_all_downstream_stages( module_lookup.assert_not_called() graph_build.assert_not_called() - def test_remote_nemotron_h_moe_fails_before_download(self): - from mobius.integrations.gguf._builder import _resolve_gguf_path - - filename = "NVIDIA-Nemotron-3.5-Lightning-30B-A3B-Q8_0.gguf" - with ( - mock.patch( - "mobius.integrations.gguf._builder._preflight_hf_gguf_file", - side_effect=NotImplementedError("nemotron_h_moe"), - ) as preflight, - mock.patch("mobius.integrations.gguf._builder.hf_hub_download") as download, - pytest.raises(NotImplementedError, match="nemotron_h_moe"), - ): - _resolve_gguf_path(f"unsloth/nemotron:{filename}") - - preflight.assert_called_once_with("unsloth/nemotron", filename, revision="main") - download.assert_not_called() - def test_remote_deferred_audio_architecture_fails_before_download(self): from mobius.integrations.gguf._builder import _resolve_gguf_path from mobius.integrations.gguf._errors import UnsupportedGGUFArchitectureError @@ -6116,24 +6451,6 @@ def test_second_cohort_exact_closure_and_mutations(self, architecture: str) -> N ) ) - def test_nemotron_h_moe_schedule_rejects_with_actionable_error(self) -> None: - from mobius.integrations.gguf._builder import ( - _raise_for_invalid_hybrid_tensor_contract, - ) - - model = self._FakeGGUF( - "nemotron_h", - { - "nemotron_h.block_count": 2, - "nemotron_h.attention.head_count_kv": [0, 2], - "nemotron_h.feed_forward_length": [0, 128], - "nemotron_h.expert_count": 4, - }, - self._second_cohort_names("nemotron_h"), - ) - with pytest.raises(ValueError, match="MoE GGUF import is deferred"): - _raise_for_invalid_hybrid_tensor_contract(model) - def test_partial_fused_and_separate_experts_are_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 b3adb6e31..d8c151d09 100644 --- a/src/mobius/integrations/gguf/_config_mapping.py +++ b/src/mobius/integrations/gguf/_config_mapping.py @@ -25,6 +25,7 @@ import dataclasses import logging import re +from collections.abc import Iterable from types import MappingProxyType from typing import TYPE_CHECKING, Any @@ -93,10 +94,12 @@ "expert_used_count": "num_experts_per_tok", "expert_feed_forward_length": "moe_intermediate_size", "expert_shared_feed_forward_length": "shared_expert_intermediate_size", + "expert_shared_count": "n_shared_experts", "expert_weights_scale": "routed_scaling_factor", "expert_weights_norm": "norm_topk_prob", "expert_group_count": "n_group", "expert_group_used_count": "topk_group", + "moe_latent_size": "moe_latent_size", # Hybrid (DeltaNet / Mamba + Attention) fields "full_attention_interval": "full_attention_interval", # SSM/DeltaNet fields (used for linear attention in hybrid models) @@ -353,6 +356,7 @@ def _nextn_predict_layers(gguf_arch: str, metadata: dict[str, Any]) -> int: def _derive_hybrid_layout( gguf_arch: str, metadata: dict[str, Any], + tensor_names: Iterable[str] | None = None, ) -> tuple[int, list[str] | None, int]: """Derive the trunk layer count and exact mixer schedule from GGUF metadata.""" total_layers = int(metadata[f"{gguf_arch}.block_count"]) @@ -397,31 +401,35 @@ def _derive_hybrid_layout( mtp_count, ) - if gguf_arch == "nemotron_h": - if mtp_count: - raise ValueError("nemotron_h GGUF import does not support folded MTP blocks") + if gguf_arch in {"nemotron_h", "nemotron_h_moe"}: kv_raw = metadata.get(f"{gguf_arch}.attention.head_count_kv") ffn_raw = metadata.get(f"{gguf_arch}.feed_forward_length") if not isinstance(kv_raw, (list, tuple, np.ndarray)) or not isinstance( ffn_raw, (list, tuple, np.ndarray) ): raise ValueError( - "nemotron_h requires per-layer attention.head_count_kv and " + f"{gguf_arch} requires per-layer attention.head_count_kv and " "feed_forward_length arrays" ) kv_heads = [int(value) for value in kv_raw] ffn_lengths = [int(value) for value in ffn_raw] if len(kv_heads) != total_layers or len(ffn_lengths) != total_layers: raise ValueError( - f"nemotron_h schedule arrays must each contain exactly {total_layers} entries" + f"{gguf_arch} schedule arrays must each contain exactly {total_layers} entries" ) if any(value < 0 for value in (*kv_heads, *ffn_lengths)): - raise ValueError("nemotron_h schedule entries must be non-negative") + raise ValueError(f"{gguf_arch} schedule entries must be non-negative") uses_moe = int(metadata.get(f"{gguf_arch}.expert_count", 0)) > 0 + names = set(tensor_names) if tensor_names is not None else None layer_types = [] - for kv_heads_i, ffn_length_i in zip(kv_heads, ffn_lengths): + for layer, (kv_heads_i, ffn_length_i) in enumerate( + zip(kv_heads[:trunk_layers], ffn_lengths[:trunk_layers]) + ): if ffn_length_i: - layer_types.append("moe" if uses_moe else "mlp") + has_router = names is not None and f"blk.{layer}.ffn_gate_inp.weight" in names + layer_types.append( + "moe" if has_router or (names is None and uses_moe) else "mlp" + ) elif kv_heads_i == 0: layer_types.append("mamba2") else: @@ -574,17 +582,18 @@ def gguf_to_config( num_attention_heads = hf_fields["num_attention_heads"] if isinstance(num_attention_heads, (list, np.ndarray)): values = [int(value) for value in num_attention_heads] - if canonical_arch != "plamo2": + if canonical_arch not in {"plamo2", "nemotron_h", "nemotron_h_moe"}: raise ValueError( f"{canonical_arch} has unsupported per-layer attention head counts" ) nonzero = {value for value in values if value} if len(nonzero) != 1: raise ValueError( - "plamo2 attention layers must use one consistent non-zero head count" + f"{canonical_arch} attention layers must use one consistent non-zero " + "head count" ) if not nonzero: - raise ValueError("plamo2 GGUF has no attention layer") + raise ValueError(f"{canonical_arch} GGUF has no attention layer") num_attention_heads = nonzero.pop() if canonical_arch in {"mamba", "mamba2"}: # Pure recurrent GGUFs deliberately write attention.head_count=0. @@ -612,6 +621,7 @@ def gguf_to_config( "lfm2moe", "jamba", "nemotron_h", + "nemotron_h_moe", "granitehybrid", "plamo2", }: @@ -652,13 +662,14 @@ def gguf_to_config( "lfm2moe", "jamba", "nemotron_h", + "nemotron_h_moe", "granitehybrid", "plamo2", } or canonical_arch in _DELTA_NET_ARCHITECTURES ): derived_layers, layer_types, derived_mtp_count = _derive_hybrid_layout( - canonical_arch, metadata + canonical_arch, metadata, model.tensor_names ) if derived_layers != int(num_hidden_layers) or derived_mtp_count != mtp_predict_layers: raise ValueError("Hybrid schedule and decoder layer metadata disagree") @@ -667,12 +678,12 @@ def gguf_to_config( metadata.get(f"{canonical_arch}.full_attention_interval", 4) ) - if canonical_arch == "nemotron_h": + if canonical_arch in {"nemotron_h", "nemotron_h_moe"}: ffn_lengths = metadata[f"{gguf_arch}.feed_forward_length"] nonzero_ffn_lengths = {int(value) for value in ffn_lengths if int(value)} if len(nonzero_ffn_lengths) > 1: raise ValueError( - "nemotron_h dense FFN layers must use one consistent feed-forward length" + f"{gguf_arch} FFN layers must use one consistent feed-forward length" ) hf_fields["intermediate_size"] = ( nonzero_ffn_lengths.pop() if nonzero_ffn_lengths else 4 * hidden_size @@ -1502,30 +1513,89 @@ def _nemotron_h_postprocess( metadata: dict[str, Any], model: Any, ) -> NemotronHConfig: - """Build the dense, no-MTP Nemotron-H subset.""" - if int(metadata.get("nemotron_h.expert_count", 0)): + """Build the exact dense or routed-MoE Nemotron-H backbone.""" + arch = model.architecture + if arch not in {"nemotron_h", "nemotron_h_moe"}: + raise ValueError(f"Unexpected GGUF architecture for Nemotron-H: {arch!r}") + inner_size = int(metadata[f"{arch}.ssm.inner_size"]) + num_heads = int(metadata[f"{arch}.ssm.time_step_rank"]) + groups = int(metadata[f"{arch}.ssm.group_count"]) + if min(inner_size, num_heads) <= 0 or inner_size % num_heads: raise ValueError( - "Nemotron-H MoE GGUF import is deferred; fused softmax MoE is incompatible " - "with its sigmoid correction-bias routing" + f"{arch}.ssm.inner_size must be positive and divisible by ssm.time_step_rank" ) - inner_size = int(metadata["nemotron_h.ssm.inner_size"]) - num_heads = int(metadata["nemotron_h.ssm.time_step_rank"]) - if min(inner_size, num_heads) <= 0 or inner_size % num_heads: + if groups <= 0 or num_heads % groups: + raise ValueError(f"{arch}.ssm.time_step_rank must be divisible by ssm.group_count") + + num_experts = int(metadata.get(f"{arch}.expert_count", 0)) + top_k = int(metadata.get(f"{arch}.expert_used_count", 0)) + layer_types = list(config.layer_types or ()) + routed_layers = [i for i, layer_type in enumerate(layer_types) if layer_type == "moe"] + if arch == "nemotron_h_moe": + if num_experts <= 1 or not 1 <= top_k <= num_experts: + raise ValueError( + f"{arch} requires expert_count > 1 and expert_used_count in " + f"[1, expert_count], got {num_experts} and {top_k}" + ) + if not routed_layers: + raise ValueError(f"{arch} metadata declares experts but has no routed MoE layer") + if config.moe_intermediate_size is None or config.moe_intermediate_size <= 0: + raise ValueError(f"{arch}.expert_feed_forward_length must be greater than zero") + if ( + config.shared_expert_intermediate_size is None + or config.shared_expert_intermediate_size <= 0 + ): + raise ValueError( + f"{arch}.expert_shared_feed_forward_length must be greater than zero" + ) + shared_count = int(metadata.get(f"{arch}.expert_shared_count", 1)) + if shared_count != 1: + raise ValueError( + f"{arch}.expert_shared_count must be exactly 1, got {shared_count}" + ) + n_group = int(metadata.get(f"{arch}.expert_group_count", 1)) + topk_group = int(metadata.get(f"{arch}.expert_group_used_count", 1)) + if (n_group, topk_group) != (1, 1): + raise ValueError( + f"{arch} grouped expert routing is unsupported; " + f"expert_group_count and expert_group_used_count must both be 1, " + f"got {n_group} and {topk_group}" + ) + latent_size = metadata.get(f"{arch}.moe_latent_size") + if latent_size is not None and int(latent_size) <= 0: + raise ValueError(f"{arch}.moe_latent_size must be greater than zero when present") + elif num_experts or top_k or routed_layers: raise ValueError( - "nemotron_h.ssm.inner_size must be positive and divisible by ssm.time_step_rank" + "nemotron_h uses the dense architecture contract; routed experts require " + "general.architecture='nemotron_h_moe'" ) + fields = _shallow_fields(config) - fields["hidden_act"] = "relu2" + fields.update( + hidden_act="relu2", + layer_types=layer_types, + num_local_experts=num_experts or None, + num_experts_per_tok=top_k or None, + norm_topk_prob=bool(metadata.get(f"{arch}.expert_weights_norm", True)), + routed_scaling_factor=float(metadata.get(f"{arch}.expert_weights_scale", 1.0)), + n_group=1, + topk_group=1, + ) return NemotronHConfig( **fields, mamba_n_heads=num_heads, mamba_d_head=inner_size // num_heads, - mamba_d_state=int(metadata["nemotron_h.ssm.state_size"]), - mamba_n_groups=int(metadata["nemotron_h.ssm.group_count"]), - mamba_d_conv=int(metadata["nemotron_h.ssm.conv_kernel"]), + mamba_d_state=int(metadata[f"{arch}.ssm.state_size"]), + mamba_n_groups=groups, + mamba_d_conv=int(metadata[f"{arch}.ssm.conv_kernel"]), mamba_expand=inner_size // config.hidden_size, mamba_conv_bias=any(".ssm_conv1d.bias" in name for name in model.tensor_names), mamba_proj_bias=False, + moe_latent_size=( + int(metadata[f"{arch}.moe_latent_size"]) + if f"{arch}.moe_latent_size" in metadata + else None + ), ) @@ -2542,6 +2612,7 @@ def _eagle3_postprocess( "jamba": _jamba_postprocess, "lfm2moe": _lfm2moe_postprocess, "nemotron_h": _nemotron_h_postprocess, + "nemotron_h_moe": _nemotron_h_postprocess, "granitehybrid": _granitehybrid_postprocess, "bert_encoder": _bert_encoder_postprocess, "modern_bert_encoder": _modern_bert_encoder_postprocess, diff --git a/src/mobius/integrations/gguf/_config_mapping_test.py b/src/mobius/integrations/gguf/_config_mapping_test.py index 26b2e761b..38a055c62 100644 --- a/src/mobius/integrations/gguf/_config_mapping_test.py +++ b/src/mobius/integrations/gguf/_config_mapping_test.py @@ -66,9 +66,9 @@ def _metadata(architecture: str) -> dict: f"{architecture}.ssm.state_size": 8, f"{architecture}.ssm.time_step_rank": 8, } - if architecture in {"nemotron_h", "granitehybrid"}: + if architecture in {"nemotron_h", "nemotron_h_moe", "granitehybrid"}: metadata[f"{architecture}.ssm.group_count"] = 2 - if architecture == "nemotron_h": + if architecture in {"nemotron_h", "nemotron_h_moe"}: metadata[f"{architecture}.feed_forward_length"] = [0, 128, 0] return metadata @@ -145,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", ["nemotron_h", "granitehybrid"]) + @pytest.mark.parametrize("architecture", ["granitehybrid"]) def test_moe_modes_fail_closed(self, architecture: str) -> None: from mobius.integrations.gguf._config_mapping import gguf_to_config @@ -155,6 +155,76 @@ 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_nemotron_h_moe_derives_exact_routed_schedule_and_defaults(self) -> None: + from mobius.integrations.gguf._config_mapping import gguf_to_config + + metadata = self._metadata("nemotron_h_moe") + metadata.update( + { + "nemotron_h_moe.expert_count": 4, + "nemotron_h_moe.expert_used_count": 2, + "nemotron_h_moe.expert_feed_forward_length": 96, + "nemotron_h_moe.expert_shared_count": 1, + "nemotron_h_moe.expert_shared_feed_forward_length": 192, + "nemotron_h_moe.expert_weights_norm": True, + "nemotron_h_moe.expert_weights_scale": 2.5, + "nemotron_h_moe.moe_latent_size": 32, + } + ) + config = gguf_to_config( + _FakeDenseGGUF( + "nemotron_h_moe", + metadata, + [ + "token_embd.weight", + "blk.1.ffn_gate_inp.weight", + "blk.1.ffn_latent_down.weight", + "blk.1.ffn_latent_up.weight", + ], + ) + ) + assert config.layer_types == ["mamba2", "moe", "mamba2"] + assert config.num_local_experts == 4 + assert config.num_experts_per_tok == 2 + assert config.moe_intermediate_size == 96 + assert config.shared_expert_intermediate_size == 192 + assert config.moe_latent_size == 32 + assert config.norm_topk_prob is True + assert config.routed_scaling_factor == pytest.approx(2.5) + + def test_nemotron_h_moe_rejects_grouped_expert_routing(self) -> None: + from mobius.integrations.gguf._config_mapping import gguf_to_config + + metadata = self._metadata("nemotron_h_moe") + metadata.update( + { + "nemotron_h_moe.expert_count": 4, + "nemotron_h_moe.expert_used_count": 2, + "nemotron_h_moe.expert_feed_forward_length": 96, + "nemotron_h_moe.expert_shared_count": 1, + "nemotron_h_moe.expert_shared_feed_forward_length": 192, + "nemotron_h_moe.expert_group_count": 2, + "nemotron_h_moe.expert_group_used_count": 1, + } + ) + with pytest.raises(ValueError, match="grouped expert routing is unsupported"): + gguf_to_config( + _FakeDenseGGUF( + "nemotron_h_moe", + metadata, + ["token_embd.weight", "blk.1.ffn_gate_inp.weight"], + ) + ) + + def test_dense_nemotron_h_rejects_moe_metadata(self) -> None: + from mobius.integrations.gguf._config_mapping import gguf_to_config + + metadata = self._metadata("nemotron_h") + metadata["nemotron_h.expert_count"] = 4 + metadata["nemotron_h.expert_used_count"] = 2 + with pytest.raises(ValueError, match="nemotron_h_moe"): + gguf_to_config(_FakeDenseGGUF("nemotron_h", metadata, ["token_embd.weight"])) + def test_jamba_derives_exact_routed_layer_schedule(self) -> None: from mobius.integrations.gguf._config_mapping import gguf_to_config diff --git a/src/mobius/integrations/gguf/_mmproj.py b/src/mobius/integrations/gguf/_mmproj.py index fb73739d6..cc0b917a2 100644 --- a/src/mobius/integrations/gguf/_mmproj.py +++ b/src/mobius/integrations/gguf/_mmproj.py @@ -1091,6 +1091,11 @@ def build_gemma4_vlm_from_gguf( text_gguf = ( _text_gguf_model if _text_gguf_model is not None else GGUFModel(resolved_text_path) ) + if text_gguf.architecture != "gemma4": + raise ValueError( + "Gemma4 VLM package construction requires a gemma4 text GGUF; " + f"got {text_gguf.architecture!r}" + ) _validate_gguf_model(text_gguf, source=str(text_gguf_path)) preserve_quantization = keep_quantized and _has_quantized_weights(text_gguf, "gemma4") diff --git a/src/mobius/integrations/gguf/_mmproj_test.py b/src/mobius/integrations/gguf/_mmproj_test.py index 69f386145..4a8f9f52a 100644 --- a/src/mobius/integrations/gguf/_mmproj_test.py +++ b/src/mobius/integrations/gguf/_mmproj_test.py @@ -329,7 +329,7 @@ def test_clip_is_allowed_only_in_mmproj_companion_context( allow_mmproj_companion=True, ) - def test_rejects_unsupported_text_architecture_before_config( + def test_rejects_non_gemma4_text_architecture_before_config( self, tmp_path: Path, ): @@ -343,7 +343,7 @@ def test_rejects_unsupported_text_architecture_before_config( "mobius.integrations.gguf._mmproj._resolve_local_path", side_effect=[str(text_path)], ) as resolve, - pytest.raises(NotImplementedError, match="nemotron_h_moe"), + pytest.raises(ValueError, match="requires a gemma4 text GGUF"), ): build_gemma4_vlm_from_gguf(text_path, "owner/repo:mmproj.gguf") diff --git a/src/mobius/integrations/gguf/_tensor_mapping.py b/src/mobius/integrations/gguf/_tensor_mapping.py index 5a2bd3cdd..e5cd3f884 100644 --- a/src/mobius/integrations/gguf/_tensor_mapping.py +++ b/src/mobius/integrations/gguf/_tensor_mapping.py @@ -411,6 +411,14 @@ "blk.{bid}.ssm_norm": "backbone.layers.{bid}.mixer.norm", "blk.{bid}.ffn_up": "backbone.layers.{bid}.mixer.up_proj", "blk.{bid}.ffn_down": "backbone.layers.{bid}.mixer.down_proj", + "blk.{bid}.ffn_gate_inp": "backbone.layers.{bid}.mixer.gate", + "blk.{bid}.exp_probs_b": "backbone.layers.{bid}.mixer.gate.e_score_correction_bias@", + "blk.{bid}.ffn_up_exps": "backbone.layers.{bid}.mixer.experts.up_proj@", + "blk.{bid}.ffn_down_exps": "backbone.layers.{bid}.mixer.experts.down_proj@", + "blk.{bid}.ffn_up_shexp": "backbone.layers.{bid}.mixer.shared_experts.up_proj", + "blk.{bid}.ffn_down_shexp": "backbone.layers.{bid}.mixer.shared_experts.down_proj", + "blk.{bid}.ffn_latent_down": "backbone.layers.{bid}.mixer.fc1_latent_proj", + "blk.{bid}.ffn_latent_up": "backbone.layers.{bid}.mixer.fc2_latent_proj", } _GRANITEHYBRID_MAPPING: dict[str, str] = { @@ -779,6 +787,7 @@ "plamo2": _PLAMO2_MAPPING, "jamba": _JAMBA_MAPPING, "nemotron_h": _NEMOTRON_H_MAPPING, + "nemotron_h_moe": _NEMOTRON_H_MAPPING, "granitehybrid": _GRANITEHYBRID_MAPPING, "bert": _BERT_MAPPING, "modern_bert": _MODERN_BERT_MAPPING, diff --git a/src/mobius/integrations/gguf/_tensor_mapping_test.py b/src/mobius/integrations/gguf/_tensor_mapping_test.py index 74a2e9e44..2841b6881 100644 --- a/src/mobius/integrations/gguf/_tensor_mapping_test.py +++ b/src/mobius/integrations/gguf/_tensor_mapping_test.py @@ -92,6 +92,21 @@ def test_llama_family_aliases(self, arch: str) -> None: "blk.1.ssm_dt.bias", "backbone.layers.1.mixer.dt_bias", ), + ( + "nemotron_h_moe", + "blk.1.exp_probs_b.bias", + "backbone.layers.1.mixer.gate.e_score_correction_bias", + ), + ( + "nemotron_h_moe", + "blk.1.ffn_up_exps.weight", + "backbone.layers.1.mixer.experts.up_proj", + ), + ( + "nemotron_h_moe", + "blk.1.ffn_up_shexp.weight", + "backbone.layers.1.mixer.shared_experts.up_proj.weight", + ), ( "granitehybrid", "blk.0.ffn_gate.weight", diff --git a/src/mobius/integrations/gguf/_upstream_data/llamacpp_pin.json b/src/mobius/integrations/gguf/_upstream_data/llamacpp_pin.json index 362da63f0..0b9c9482d 100644 --- a/src/mobius/integrations/gguf/_upstream_data/llamacpp_pin.json +++ b/src/mobius/integrations/gguf/_upstream_data/llamacpp_pin.json @@ -118,7 +118,7 @@ "muse-glimmer": {"cohort": "C01-dense-transformer", "cpp_loader": true, "dual_moe": false}, "nanbeige": {"cohort": "C01-dense-transformer", "cpp_loader": true, "dual_moe": false, "tensor_families": ["token_embd", "output_norm", "output", "rope_freqs", "blk.{bid}.attn_norm", "blk.{bid}.attn_q", "blk.{bid}.attn_k", "blk.{bid}.attn_v", "blk.{bid}.attn_output", "blk.{bid}.attn_rot_embd", "blk.{bid}.ffn_norm", "blk.{bid}.ffn_gate", "blk.{bid}.ffn_down", "blk.{bid}.ffn_up"], "tensor_names": ["token_embd.weight", "output_norm.weight", "output.weight", "blk.{bid}.attn_norm.weight", "blk.{bid}.attn_output.weight", "rope_freqs.weight", "blk.{bid}.ffn_norm.weight", "blk.{bid}.ffn_gate.weight", "blk.{bid}.ffn_down.weight", "blk.{bid}.ffn_up.weight", "blk.{bid}.attn_qkv.weight", "blk.{bid}.attn_qkv.bias", "blk.{bid}.attn_q.weight", "blk.{bid}.attn_q.bias", "blk.{bid}.attn_k.weight", "blk.{bid}.attn_k.bias", "blk.{bid}.attn_v.weight", "blk.{bid}.attn_v.bias"], "tensor_closure_status": "audited-direct-loader-conditional-union", "converter_inventory_status": "exact-pinned-MODEL_TENSORS-family-inventory", "loader_source": "src/models/nanbeige.cpp", "required_metadata": ["attention.layer_norm_rms_epsilon"], "optional_metadata": ["num_loops", "skip_loop_final_norm"], "loader_helpers": ["create_tensor_qkv"]}, "nemotron": {"cohort": "C01-dense-transformer", "cpp_loader": true, "dual_moe": false}, - "nemotron_h": {"cohort": "C06-hybrid-attn-ssm", "cpp_loader": true, "dual_moe": true, "tensor_names": ["token_embd.weight", "output_norm.weight", "output.weight", "blk.{bid}.attn_norm.weight", "blk.{bid}.ssm_in.weight", "blk.{bid}.ssm_conv1d.weight", "blk.{bid}.ssm_conv1d.bias", "blk.{bid}.ssm_dt.bias", "blk.{bid}.ssm_a", "blk.{bid}.ssm_d", "blk.{bid}.ssm_norm.weight", "blk.{bid}.ssm_out.weight", "blk.{bid}.attn_q.weight", "blk.{bid}.attn_k.weight", "blk.{bid}.attn_v.weight", "blk.{bid}.attn_output.weight", "blk.{bid}.attn_output.bias", "blk.{bid}.ffn_up.weight", "blk.{bid}.ffn_up.bias", "blk.{bid}.ffn_down.weight", "blk.{bid}.ffn_down.bias", "blk.{bid}.ffn_gate_inp.weight", "blk.{bid}.ffn_exp_probs_b.bias", "blk.{bid}.ffn_latent_down.weight", "blk.{bid}.ffn_latent_up.weight", "blk.{bid}.ffn_up_exps.weight", "blk.{bid}.ffn_down_exps.weight", "blk.{bid}.ffn_up_shexp.weight", "blk.{bid}.ffn_down_shexp.weight"]}, + "nemotron_h": {"cohort": "C06-hybrid-attn-ssm", "cpp_loader": true, "dual_moe": true, "tensor_names": ["token_embd.weight", "output_norm.weight", "output.weight", "blk.{bid}.attn_norm.weight", "blk.{bid}.ssm_in.weight", "blk.{bid}.ssm_conv1d.weight", "blk.{bid}.ssm_conv1d.bias", "blk.{bid}.ssm_dt.bias", "blk.{bid}.ssm_a", "blk.{bid}.ssm_d", "blk.{bid}.ssm_norm.weight", "blk.{bid}.ssm_out.weight", "blk.{bid}.attn_q.weight", "blk.{bid}.attn_k.weight", "blk.{bid}.attn_v.weight", "blk.{bid}.attn_output.weight", "blk.{bid}.attn_output.bias", "blk.{bid}.ffn_up.weight", "blk.{bid}.ffn_up.bias", "blk.{bid}.ffn_down.weight", "blk.{bid}.ffn_down.bias", "blk.{bid}.ffn_gate_inp.weight", "blk.{bid}.exp_probs_b.bias", "blk.{bid}.ffn_latent_down.weight", "blk.{bid}.ffn_latent_up.weight", "blk.{bid}.ffn_up_exps.weight", "blk.{bid}.ffn_down_exps.weight", "blk.{bid}.ffn_up_shexp.weight", "blk.{bid}.ffn_down_shexp.weight"]}, "nemotron_h_moe": {"cohort": "C06-hybrid-attn-ssm", "cpp_loader": true, "dual_moe": true}, "neo-bert": {"cohort": "C03-encoder-embed-rerank", "cpp_loader": true, "dual_moe": false}, "nomic-bert": {"cohort": "C03-encoder-embed-rerank", "cpp_loader": true, "dual_moe": true}, diff --git a/tests/_test_configs.py b/tests/_test_configs.py index 7fce0b0d1..f07350ea0 100644 --- a/tests/_test_configs.py +++ b/tests/_test_configs.py @@ -1475,6 +1475,31 @@ def _base_config(config_cls=None, **overrides) -> ArchitectureConfig: }, True, ), + # Nemotron-H latent MoE: gate/shared expert use hidden states while routed + # experts operate in the projected latent space. + ( + "nemotron_h", + { + "hidden_act": "relu2", + "layer_types": ["mamba2", "moe", "full_attention", "moe"], + "_config_cls": NemotronHConfig, + "num_hidden_layers": 4, + "mamba_n_heads": TINY_KV_HEADS, + "mamba_d_head": TINY_HEAD_DIM, + "mamba_d_state": 16, + "mamba_n_groups": 1, + "mamba_d_conv": 4, + "mamba_expand": 2, + "num_local_experts": 4, + "num_experts_per_tok": 2, + "moe_intermediate_size": TINY_INTERMEDIATE, + "moe_latent_size": TINY_HIDDEN // 2, + "shared_expert_intermediate_size": TINY_INTERMEDIATE * 2, + "norm_topk_prob": True, + "routed_scaling_factor": 2.5, + }, + True, + ), # nemotron_h MoE variant: hybrid Mamba2+MoE+Attention (Nemotron-3 30B/120B) ( "nemotron_h", diff --git a/tests/synthetic_parity_test.py b/tests/synthetic_parity_test.py index a698408ce..9fd89e0a5 100644 --- a/tests/synthetic_parity_test.py +++ b/tests/synthetic_parity_test.py @@ -232,7 +232,6 @@ # NemotronH Mamba2 layers diverge (cos=0.65): LinearAttention gated-SSM # recurrence on CPU produces different results than HF's naive Mamba2. # Attention-only layers match perfectly (cos=0.9999). - "nemotron_h": "Mamba2 SSM recurrence diverges on CPU (LinearAttention vs HF naive)", } # Fields that are properties in HF configs and cannot be set directly, @@ -602,12 +601,16 @@ def _create_hf_config(model_type: str, config_overrides: dict): for lt in layer_types ] - # NemotronH uses layers_block_type with HF values {"mamba", "attention", "moe"}. - # Convert our internal layer_types names (mamba2, full_attention, mlp) to HF names. + # NemotronH uses its public layer-type vocabulary rather than Mobius names. # Also translate mobius Mamba field names to HF NemotronHConfig field names. if hf_model_type in ("nemotron_h",) and "layer_types" in hf_kwargs: layer_types = hf_kwargs.pop("layer_types") - _nemotron_type_map = {"mamba2": "mamba", "full_attention": "attention", "mlp": "moe"} + _nemotron_type_map = { + "mamba2": "mamba", + "full_attention": "attention", + "mlp": "mlp", + "moe": "moe", + } hf_kwargs["layers_block_type"] = [_nemotron_type_map.get(lt, lt) for lt in layer_types] # Mobius NemotronHConfig → HF NemotronHConfig field name mapping _nemotron_field_map = { @@ -617,10 +620,17 @@ def _create_hf_config(model_type: str, config_overrides: dict): "mamba_n_groups": "n_groups", "mamba_d_conv": "conv_kernel", "mamba_expand": "expand", + "rms_norm_eps": "layer_norm_epsilon", } for old_name, new_name in _nemotron_field_map.items(): if old_name in hf_kwargs: hf_kwargs[new_name] = hf_kwargs.pop(old_name) + if "hidden_act" in hf_kwargs: + hf_kwargs["mlp_hidden_act"] = hf_kwargs.pop("hidden_act") + if "shared_expert_intermediate_size" in hf_kwargs: + hf_kwargs["moe_shared_expert_intermediate_size"] = hf_kwargs.pop( + "shared_expert_intermediate_size" + ) # HF NemotronH has an explicit head_dim (default 128) that is not # derived from hidden_size / num_attention_heads. Set it to match. if "head_dim" not in hf_kwargs: @@ -686,6 +696,7 @@ def _create_hf_config(model_type: str, config_overrides: dict): "num_experts_per_tok": "moe_topk", "moe_intermediate_size": "expert_ffn_hidden_size", }, + "nemotron_h": {"num_local_experts": "n_routed_experts"}, } if hf_model_type in expert_field_aliases: for src_field, dst_field in expert_field_aliases[hf_model_type].items(): @@ -1171,6 +1182,23 @@ def test_synthetic_parity(model_type: str, config_overrides: dict): # 3. Create HF model hf_config = _create_hf_config(model_type, config_overrides) hf_model = _create_hf_model(model_type, hf_config, seed) + if model_type == "nemotron_h": + # Force correction bias to determine the selected experts. This catches + # implementations that incorrectly use biased scores as final weights. + for layer in hf_model.model.layers: + if getattr(layer, "block_type", None) != "moe": + continue + bias = layer.mixer.gate.e_score_correction_bias + with torch.no_grad(): + bias.copy_( + torch.linspace( + 4.0, + 1.0, + bias.numel(), + dtype=bias.dtype, + device=bias.device, + ) + ) # 4. Transfer HF weights to ONNX try: