From b0b94f92cd39aea883bdbc008a3ff3ff31b56121 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Fri, 14 Aug 2026 20:23:27 -0700 Subject: [PATCH] Support Nemotron 3.5 Q8 GGUF conversion Add a reusable GGUF architecture-adapter seam and a strict nemotron_h_moe adapter that validates the pinned 52-layer backbone, excludes the auxiliary MTP block, and enforces complete source-to-initializer mapping. Preserve Q8_0 weights in MatMulNBits, including stacked routed experts, reconstruct the pinned Pixtral tokenizer contract, and add a fresh-process direct-ORT acceptance runner with exact llama.cpp generation evidence. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu --- docs/api/build_from_gguf.md | 230 +++--- docs/cli_reference.md | 8 +- .../validate_gguf_q8.py | 665 +++++++++++++++++ src/mobius/components/_mamba_block.py | 7 +- src/mobius/integrations/gguf/_architecture.py | 189 +++++ .../integrations/gguf/_architecture_test.py | 406 ++++++++++ src/mobius/integrations/gguf/_builder.py | 443 +++++++---- src/mobius/integrations/gguf/_builder_test.py | 111 +-- .../integrations/gguf/_config_mapping.py | 27 +- src/mobius/integrations/gguf/_mmproj_test.py | 4 +- .../integrations/gguf/_nemotron_h_moe.py | 697 ++++++++++++++++++ src/mobius/integrations/gguf/_repacker.py | 62 ++ .../integrations/gguf/_repacker_test.py | 100 +++ src/mobius/integrations/gguf/_tokenizer.py | 137 +++- .../integrations/gguf/_tokenizer_test.py | 242 ++++++ src/mobius/models/nemotron_h.py | 76 +- src/mobius/models/nemotron_h_test.py | 122 +++ 17 files changed, 3119 insertions(+), 407 deletions(-) create mode 100644 examples/olive/nemotron-3_5-lightning-30b/validate_gguf_q8.py create mode 100644 src/mobius/integrations/gguf/_architecture.py create mode 100644 src/mobius/integrations/gguf/_architecture_test.py create mode 100644 src/mobius/integrations/gguf/_nemotron_h_moe.py create mode 100644 src/mobius/models/nemotron_h_test.py diff --git a/docs/api/build_from_gguf.md b/docs/api/build_from_gguf.md index 82b1afdce..f278199c5 100644 --- a/docs/api/build_from_gguf.md +++ b/docs/api/build_from_gguf.md @@ -91,163 +91,137 @@ LLM architectures are supported. Sharded GGUF files are rejected. A single shard has only part of the tensor table, and treating it as a complete checkpoint would create a corrupt model. -## NVIDIA Nemotron 3.5 Lightning waiver +## NVIDIA Nemotron 3.5 Lightning Q8_0 -Direct conversion of GGUF architecture `nemotron_h_moe` is intentionally -disabled. The following evidence is pinned: +Mobius supports the pinned single-file `Q8_0` GGUF production slice: -- 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. +- Repository: `unsloth/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-GGUF` +- Revision: `f2d3fe3694501008786e81e5f20360cbf715496a` +- File: `NVIDIA-Nemotron-3.5-Lightning-30B-A3B-Q8_0.gguf` +- Size: `35,004,643,392` bytes +- SHA-256: `dc5276dd0619c04e277504d2358a793e31ccbe39e894d767d0d14f2a221e2ca4` -### Quantization findings +The architecture adapter validates the complete 417-tensor header before graph +construction. It maps exactly 401 backbone sources, explicitly excludes the +16 tensors in auxiliary MTP block 52, and produces 6,243 logical decoder +weights. The backbone schedule must be exactly 23 Mamba + 23 MoE + 6 attention +layers. Block 52 remains a separate combined attention+MoE MTP block rather +than becoming a false 53rd decoder layer. -| 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: +### Pinned build ```powershell $repo = "unsloth/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-GGUF" $revision = "f2d3fe3694501008786e81e5f20360cbf715496a" $file = "NVIDIA-Nemotron-3.5-Lightning-30B-A3B-Q8_0.gguf" +$official = "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16" +$officialRevision = "d468880b6ad3c6e0d21377ce7242adaea4cc884d" python -m pip install ` --index-url https://packagefeedproxy.microsoft.io/pypi/simple ` huggingface_hub hf download $repo $file --revision $revision --local-dir .\nemotron-gguf +hf download $official tokenizer.json tokenizer_config.json ` + special_tokens_map.json chat_template.jinja ` + --revision $officialRevision --local-dir .\nemotron-tokenizer -# 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 + --external-data onnx --output .\nemotron-gguf-onnx +Copy-Item .\nemotron-tokenizer\tokenizer_config.json .\nemotron-gguf-onnx\ +Copy-Item .\nemotron-tokenizer\special_tokens_map.json .\nemotron-gguf-onnx\ +Copy-Item .\nemotron-tokenizer\chat_template.jinja .\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. +The Q8 blocks are affine-repacked exactly into +`MatMulNBits(bits=8, block_size=32)`. Routed expert tensors are expanded along +their leading expert axis without dequantization. The resulting graph has +6,005 `MatMulNBits` nodes and one `GatherBlockQuantized` embedding node, with +no `QuantizeLinear`/`DequantizeLinear` round trip. -To execute the pinned GGUF without changing its quantization, use current -llama.cpp instead: +On the 63.3 GiB Windows acceptance host, the pinned build completed in +227.001 seconds plus 145.792 seconds to save. The package is +36,920,438,736 bytes and the build process peaked at 54,818,070,528 bytes of +working set. The weighted graph contains all 18,255 mapped weight +initializers, including 6,006 Q8 weights. -```powershell -.\llama-cli.exe ` - --model ".\nemotron-gguf\$file" ` - --temp 0.6 --top-p 0.95 --min-p 0.01 -``` +### Tokenizer and runtime contract + +The embedded GPT-2/Pixtral vocabulary and merges are reconstructed as a +ByteLevel tokenizer. The source metadata's padding ID 999 is rejected because +it names ``, not the model's runtime padding convention. The +package keeps two explicit contracts: + +- Pinned tokenizer asset: BOS 1 and `<|im_end|>` as EOS/padding ID 11. +- Direct-ORT model contract: padding ID 0 and EOS IDs `[2, 11]`. + +The reconstructed tokenizer has the exact official vocabulary, pre-tokenizer, +decoder, post-processor, special-token flags, and encode/decode behavior. The +GGUF's embedded chat template is not the template at the pinned official +revision, so it is not silently emitted as authoritative. The recipe copies +the official `tokenizer_config.json`, `special_tokens_map.json`, and +`chat_template.jinja` sidecars by immutable revision instead. -### Option A: official BF16, then Olive +The direct-ORT runner uses unpadded prompts. Its validation also compares every +real-token logit from an unpadded prefill with the same prompt followed by +right padding and an explicit attention mask. It never continues cached +generation from recurrent states advanced through padding. -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: +Generic ORT GenAI configuration cannot currently bind arbitrary non-KV +recurrent cache inputs such as the graph's convolution and SSM states. Use +direct ONNX Runtime generation rather than treating session creation as +generation evidence. + +### Reproduce semantic acceptance + +The validator builds and saves in one process, then loads and generates in a +fresh process. It records package/operator/initializer counts, mapping +completeness, runtime versions, timings, and peak working sets: ```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 ` - .\nemotron-bf16-onnx +python examples\olive\nemotron-3_5-lightning-30b\validate_gguf_q8.py ` + --phase all ` + --gguf ".\nemotron-gguf\$file" ` + --official-tokenizer-dir .\nemotron-tokenizer ` + --output .\nemotron-gguf-onnx ` + --device cpu ``` -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" -} +The independent greedy reference was produced by llama.cpp commit +`9d57ce456c94d241dde672b2db9cf18879766568` from prompt +`The capital of France is`: + +```text +llama-server -m NVIDIA-Nemotron-3.5-Lightning-30B-A3B-Q8_0.gguf \ + -c 128 -t 12 -tb 12 -b 64 -ub 64 -ngl 0 \ + --host 127.0.0.1 --port 18081 --no-warmup +POST /completion +{"prompt":"The capital of France is","n_predict":8,"temperature":0, + "seed":1,"cache_prompt":false,"n_probs":1} ``` -```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\ -``` +- Prompt IDs: `[1784, 8961, 1307, 5498, 1395]` +- Generated IDs: `[6993, 1046, 1256, 1010, 1784, 8961, 1307, 10787]` +- Text: ` Paris. \nThe capital of Germany` +- Throughput: 5.58 prompt tok/s and 5.95 generated tok/s +- Peak working set from a separate CLI run: 16.81 GiB -The candidate can be checked for direct ORT session loading: +The fresh-process ONNX Runtime 1.28.0 CPU run at commit `45de2a8b06` +loaded the package in 38.329 seconds and completed the five-token prefill in +99.447 seconds. Its seven cached decode calls ran at 0.01617 steps/s and the +process peaked at 38,148,943,872 bytes of working set. The right-padded +explicit-mask check had zero maximum absolute difference across all real-token +logits, and cached greedy generation matched every llama.cpp token and the +decoded text above. -```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()]) -``` +### Explicitly unsupported variants + +| GGUF file | Relevant tensor inventory | Status | +|---|---|---| +| `...-Q8_0.gguf` | 32.904B parameters in `Q8_0` | Supported through exact affine repacking | +| `...-MXFP4_MOE.gguf` | 14.687B `MXFP4`, 12.772B `Q5_1`, 5.445B `Q8_0` | Rejected; the large `Q5_1` source has no validated preserved runtime mapping | +| `...-UD-Q4_K_M.gguf` | 15.326B `Q5_0`, 12.772B `Q5_1`, 4.806B `Q8_0` | Rejected; the preset name does not describe its large source qtypes | +| `BF16/...-0000*-of-00002.gguf` | 329 tensors in shard 1 and 88 tensors in shard 2 | Rejected by the generic shard-assembly guard | -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. +Dequantizing and requantizing Q5-family tensors would change their +quantization and is not described as preservation. diff --git a/docs/cli_reference.md b/docs/cli_reference.md index b4f516eb5..0eae2302e 100644 --- a/docs/cli_reference.md +++ b/docs/cli_reference.md @@ -343,10 +343,10 @@ Quantized files containing only qtypes with no supported preservation target float. Re-run with `--dequantize` to request explicit float conversion. 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. The pinned Nemotron 3.5 Lightning single-file Q8_0 artifact is +supported with exact Q8 repacking; mixed Q5 variants and split BF16 GGUFs +remain rejected. See +[`build_from_gguf()`](api/build_from_gguf.md#nvidia-nemotron-35-lightning-q8_0). --- diff --git a/examples/olive/nemotron-3_5-lightning-30b/validate_gguf_q8.py b/examples/olive/nemotron-3_5-lightning-30b/validate_gguf_q8.py new file mode 100644 index 000000000..defd6bee9 --- /dev/null +++ b/examples/olive/nemotron-3_5-lightning-30b/validate_gguf_q8.py @@ -0,0 +1,665 @@ +#!/usr/bin/env python +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Build and validate the pinned Nemotron 3.5 Lightning Q8_0 GGUF package.""" + +from __future__ import annotations + +import argparse +import gc +import hashlib +import json +import math +import shutil +import subprocess +import sys +import time +from collections import Counter +from pathlib import Path +from typing import Any + +import numpy as np +from inference import ( + _as_numpy, + _create_session, + _initial_states, + _run_session, + _token_feeds, + _update_states, +) + +GGUF_REPO = "unsloth/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-GGUF" +GGUF_REVISION = "f2d3fe3694501008786e81e5f20360cbf715496a" +GGUF_FILENAME = "NVIDIA-Nemotron-3.5-Lightning-30B-A3B-Q8_0.gguf" +GGUF_SIZE = 35_004_643_392 +GGUF_SHA256 = "dc5276dd0619c04e277504d2358a793e31ccbe39e894d767d0d14f2a221e2ca4" +LLAMA_CPP_COMMIT = "9d57ce456c94d241dde672b2db9cf18879766568" +LLAMA_CPP_SERVER_COMMAND = ( + f"llama-server -m {GGUF_FILENAME} -c 128 -t 12 -tb 12 -b 64 -ub 64 " + "-ngl 0 --host 127.0.0.1 --port 18081 --no-warmup" +) +PROMPT = "The capital of France is" +PROMPT_IDS = [1784, 8961, 1307, 5498, 1395] +EXPECTED_IDS = [6993, 1046, 1256, 1010, 1784, 8961, 1307, 10787] +EXPECTED_TEXT = " Paris. \nThe capital of Germany" +LLAMA_CPP_REFERENCE = { + "commit": LLAMA_CPP_COMMIT, + "compiler": "MSVC 19.44.35228.0", + "server_command": LLAMA_CPP_SERVER_COMMAND, + "completion_request": { + "prompt": PROMPT, + "n_predict": 8, + "temperature": 0, + "seed": 1, + "cache_prompt": False, + "n_probs": 1, + }, + "prompt_ids": PROMPT_IDS, + "generated_ids": EXPECTED_IDS, + "generated_text": EXPECTED_TEXT, + "prompt_tokens_per_second": 5.58, + "generation_tokens_per_second": 5.95, + "peak_working_set_gib": 16.81, +} +OFFICIAL_TOKENIZER_SHA256 = { + "chat_template.jinja": "58933db77d3099b4f78c55a38347a72e1ea05b97d6bd8f38775303dc0194e0a9", + "special_tokens_map.json": ( + "e9435fefd6d838fd9fcbbc44b97a8e3ff322be7f6dfb7e4fd2468586574bb52b" + ), + "tokenizer_config.json": ( + "10f93eabcb9b1602fbb991d6308e787ce1df28ee9cd7a1c6d1e8c3f338b957bc" + ), + "tokenizer.json": "623c34567aebb18582765289fbe23d901c62704d6518d71866e0e58db892b5b7", +} +GGUF_CHAT_TEMPLATE_SHA256 = "cbb337473ffde036fd4b6e7e7763dcb97c7cd8b4a311cd52d361d2766b00eb7c" + + +def _memory_sample() -> dict[str, int]: + try: + import psutil + except ImportError: + import resource + + rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss * 1024 + return {"rss": rss, "peak_working_set": rss} + + info = psutil.Process().memory_info() + return { + "rss": int(info.rss), + "peak_working_set": int(getattr(info, "peak_wset", info.rss)), + } + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + while chunk := stream.read(8 * 1024 * 1024): + digest.update(chunk) + return digest.hexdigest() + + +def _write_json(path: Path, value: Any) -> None: + path.write_text(json.dumps(value, indent=2, sort_keys=True) + "\n", encoding="utf-8") + + +def _write_tokenizer_contract( + gguf_path: Path, + official_dir: Path, + output_dir: Path, +) -> dict[str, Any]: + from tokenizers import Tokenizer + + from mobius.integrations.gguf._reader import GGUFModel + + metadata = GGUFModel(gguf_path).metadata + tokens = metadata["tokenizer.ggml.tokens"] + expected_tokens = {0: "", 1: "", 2: "", 11: "<|im_end|>"} + actual_tokens = {index: tokens[index] for index in expected_tokens} + if actual_tokens != expected_tokens: + raise ValueError(f"Unexpected pinned special-token strings: {actual_tokens}") + + actual_hashes = {} + for filename, expected_sha256 in OFFICIAL_TOKENIZER_SHA256.items(): + source = official_dir / filename + if not source.is_file(): + raise FileNotFoundError(f"Missing pinned official tokenizer asset: {source}") + actual_hashes[filename] = _sha256(source) + if actual_hashes[filename] != expected_sha256: + raise ValueError( + f"Pinned official tokenizer asset {filename!r} has SHA-256 " + f"{actual_hashes[filename]}, expected {expected_sha256}" + ) + + rebuilt_path = output_dir / "tokenizer.json" + rebuilt = Tokenizer.from_file(str(rebuilt_path)) + official = Tokenizer.from_file(str(official_dir / "tokenizer.json")) + if rebuilt.get_vocab(with_added_tokens=True) != official.get_vocab(with_added_tokens=True): + raise ValueError( + "Reconstructed tokenizer vocabulary differs from the pinned official asset" + ) + samples = ( + PROMPT, + EXPECTED_TEXT, + "café déjà vu — 你好 🌍", + " leading\tspaces\r\nnewlines ", + "<|im_start|>assistant\nx<|im_end|>", + bytes(range(1, 128)).decode("latin1"), + ) + for sample in samples: + official_encoding = official.encode(sample) + rebuilt_encoding = rebuilt.encode(sample) + if rebuilt_encoding.ids != official_encoding.ids: + raise ValueError(f"Reconstructed tokenizer IDs differ for {sample!r}") + if rebuilt.decode( + rebuilt_encoding.ids, + skip_special_tokens=False, + ) != official.decode(official_encoding.ids, skip_special_tokens=False): + raise ValueError(f"Reconstructed tokenizer decode differs for {sample!r}") + + official_json = json.loads((official_dir / "tokenizer.json").read_text(encoding="utf-8")) + rebuilt_json = json.loads(rebuilt_path.read_text(encoding="utf-8")) + for section in ("pre_tokenizer", "decoder", "post_processor", "added_tokens"): + if rebuilt_json[section] != official_json[section]: + raise ValueError(f"Reconstructed tokenizer {section!r} differs from official") + + chat_template = metadata.get("tokenizer.chat_template") + if not isinstance(chat_template, str) or not chat_template: + raise ValueError("Pinned GGUF tokenizer has no chat template") + gguf_template_sha256 = hashlib.sha256( + chat_template.replace("\r\n", "\n").encode() + ).hexdigest() + if gguf_template_sha256 != GGUF_CHAT_TEMPLATE_SHA256: + raise ValueError(f"Pinned GGUF chat-template SHA-256 differs: {gguf_template_sha256}") + if gguf_template_sha256 == actual_hashes["chat_template.jinja"]: + raise ValueError( + "GGUF and official chat templates unexpectedly have the same provenance" + ) + + for filename in ( + "chat_template.jinja", + "special_tokens_map.json", + "tokenizer_config.json", + ): + shutil.copyfile(official_dir / filename, output_dir / filename) + return { + "asset_bos_token_id": 1, + "asset_eos_token_id": 11, + "asset_padding_token_id": 11, + "asset_unk_token_id": 0, + "reconstructed_matches_official_tokenizer": True, + "official_revision": "d468880b6ad3c6e0d21377ce7242adaea4cc884d", + "official_asset_sha256": actual_hashes, + "gguf_chat_template_sha256_rejected": gguf_template_sha256, + } + + +def _graph_audit(model, expected_weight_initializers: set[str]) -> dict[str, Any]: + op_counts = Counter( + f"{node.domain or 'ai.onnx'}::{node.op_type}" for node in model.graph.all_nodes() + ) + initializers = model.graph.initializers + quantized_weights = [ + name + for name, value in initializers.items() + if value.dtype.name == "UINT8" and name.endswith((".weight", ".qweight")) + ] + unset = [name for name, value in initializers.items() if value.const_value is None] + if unset: + raise ValueError(f"Weighted graph has {len(unset)} unset initializers: {unset[:10]}") + missing_weights = expected_weight_initializers - set(initializers) + folded_gate_weights = { + name for name in expected_weight_initializers if name.endswith(".moe.gate.weight") + } + if missing_weights != folded_gate_weights: + raise ValueError( + f"Weighted graph is missing {len(missing_weights)} mapped initializers: " + f"{sorted(missing_weights)[:10]}" + ) + # Gate matmuls explicitly cast their weights to float32 to preserve the + # official routing contract. Constant folding consumes those 23 source + # names and materializes replacement constants in the weighted graph. + post_fold_initializers = sorted(set(initializers) - expected_weight_initializers) + if op_counts["com.microsoft::MatMulNBits"] != 6005: + raise ValueError(f"Unexpected MatMulNBits count: {dict(op_counts)}") + if op_counts["com.microsoft::GatherBlockQuantized"] != 1: + raise ValueError(f"Unexpected GatherBlockQuantized count: {dict(op_counts)}") + if len(expected_weight_initializers) != 18_255: + raise ValueError( + "Expected 18,255 mapped weight initializers, got " + f"{len(expected_weight_initializers)}" + ) + if len(quantized_weights) != 6006: + raise ValueError( + f"Expected 6,006 Q8 weight initializers, got {len(quantized_weights)}" + ) + forbidden = { + name: count + for name, count in op_counts.items() + if name.endswith(("::QuantizeLinear", "::DequantizeLinear")) + } + if forbidden: + raise ValueError(f"Unexpected dequantize/requantize operators: {forbidden}") + return { + "op_histogram": dict(sorted(op_counts.items())), + "initializer_count": len(initializers), + "mapped_weight_initializer_count": len(expected_weight_initializers), + "folded_gate_weight_count": len(folded_gate_weights), + "post_fold_initializer_count": len(post_fold_initializers), + "post_fold_initializers": post_fold_initializers, + "q8_weight_initializer_count": len(quantized_weights), + "matmul_nbits_count": op_counts["com.microsoft::MatMulNBits"], + "gather_block_quantized_count": op_counts["com.microsoft::GatherBlockQuantized"], + "forbidden_qdq_ops": forbidden, + } + + +def _mapping_audit(gguf_path: Path) -> tuple[dict[str, Any], set[str]]: + from mobius.integrations.gguf._architecture import ( + GGUFMappingAudit, + create_architecture_adapter, + ) + from mobius.integrations.gguf._reader import GGUFModel + + source = GGUFModel(gguf_path) + adapter = create_architecture_adapter(source.architecture, source) + if adapter is None: + raise ValueError(f"No adapter for {source.architecture!r}") + adapter.validate_model(source=str(gguf_path)) + audit = GGUFMappingAudit() + qtypes: Counter[str] = Counter() + base_qtypes: Counter[str] = Counter() + mtp_qtypes: Counter[str] = Counter() + qtype_parameters: Counter[str] = Counter() + base_qtype_parameters: Counter[str] = Counter() + mtp_qtype_parameters: Counter[str] = Counter() + q8_targets = 0 + expected_initializers: set[str] = set() + for record in source._reader.tensors: + shape = tuple(int(dim) for dim in reversed(record.shape)) + mapping = adapter.map_tensor(record.name, shape) + audit.record(record.name, mapping) + qtype = record.tensor_type.name + parameters = math.prod(shape) + qtypes[qtype] += 1 + qtype_parameters[qtype] += parameters + (mtp_qtypes if record.name.startswith("blk.52.") else base_qtypes)[qtype] += 1 + (mtp_qtype_parameters if record.name.startswith("blk.52.") else base_qtype_parameters)[ + qtype + ] += parameters + if mapping is not None and mapping.exclusion is None and qtype == "Q8_0": + q8_targets += len(mapping.targets) + for target in mapping.targets: + stem = target.initializer_name.removesuffix(".weight") + expected_initializers.add( + f"{stem}.qweight" + if target.initializer_name == "model.embed_tokens.weight" + else target.initializer_name + ) + expected_initializers.add(f"{stem}.scales") + expected_initializers.add(f"{stem}.zero_points") + elif mapping is not None and mapping.exclusion is None: + expected_initializers.update(target.initializer_name for target in mapping.targets) + adapter.validate_mapping_audit(audit) + metadata = source.metadata + tokens = metadata["tokenizer.ggml.tokens"] + merges = metadata["tokenizer.ggml.merges"] + return { + "source_tensors": len(source.tensor_names), + "mapped_sources": len(audit.mapped_sources), + "mtp_exclusions": len(audit.excluded_sources), + "logical_targets": len(audit.target_sources), + "q8_logical_targets": q8_targets, + "mapped_weight_initializers": len(expected_initializers), + "qtype_tensors": dict(sorted(qtypes.items())), + "base_qtype_tensors": dict(sorted(base_qtypes.items())), + "mtp_qtype_tensors": dict(sorted(mtp_qtypes.items())), + "qtype_parameters": dict(sorted(qtype_parameters.items())), + "base_qtype_parameters": dict(sorted(base_qtype_parameters.items())), + "mtp_qtype_parameters": dict(sorted(mtp_qtype_parameters.items())), + "tokenizer": { + "profile": [ + metadata["tokenizer.ggml.model"], + metadata["tokenizer.ggml.pre"], + ], + "bos_token_id": metadata["tokenizer.ggml.bos_token_id"], + "eos_token_id": metadata["tokenizer.ggml.eos_token_id"], + "rejected_gguf_padding_token_id": metadata["tokenizer.ggml.padding_token_id"], + "token_count": len(tokens), + "merge_count": len(merges), + "token_sha256": hashlib.sha256("\n".join(tokens).encode()).hexdigest(), + "merge_sha256": hashlib.sha256("\n".join(merges).encode()).hexdigest(), + }, + }, expected_initializers + + +def _build(gguf_path: Path, official_tokenizer_dir: Path, output_dir: Path) -> None: + if gguf_path.stat().st_size != GGUF_SIZE: + raise ValueError( + f"Pinned GGUF size mismatch: {gguf_path.stat().st_size} != {GGUF_SIZE}" + ) + source_sha256 = _sha256(gguf_path) + if source_sha256 != GGUF_SHA256: + raise ValueError(f"Pinned GGUF SHA-256 mismatch: {source_sha256}") + + from mobius import build_from_gguf + from mobius.integrations.gguf import write_gguf_tokenizer_json + + output_dir.mkdir(parents=True, exist_ok=True) + baseline_memory = _memory_sample() + build_started = time.perf_counter() + package = build_from_gguf( + gguf_path, + keep_quantized=True, + execution_provider="cpu", + ) + build_seconds = time.perf_counter() - build_started + build_memory = _memory_sample() + mapping_audit, expected_initializers = _mapping_audit(gguf_path) + expected_mapping = { + "source_tensors": 417, + "mapped_sources": 401, + "mtp_exclusions": 16, + "logical_targets": 6243, + "q8_logical_targets": 6006, + "mapped_weight_initializers": 18_255, + } + if {key: mapping_audit[key] for key in expected_mapping} != expected_mapping: + raise ValueError(f"Unexpected mapping audit: {mapping_audit}") + graph_audit = _graph_audit(package["model"], expected_initializers) + + save_started = time.perf_counter() + package.save( + str(output_dir), + external_data="onnx", + progress_bar=False, + ) + save_seconds = time.perf_counter() - save_started + save_memory = _memory_sample() + tokenizer_path = write_gguf_tokenizer_json(gguf_path, output_dir) + if tokenizer_path is None: + raise ValueError("Pinned GGUF tokenizer.json was not emitted") + tokenizer_contract = _write_tokenizer_contract( + gguf_path, + official_tokenizer_dir, + output_dir, + ) + + # The GGUF metadata's PAD=999 is not a valid runtime padding contract for + # this model. Keep the official decoder contract and the tokenizer's role + # token separate and explicit. + _write_json( + output_dir / "config.json", + { + "model_type": "nemotron_h", + "bos_token_id": 1, + "eos_token_id": 2, + "pad_token_id": 0, + }, + ) + _write_json( + output_dir / "generation_config.json", + { + "bos_token_id": 1, + "eos_token_id": [2, 11], + "pad_token_id": 0, + }, + ) + + package_bytes = sum( + path.stat().st_size for path in output_dir.rglob("*") if path.is_file() + ) + del package + gc.collect() + released_memory = _memory_sample() + report = { + "phase": "build", + "source": { + "repo": GGUF_REPO, + "revision": GGUF_REVISION, + "filename": GGUF_FILENAME, + "size": GGUF_SIZE, + "sha256": source_sha256, + }, + "mapping": mapping_audit, + "tokenizer_contract": tokenizer_contract, + "graph": graph_audit, + "package_bytes": package_bytes, + "timing_seconds": { + "build": build_seconds, + "save": save_seconds, + }, + "memory_bytes": { + "baseline": baseline_memory, + "after_build": build_memory, + "after_save": save_memory, + "after_release": released_memory, + }, + } + _write_json(output_dir / "gguf_q8_build_report.json", report) + print(json.dumps(report, indent=2, sort_keys=True)) + + +def _prefill( + session, + output_names: list[str], + input_ids: list[int], + attention_mask: np.ndarray, +): + states = _initial_states(session) + feeds = _token_feeds( + session, + np.asarray([input_ids], dtype=np.int64), + total_length=len(input_ids), + position_ids=np.arange(len(input_ids), dtype=np.int64)[None, :], + states=states, + ) + feeds["attention_mask"] = attention_mask + outputs = _run_session(session, output_names, feeds) + _update_states(states, output_names, outputs) + logits = _as_numpy(outputs[output_names.index("logits")]).astype(np.float32) + return logits, states + + +def _run(output_dir: Path, device: str) -> None: + import onnxruntime as ort + from tokenizers import Tokenizer + + baseline_memory = _memory_sample() + load_started = time.perf_counter() + session = _create_session(output_dir / "model.onnx", device, profile=False) + load_seconds = time.perf_counter() - load_started + load_memory = _memory_sample() + output_names = [output.name for output in session.get_outputs()] + + prefill_started = time.perf_counter() + logits, states = _prefill( + session, + output_names, + PROMPT_IDS, + np.ones((1, len(PROMPT_IDS)), dtype=np.int64), + ) + prefill_seconds = time.perf_counter() - prefill_started + prefill_memory = _memory_sample() + + padded_ids = [*PROMPT_IDS, 0, 0] + padded_mask = np.asarray([[1] * len(PROMPT_IDS) + [0, 0]], dtype=np.int64) + padded_logits, _ = _prefill(session, output_names, padded_ids, padded_mask) + padded_real_token_max_abs = float( + np.max(np.abs(padded_logits[:, : len(PROMPT_IDS)] - logits)) + ) + np.testing.assert_allclose( + padded_logits[:, : len(PROMPT_IDS)], + logits, + rtol=1e-5, + atol=1e-5, + ) + + generated = [] + decode_seconds = [] + decode_memory = [] + past_length = len(PROMPT_IDS) + next_logits = logits[0, -1] + for index in range(len(EXPECTED_IDS)): + token_id = int(np.argmax(next_logits)) + generated.append(token_id) + if index + 1 == len(EXPECTED_IDS): + break + started = time.perf_counter() + feeds = _token_feeds( + session, + np.asarray([[token_id]], dtype=np.int64), + total_length=past_length + 1, + position_ids=np.asarray([[past_length]], dtype=np.int64), + states=states, + ) + outputs = _run_session(session, output_names, feeds) + _update_states(states, output_names, outputs) + past_length += 1 + next_logits = _as_numpy(outputs[output_names.index("logits")])[0, -1].astype( + np.float32 + ) + decode_seconds.append(time.perf_counter() - started) + decode_memory.append(_memory_sample()) + + if generated != EXPECTED_IDS: + raise AssertionError( + f"Greedy tokens differ from llama.cpp {LLAMA_CPP_COMMIT}: " + f"actual={generated}, expected={EXPECTED_IDS}" + ) + tokenizer = Tokenizer.from_file(str(output_dir / "tokenizer.json")) + prompt_ids = tokenizer.encode(PROMPT).ids + if prompt_ids != PROMPT_IDS: + raise AssertionError( + f"Tokenizer prompt IDs differ: actual={prompt_ids}, expected={PROMPT_IDS}" + ) + generated_text = tokenizer.decode(generated) + if generated_text != EXPECTED_TEXT: + raise AssertionError( + f"Decoded text differs: actual={generated_text!r}, expected={EXPECTED_TEXT!r}" + ) + if tokenizer.token_to_id("<|im_end|>") != 11: + raise AssertionError("Tokenizer role token <|im_end|> must remain ID 11") + tokenizer_config = json.loads( + (output_dir / "tokenizer_config.json").read_text(encoding="utf-8") + ) + asset_special_ids = { + role: tokenizer.token_to_id(tokenizer_config[f"{config_name}_token"]) + for role, config_name in ( + ("bos", "bos"), + ("eos", "eos"), + ("padding", "pad"), + ("unknown", "unk"), + ) + } + if asset_special_ids != {"bos": 1, "eos": 11, "padding": 11, "unknown": 0}: + raise AssertionError(f"Tokenizer asset contract differs: {asset_special_ids}") + + report = { + "phase": "runtime", + "onnxruntime_version": ort.__version__, + "onnxruntime_build_info": ort.get_build_info(), + "available_providers": ort.get_available_providers(), + "session_providers": session.get_providers(), + "device": device, + "prompt": PROMPT, + "prompt_ids": PROMPT_IDS, + "generated_ids": generated, + "generated_text": generated_text, + "llama_cpp_commit": LLAMA_CPP_COMMIT, + "llama_cpp_reference": LLAMA_CPP_REFERENCE, + "padding_contract": { + "gguf_padding_id_rejected": 999, + "runtime_padding_id": 0, + "runtime_eos_ids": [2, 11], + "tokenizer_asset_special_ids": asset_special_ids, + "right_padded_real_token_logits_max_abs": padded_real_token_max_abs, + "right_padded_real_token_logits_atol": 1e-5, + }, + "timing_seconds": { + "session_load": load_seconds, + "prefill": prefill_seconds, + "cached_decode_steps": decode_seconds, + }, + "throughput_tokens_per_second": { + "prefill": len(PROMPT_IDS) / prefill_seconds, + "cached_decode_steps": len(decode_seconds) / sum(decode_seconds), + }, + "memory_bytes": { + "baseline": baseline_memory, + "after_session_load": load_memory, + "after_prefill": prefill_memory, + "after_cached_decode_steps": decode_memory, + }, + } + _write_json(output_dir / "gguf_q8_runtime_report.json", report) + print(json.dumps(report, indent=2, sort_keys=True)) + + +def _all( + gguf_path: Path, + official_tokenizer_dir: Path, + output_dir: Path, + device: str, +) -> None: + common = [sys.executable, str(Path(__file__).resolve())] + subprocess.run( + [ + *common, + "--phase", + "build", + "--gguf", + str(gguf_path), + "--official-tokenizer-dir", + str(official_tokenizer_dir), + "--output", + str(output_dir), + ], + check=True, + ) + subprocess.run( + [ + *common, + "--phase", + "run", + "--output", + str(output_dir), + "--device", + device, + ], + check=True, + ) + combined = { + "build": json.loads( + (output_dir / "gguf_q8_build_report.json").read_text(encoding="utf-8") + ), + "runtime": json.loads( + (output_dir / "gguf_q8_runtime_report.json").read_text(encoding="utf-8") + ), + } + _write_json(output_dir / "gguf_q8_acceptance_report.json", combined) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--phase", choices=["all", "build", "run"], default="all") + parser.add_argument("--gguf", type=Path) + parser.add_argument("--official-tokenizer-dir", type=Path) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--device", choices=["cpu", "cuda"], default="cpu") + args = parser.parse_args() + + if args.phase in {"all", "build"} and args.gguf is None: + parser.error("--gguf is required for build/all phases") + if args.phase in {"all", "build"} and args.official_tokenizer_dir is None: + parser.error("--official-tokenizer-dir is required for build/all phases") + if args.phase == "build": + _build(args.gguf, args.official_tokenizer_dir, args.output) + elif args.phase == "run": + _run(args.output, args.device) + else: + _all(args.gguf, args.official_tokenizer_dir, args.output, args.device) + + +if __name__ == "__main__": + main() diff --git a/src/mobius/components/_mamba_block.py b/src/mobius/components/_mamba_block.py index a4a553494..c0a36d11b 100644 --- a/src/mobius/components/_mamba_block.py +++ b/src/mobius/components/_mamba_block.py @@ -294,8 +294,11 @@ def __init__( eps: float = 1e-5, norm_group_size: int | None = None, time_step_min: float = 0.0, + linear_class: type | None = None, ): super().__init__() + if linear_class is None: + linear_class = Linear self.d_model = d_model self.d_inner = d_inner self.num_heads = num_heads @@ -310,7 +313,7 @@ def __init__( self.conv_dim = d_inner + 2 * n_groups * d_state proj_size = d_inner + self.conv_dim + num_heads - self.in_proj = Linear(d_model, proj_size, bias=proj_bias) + self.in_proj = linear_class(d_model, proj_size, bias=proj_bias) self.conv1d = _Mamba2DepthwiseConv1d( self.conv_dim, conv_kernel, @@ -327,7 +330,7 @@ def __init__( eps=eps, group_size=norm_group_size, ) - self.out_proj = Linear(d_inner, d_model, bias=proj_bias) + self.out_proj = linear_class(d_inner, d_model, bias=proj_bias) def forward( self, diff --git a/src/mobius/integrations/gguf/_architecture.py b/src/mobius/integrations/gguf/_architecture.py new file mode 100644 index 000000000..a1410046e --- /dev/null +++ b/src/mobius/integrations/gguf/_architecture.py @@ -0,0 +1,189 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Architecture adapters for GGUF imports. + +Adapters keep source-architecture details out of the generic builder. They +translate GGUF metadata and tensor records into Mobius config and initializer +contracts while the builder owns parsing, graph construction, and weight +application. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from importlib import import_module +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + import torch + + from mobius._configs import ArchitectureConfig, QuantizationConfig + from mobius._model_package import ModelPackage + from mobius.integrations.gguf._reader import GGUFModel + from mobius.integrations.gguf._repacker import RepackedTensor + + +@dataclass(frozen=True) +class GGUFTensorTarget: + """One state-dict and ONNX target produced from a GGUF source tensor.""" + + state_dict_name: str + initializer_name: str + source_index: int | None = None + + +@dataclass(frozen=True) +class GGUFTensorMapping: + """Disposition of one GGUF tensor.""" + + targets: tuple[GGUFTensorTarget, ...] = () + exclusion: str | None = None + + def __post_init__(self) -> None: + if bool(self.targets) == bool(self.exclusion): + raise ValueError("A GGUF tensor mapping must have targets or one exclusion") + + @classmethod + def excluded(cls, reason: str) -> GGUFTensorMapping: + return cls(exclusion=reason) + + +@dataclass +class GGUFMappingAudit: + """Completeness accounting collected while loading a GGUF tensor table.""" + + mapped_sources: set[str] = field(default_factory=set) + excluded_sources: dict[str, str] = field(default_factory=dict) + unmapped_sources: set[str] = field(default_factory=set) + target_sources: dict[str, str] = field(default_factory=dict) + + def record(self, source_name: str, mapping: GGUFTensorMapping | None) -> None: + if mapping is None: + self.unmapped_sources.add(source_name) + return + if mapping.exclusion is not None: + self.excluded_sources[source_name] = mapping.exclusion + return + + self.mapped_sources.add(source_name) + for target in mapping.targets: + previous = self.target_sources.setdefault(target.initializer_name, source_name) + if previous != source_name: + raise ValueError( + f"GGUF sources {previous!r} and {source_name!r} both map to " + f"initializer {target.initializer_name!r}" + ) + + +class GGUFArchitectureAdapter: + """Base contract for source-architecture-specific GGUF behavior.""" + + architecture: str + model_type: str + + def __init__(self, model: GGUFModel) -> None: + self.model = model + + def validate_model(self, *, source: str) -> None: + """Validate the source tensor table and supported quantization.""" + raise NotImplementedError + + def build_config(self) -> ArchitectureConfig: + raise NotImplementedError + + def quantization_config(self) -> QuantizationConfig: + raise NotImplementedError + + def map_tensor( + self, + source_name: str, + source_shape: tuple[int, ...], + ) -> GGUFTensorMapping | None: + raise NotImplementedError + + def transform_tensor( + self, + source_name: str, + target: GGUFTensorTarget, + tensor: torch.Tensor, + ) -> torch.Tensor: + return tensor + + def transform_repacked( + self, + source_name: str, + target: GGUFTensorTarget, + tensor: RepackedTensor, + ) -> RepackedTensor: + return tensor + + def validate_mapping_audit(self, audit: GGUFMappingAudit) -> None: + if audit.unmapped_sources: + examples = sorted(audit.unmapped_sources)[:10] + raise ValueError( + f"{self.architecture} GGUF has {len(audit.unmapped_sources)} " + f"unmapped source tensor(s): {examples}" + ) + + +_ADAPTER_TYPES: dict[str, type[GGUFArchitectureAdapter]] = {} +_BUILTINS_LOADED = False +_BUILTIN_MODULES = ("mobius.integrations.gguf._nemotron_h_moe",) + + +def register_architecture_adapter( + adapter_type: type[GGUFArchitectureAdapter], +) -> type[GGUFArchitectureAdapter]: + """Register an architecture adapter class.""" + architecture = adapter_type.architecture + if architecture in _ADAPTER_TYPES: + raise ValueError(f"GGUF architecture adapter already registered: {architecture!r}") + _ADAPTER_TYPES[architecture] = adapter_type + return adapter_type + + +def _load_builtin_adapters() -> None: + global _BUILTINS_LOADED + if _BUILTINS_LOADED: + return + for module_name in _BUILTIN_MODULES: + import_module(module_name) + _BUILTINS_LOADED = True + + +def create_architecture_adapter( + architecture: str, + model: GGUFModel, +) -> GGUFArchitectureAdapter | None: + """Create the registered adapter for *architecture*, if one exists.""" + _load_builtin_adapters() + adapter_type = _ADAPTER_TYPES.get(architecture) + return adapter_type(model) if adapter_type is not None else None + + +def validate_package_state_dict( + package: ModelPackage, + state_dict: dict[str, torch.Tensor], +) -> None: + """Require exact state-dict coverage of all unset package initializers.""" + required: set[str] = set() + for model in package.values(): + required.update( + name + for name, initializer in model.graph.initializers.items() + if initializer.const_value is None + ) + + provided = set(state_dict) + missing = required - provided + unexpected = provided - required + if not missing and not unexpected: + return + + details = [] + if missing: + details.append(f"{len(missing)} missing: {sorted(missing)[:10]}") + if unexpected: + details.append(f"{len(unexpected)} unexpected: {sorted(unexpected)[:10]}") + raise ValueError("GGUF state-dict/initializer mismatch; " + "; ".join(details)) diff --git a/src/mobius/integrations/gguf/_architecture_test.py b/src/mobius/integrations/gguf/_architecture_test.py new file mode 100644 index 000000000..1eff962f9 --- /dev/null +++ b/src/mobius/integrations/gguf/_architecture_test.py @@ -0,0 +1,406 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +from __future__ import annotations + +import hashlib +import os +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pytest +import torch +from gguf import GGMLQuantizationType + +from mobius._configs import NemotronHConfig +from mobius.integrations.gguf._architecture import ( + GGUFArchitectureAdapter, + GGUFMappingAudit, + create_architecture_adapter, +) +from mobius.integrations.gguf._config_mapping import gguf_to_config + +_LAYER_TYPES = ( + "mamba2", + "moe", + "mamba2", + "moe", + "mamba2", + "full_attention", + "moe", + "mamba2", + "moe", + "mamba2", + "moe", + "mamba2", + "full_attention", + "moe", + "mamba2", + "moe", + "mamba2", + "moe", + "mamba2", + "full_attention", + "moe", + "mamba2", + "moe", + "mamba2", + "moe", + "mamba2", + "full_attention", + "moe", + "mamba2", + "moe", + "mamba2", + "moe", + "mamba2", + "full_attention", + "moe", + "mamba2", + "moe", + "mamba2", + "moe", + "mamba2", + "moe", + "mamba2", + "full_attention", + "moe", + "mamba2", + "moe", + "mamba2", + "moe", + "mamba2", + "moe", + "mamba2", + "moe", +) +_SUFFIXES = { + "mamba2": ( + "attn_norm.weight", + "ssm_a", + "ssm_conv1d.bias", + "ssm_conv1d.weight", + "ssm_d", + "ssm_dt.bias", + "ssm_in.weight", + "ssm_norm.weight", + "ssm_out.weight", + ), + "moe": ( + "attn_norm.weight", + "exp_probs_b.bias", + "ffn_down_exps.weight", + "ffn_down_shexp.weight", + "ffn_gate_inp.weight", + "ffn_up_exps.weight", + "ffn_up_shexp.weight", + ), + "full_attention": ( + "attn_k.weight", + "attn_norm.weight", + "attn_output.weight", + "attn_q.weight", + "attn_v.weight", + ), +} +_MTP_SUFFIXES = ( + "attn_k.weight", + "attn_norm.weight", + "attn_output.weight", + "attn_q.weight", + "attn_v.weight", + "exp_probs_b.bias", + "ffn_down_exps.weight", + "ffn_down_shexp.weight", + "ffn_gate_inp.weight", + "ffn_up_exps.weight", + "ffn_up_shexp.weight", + "nextn.eh_proj.weight", + "nextn.enorm.weight", + "nextn.hnorm.weight", + "nextn.shared_head_norm.weight", + "post_attention_norm.weight", +) +_Q8_SUFFIXES = { + "ssm_in.weight", + "ssm_out.weight", + "attn_q.weight", + "attn_k.weight", + "attn_v.weight", + "attn_output.weight", + "ffn_up_exps.weight", + "ffn_down_exps.weight", + "ffn_up_shexp.weight", + "ffn_down_shexp.weight", +} + + +def _source_shape(layer_type: str, suffix: str) -> tuple[int, ...]: + h, q, kv = 64, 64, 32 + experts, moe_inner, shared_inner = 128, 32, 64 + d_inner, conv_dim, mamba_heads = 64, 128, 2 + shapes = { + "attn_norm.weight": (h,), + "ssm_a": (mamba_heads, 1), + "ssm_conv1d.bias": (conv_dim,), + "ssm_conv1d.weight": (conv_dim, 4), + "ssm_d": (mamba_heads, 1), + "ssm_dt.bias": (mamba_heads,), + "ssm_in.weight": (d_inner + conv_dim + mamba_heads, h), + "ssm_norm.weight": (2, 32), + "ssm_out.weight": (h, d_inner), + "exp_probs_b.bias": (experts,), + "ffn_down_exps.weight": (experts, h, moe_inner), + "ffn_down_shexp.weight": (h, shared_inner), + "ffn_gate_inp.weight": (experts, h), + "ffn_up_exps.weight": (experts, moe_inner, h), + "ffn_up_shexp.weight": (shared_inner, h), + "attn_k.weight": (kv, h), + "attn_output.weight": (h, q), + "attn_q.weight": (q, h), + "attn_v.weight": (kv, h), + } + assert suffix in _SUFFIXES[layer_type] + return shapes[suffix] + + +def _record(name: str, qtype, shape: tuple[int, ...]): + return SimpleNamespace( + name=name, + tensor_type=qtype, + # ReaderTensor shapes use GGML order. + shape=np.asarray(tuple(reversed(shape)), dtype=np.uint64), + ) + + +def _synthetic_pinned_header(): + records = [ + _record("output.weight", GGMLQuantizationType.Q8_0, (256, 64)), + _record("output_norm.weight", GGMLQuantizationType.F32, (64,)), + _record("token_embd.weight", GGMLQuantizationType.Q8_0, (256, 64)), + ] + for index, layer_type in enumerate(_LAYER_TYPES): + for suffix in _SUFFIXES[layer_type]: + qtype = ( + GGMLQuantizationType.Q8_0 + if suffix in _Q8_SUFFIXES + else GGMLQuantizationType.F32 + ) + records.append( + _record( + f"blk.{index}.{suffix}", + qtype, + _source_shape(layer_type, suffix), + ) + ) + + mtp_shapes = { + suffix: _source_shape( + ( + "full_attention" + if suffix.startswith("attn_") and suffix != "attn_norm.weight" + else "moe" + ), + suffix, + ) + for suffix in _MTP_SUFFIXES + if suffix in _Q8_SUFFIXES + or suffix in {"attn_norm.weight", "exp_probs_b.bias", "ffn_gate_inp.weight"} + } + mtp_shapes.update( + { + "nextn.eh_proj.weight": (64, 128), + "nextn.enorm.weight": (64,), + "nextn.hnorm.weight": (64,), + "nextn.shared_head_norm.weight": (64,), + "post_attention_norm.weight": (64,), + } + ) + for suffix in _MTP_SUFFIXES: + if suffix == "ffn_gate_inp.weight": + qtype = GGMLQuantizationType.BF16 + elif suffix in _Q8_SUFFIXES or suffix == "nextn.eh_proj.weight": + qtype = GGMLQuantizationType.Q8_0 + else: + qtype = GGMLQuantizationType.F32 + records.append(_record(f"blk.52.{suffix}", qtype, mtp_shapes[suffix])) + + metadata = { + "nemotron_h_moe.attention.head_count": 2, + "nemotron_h_moe.attention.head_count_kv": [ + 1 if layer_type == "full_attention" else 0 for layer_type in _LAYER_TYPES + ] + + [1], + "nemotron_h_moe.attention.key_length": 32, + "nemotron_h_moe.attention.layer_norm_rms_epsilon": 1e-5, + "nemotron_h_moe.block_count": 53, + "nemotron_h_moe.context_length": 128, + "nemotron_h_moe.embedding_length": 64, + "nemotron_h_moe.expert_count": 128, + "nemotron_h_moe.expert_feed_forward_length": 32, + "nemotron_h_moe.expert_group_count": 1, + "nemotron_h_moe.expert_group_used_count": 1, + "nemotron_h_moe.expert_shared_count": 1, + "nemotron_h_moe.expert_shared_feed_forward_length": 64, + "nemotron_h_moe.expert_used_count": 1, + "nemotron_h_moe.expert_weights_norm": True, + "nemotron_h_moe.expert_weights_scale": 2.5, + "nemotron_h_moe.nextn_predict_layers": 1, + "nemotron_h_moe.ssm.conv_kernel": 4, + "nemotron_h_moe.ssm.group_count": 2, + "nemotron_h_moe.ssm.inner_size": 64, + "nemotron_h_moe.ssm.state_size": 16, + "nemotron_h_moe.ssm.time_step_rank": 2, + "nemotron_h_moe.vocab_size": 256, + "tokenizer.ggml.model": "gpt2", + "tokenizer.ggml.pre": "pixtral", + "tokenizer.ggml.bos_token_id": 1, + "tokenizer.ggml.eos_token_id": 11, + "tokenizer.ggml.padding_token_id": 999, + } + return SimpleNamespace( + architecture="nemotron_h_moe", + metadata=metadata, + tensor_names=[record.name for record in records], + _reader=SimpleNamespace(tensors=records), + _path=Path("synthetic-nemotron-q8.gguf"), + ) + + +def test_architecture_adapter_requires_source_validation() -> None: + adapter = GGUFArchitectureAdapter(SimpleNamespace()) + + with pytest.raises(NotImplementedError): + adapter.validate_model(source="synthetic") + + +def test_nemotron_adapter_derives_exact_backbone_and_mapping() -> None: + model = _synthetic_pinned_header() + adapter = create_architecture_adapter(model.architecture, model) + assert adapter is not None + adapter.validate_model(source="synthetic") + + config = gguf_to_config(model, adapter=adapter) + assert isinstance(config, NemotronHConfig) + assert config.model_type == "nemotron_h" + assert config.layer_types == list(_LAYER_TYPES) + assert config.layer_types.count("mamba2") == 23 + assert config.layer_types.count("moe") == 23 + assert config.layer_types.count("full_attention") == 6 + assert config.num_hidden_layers == 52 + assert config.bos_token_id == 1 + assert config.eos_token_id == 2 + assert config.pad_token_id == 0 + assert config.num_key_value_heads == 1 + assert config.mamba_n_heads == 2 + assert config.mamba_d_head == 32 + + audit = GGUFMappingAudit() + for record in model._reader.tensors: + shape = tuple(int(dim) for dim in reversed(record.shape)) + audit.record(record.name, adapter.map_tensor(record.name, shape)) + adapter.validate_mapping_audit(audit) + + assert len(audit.mapped_sources) == 401 + assert len(audit.excluded_sources) == 16 + assert set(audit.excluded_sources) == {f"blk.52.{suffix}" for suffix in _MTP_SUFFIXES} + assert len(audit.target_sources) == 6243 + + +def test_nemotron_adapter_expands_experts_and_transforms_mamba_values() -> None: + model = _synthetic_pinned_header() + adapter = create_architecture_adapter(model.architecture, model) + assert adapter is not None + + mapping = adapter.map_tensor("blk.1.ffn_up_exps.weight", (128, 32, 64)) + assert mapping is not None + assert [target.source_index for target in mapping.targets] == list(range(128)) + assert mapping.targets[0].state_dict_name == ( + "backbone.layers.1.mixer.experts.0.up_proj.weight" + ) + assert mapping.targets[1].initializer_name == ( + "model.layers.1.moe.experts.1.up_proj.weight" + ) + + target = adapter.map_tensor("blk.0.ssm_a", (2, 1)).targets[0] + transformed = adapter.transform_tensor( + "blk.0.ssm_a", + target, + torch.tensor([[-1.0], [-np.e]], dtype=torch.float32), + ) + torch.testing.assert_close(transformed, torch.tensor([0.0, 1.0])) + + conv_target = adapter.map_tensor("blk.0.ssm_conv1d.weight", (128, 4)).targets[0] + conv = adapter.transform_tensor( + "blk.0.ssm_conv1d.weight", + conv_target, + torch.zeros(128, 4), + ) + assert conv.shape == (128, 1, 4) + + with pytest.raises(ValueError, match="expected"): + adapter.map_tensor("blk.0.ssm_out.weight", (63, 64)) + + +def test_nemotron_adapter_rejects_unpreserved_base_qtype_with_evidence() -> None: + model = _synthetic_pinned_header() + record = next( + record for record in model._reader.tensors if record.name == "blk.1.ffn_up_exps.weight" + ) + record.tensor_type = GGMLQuantizationType.Q5_1 + adapter = create_architecture_adapter(model.architecture, model) + assert adapter is not None + + with pytest.raises(NotImplementedError, match=r"Q5_1=262,144 parameters"): + adapter.validate_model(source="synthetic") + + +def test_nemotron_adapter_rejects_qtype_location_swap() -> None: + model = _synthetic_pinned_header() + q8_record = next( + record for record in model._reader.tensors if record.name == "blk.1.ffn_up_exps.weight" + ) + float_record = next( + record + for record in model._reader.tensors + if record.name == "blk.1.ffn_gate_inp.weight" + ) + q8_record.tensor_type, float_record.tensor_type = ( + float_record.tensor_type, + q8_record.tensor_type, + ) + adapter = create_architecture_adapter(model.architecture, model) + assert adapter is not None + + with pytest.raises(ValueError, match="exact Q8 preservation"): + adapter.validate_model(source="synthetic") + + +@pytest.mark.integration +def test_pinned_nemotron_q8_header_and_mapping_from_real_artifact() -> None: + path_value = os.environ.get("MOBIUS_NEMOTRON_Q8_GGUF") + if not path_value: + pytest.skip("Set MOBIUS_NEMOTRON_Q8_GGUF to the pinned Q8_0 artifact") + path = Path(path_value) + assert path.stat().st_size == 35_004_643_392 + + digest = hashlib.sha256() + with path.open("rb") as stream: + while chunk := stream.read(8 * 1024 * 1024): + digest.update(chunk) + assert digest.hexdigest() == ( + "dc5276dd0619c04e277504d2358a793e31ccbe39e894d767d0d14f2a221e2ca4" + ) + + from mobius.integrations.gguf._reader import GGUFModel + + model = GGUFModel(path) + adapter = create_architecture_adapter(model.architecture, model) + assert adapter is not None + adapter.validate_model(source=str(path)) + audit = GGUFMappingAudit() + for record in model._reader.tensors: + source_shape = tuple(int(dim) for dim in reversed(record.shape)) + audit.record(record.name, adapter.map_tensor(record.name, source_shape)) + adapter.validate_mapping_audit(audit) diff --git a/src/mobius/integrations/gguf/_builder.py b/src/mobius/integrations/gguf/_builder.py index 0f3e9de3c..f9d8b61ed 100644 --- a/src/mobius/integrations/gguf/_builder.py +++ b/src/mobius/integrations/gguf/_builder.py @@ -21,7 +21,7 @@ import logging import re from collections import Counter -from collections.abc import Iterable, Mapping +from collections.abc import Mapping from pathlib import Path from typing import TYPE_CHECKING @@ -39,6 +39,8 @@ _HUB_PREFLIGHT_TRANSPORT_ERRORS += (_HttpxTransportError,) if TYPE_CHECKING: + from mobius.integrations.gguf._architecture import GGUFArchitectureAdapter + from mobius.integrations.gguf._reader import GGUFModel from mobius.tasks import ModelTask logger = logging.getLogger(__name__) @@ -47,76 +49,6 @@ r"-(?P\d{5})-of-(?P\d{5})\.gguf$", re.IGNORECASE, ) -_NEMOTRON_H_MOE_ARCHITECTURE = "nemotron_h_moe" - - -def _summarize_nemotron_h_moe_layout( - tensor_names: Iterable[str], -) -> tuple[Counter[str], tuple[int, ...], dict[int, frozenset[str]]]: - """Summarize base-layer and MTP mixer types from Nemotron-H GGUF names.""" - layer_kinds: dict[int, set[str]] = {} - mtp_blocks: set[int] = set() - for name in tensor_names: - match = re.match(r"^blk\.(\d+)\.(.+)$", name) - if match is None: - continue - block_index = int(match.group(1)) - suffix = match.group(2) - kinds = layer_kinds.setdefault(block_index, set()) - if suffix.startswith("nextn."): - mtp_blocks.add(block_index) - elif suffix.startswith("ssm_"): - kinds.add("mamba") - elif suffix.startswith(("ffn_", "exp_probs_")): - kinds.add("moe") - elif suffix.startswith(("attn_q.", "attn_k.", "attn_v.", "attn_output.")): - kinds.add("attention") - - base_counts: Counter[str] = Counter() - for block_index, kinds in layer_kinds.items(): - if block_index not in mtp_blocks: - base_counts.update(kinds) - mtp_kinds = { - block_index: frozenset(layer_kinds.get(block_index, set())) - for block_index in sorted(mtp_blocks) - } - return base_counts, tuple(sorted(mtp_blocks)), mtp_kinds - - -def _raise_for_unsupported_gguf_architecture( - architecture: str, - *, - source: str, - tensor_names: Iterable[str] | None = None, -) -> None: - """Reject GGUF architectures that do not have semantic conversion evidence.""" - if architecture != _NEMOTRON_H_MOE_ARCHITECTURE: - return - - layout = "" - if tensor_names is not None: - counts, mtp_blocks, mtp_kinds = _summarize_nemotron_h_moe_layout(tensor_names) - mtp_kind_names = {index: sorted(kinds) for index, kinds in mtp_kinds.items()} - layout = ( - " Detected base schedule: " - f"{counts['mamba']} Mamba + {counts['moe']} MoE + " - f"{counts['attention']} attention layers; auxiliary MTP blocks: " - f"{list(mtp_blocks)} with mixer types {mtp_kind_names}." - ) - - raise NotImplementedError( - "Direct GGUF conversion for architecture 'nemotron_h_moe' is intentionally " - f"disabled for {source!r}.{layout} 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 current 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. No ONNX artifacts were emitted. 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." - ) def _raise_for_sharded_gguf( @@ -176,21 +108,23 @@ def _preflight_hf_gguf(api: HfApi, repo_id: str, filename: str) -> None: else: architecture = getattr(gguf_metadata, "architecture", None) if isinstance(architecture, str): - _raise_for_unsupported_gguf_architecture( - architecture, - source=source, - ) + logger.debug("Hub GGUF architecture preflight for %s: %s", source, architecture) -def _validate_gguf_model(gguf_model, *, source: str) -> None: +def _validate_gguf_model( + gguf_model: GGUFModel, + *, + source: str, +) -> GGUFArchitectureAdapter | None: """Validate a parsed GGUF before config extraction or graph construction.""" + from mobius.integrations.gguf._architecture import create_architecture_adapter + split_count = int(gguf_model.get_metadata("split.count", 1)) _raise_for_sharded_gguf(source=source, split_count=split_count) - _raise_for_unsupported_gguf_architecture( - gguf_model.architecture, - source=source, - tensor_names=gguf_model.tensor_names, - ) + adapter = create_architecture_adapter(gguf_model.architecture, gguf_model) + if adapter is not None: + adapter.validate_model(source=source) + return adapter def _looks_like_hf_repo_id(value: str) -> bool: @@ -363,15 +297,17 @@ def build_from_gguf( # 1. Parse GGUF file (auto-download from HF Hub when given "owner/repo[:filename]") gguf_path = _resolve_gguf_path(gguf_path) gguf_model = GGUFModel(gguf_path) - _validate_gguf_model(gguf_model, source=str(gguf_path)) + adapter = _validate_gguf_model(gguf_model, source=str(gguf_path)) gguf_arch = gguf_model.architecture logger.info("Loaded GGUF file: %s (arch=%s)", gguf_path, gguf_arch) - preserve_quantization = keep_quantized and _has_quantized_weights(gguf_model, gguf_arch) + preserve_quantization = keep_quantized and _has_quantized_weights( + gguf_model, gguf_arch, adapter=adapter + ) if keep_quantized and not preserve_quantization: logger.info("GGUF contains no mapped quantized weights; using the float import path") # 2. Extract config from GGUF metadata - config = gguf_to_config(gguf_model) + config = gguf_to_config(gguf_model, adapter=adapter) model_type = getattr(config, "_gguf_model_type", None) if model_type is None: model_type = GGUF_ARCH_TO_MODEL_TYPE.get(gguf_arch, gguf_arch) @@ -383,27 +319,30 @@ def build_from_gguf( # 3. Quantized path: detect dominant type and set config if preserve_quantization: - from mobius._configs import QuantizationConfig - from mobius._flags import flags - from mobius.integrations.gguf._tencent_q1_0 import is_tencent_q1_0_layout - - bits, block_size, is_sym = _detect_quant_params(gguf_model, gguf_arch) - # Float zero-point only when actually using Tencent's native 2-bit form. - float_zp = is_tencent_q1_0_layout(gguf_model) and flags.tencent_q1_0_use_native_2bit - quantize_embeddings = _can_quantize_embedding( - gguf_model, - gguf_arch, - bits=bits, - block_size=block_size, - ) - quantize_lm_head = ( - quantize_embeddings - if config.tie_word_embeddings - else _can_quantize_lm_head(gguf_model, gguf_arch) - ) - config = dataclasses.replace( - config, - quantization=QuantizationConfig( + if adapter is not None: + quantization = adapter.quantization_config() + else: + from mobius._configs import QuantizationConfig + from mobius._flags import flags + from mobius.integrations.gguf._tencent_q1_0 import is_tencent_q1_0_layout + + bits, block_size, is_sym = _detect_quant_params(gguf_model, gguf_arch) + # Float zero-point only when actually using Tencent's native 2-bit form. + float_zp = ( + is_tencent_q1_0_layout(gguf_model) and flags.tencent_q1_0_use_native_2bit + ) + quantize_embeddings = _can_quantize_embedding( + gguf_model, + gguf_arch, + bits=bits, + block_size=block_size, + ) + quantize_lm_head = ( + quantize_embeddings + if config.tie_word_embeddings + else _can_quantize_lm_head(gguf_model, gguf_arch) + ) + quantization = QuantizationConfig( bits=bits, group_size=block_size, quant_method="gguf", @@ -412,17 +351,20 @@ def build_from_gguf( quantize_embeddings=quantize_embeddings, quantize_lm_head=quantize_lm_head, tie_word_embeddings=quantize_lm_head and config.tie_word_embeddings, - ), + ) + config = dataclasses.replace( + config, + quantization=quantization, ) logger.info( "Quantized mode: bits=%d, block_size=%d, symmetric=%s, " "float_zp=%s, embedding=%s, lm_head=%s", - bits, - block_size, - is_sym, - float_zp, - quantize_embeddings, - quantize_lm_head, + quantization.bits, + quantization.group_size, + quantization.sym, + quantization.float_zero_point, + quantization.quantize_embeddings, + quantization.quantize_lm_head, ) # 4. Look up module class and resolve task @@ -439,7 +381,7 @@ def build_from_gguf( # 5. Build ONNX graph module = module_class(config) - if preserve_quantization: + if preserve_quantization and adapter is None: _replace_native_block_linears(module, gguf_model, gguf_arch) pkg = build_from_module( module, config, resolved_task, execution_provider=execution_provider @@ -452,9 +394,19 @@ def build_from_gguf( # 6. Load tensors from GGUF → state_dict if preserve_quantization: - state_dict = _load_quantized_state_dict(gguf_model, gguf_arch, module, config) + state_dict = _load_quantized_state_dict( + gguf_model, + gguf_arch, + module, + config, + adapter=adapter, + ) else: - state_dict = _load_dequantized_state_dict(gguf_model, gguf_arch) + state_dict = _load_dequantized_state_dict( + gguf_model, + gguf_arch, + adapter=adapter, + ) logger.info( "Mapped %d state_dict entries from GGUF tensors", @@ -465,28 +417,29 @@ def build_from_gguf( # For the quantized path, only float tensors go through # process_tensors; quantized Q/K tensors were permuted in # _load_quantized_state_dict already. - if preserve_quantization: - float_keys = { - k - for k in state_dict - if not ( - k.endswith((".scales", ".zero_points")) - or _is_quantized_weight(k, state_dict) - or _is_native_block_weight(k, state_dict) - ) - } - float_dict = {k: state_dict[k] for k in float_keys} - quant_dict = {k: state_dict[k] for k in state_dict if k not in float_keys} - float_dict = process_tensors(float_dict, config) - state_dict = {**float_dict, **quant_dict} - else: - state_dict = process_tensors(state_dict, config) + if adapter is None: + if preserve_quantization: + float_keys = { + k + for k in state_dict + if not ( + k.endswith((".scales", ".zero_points")) + or _is_quantized_weight(k, state_dict) + or _is_native_block_weight(k, state_dict) + ) + } + float_dict = {k: state_dict[k] for k in float_keys} + quant_dict = {k: state_dict[k] for k in state_dict if k not in float_keys} + float_dict = process_tensors(float_dict, config) + state_dict = {**float_dict, **quant_dict} + else: + state_dict = process_tensors(state_dict, config) - # 7b. Normalize GGUF-specific weight shapes to match HF conventions. - # This converts GGUF tensor quirks (stacked experts, 1D gates, 2D - # conv weights, suffix artifacts) into the shapes that HF models - # produce, so preprocess_weights only needs to handle HF→ONNX. - state_dict = _normalize_gguf_weights(state_dict) + # Normalize GGUF-specific weight shapes to match HF conventions. + # This converts GGUF tensor quirks (stacked experts, 1D gates, 2D + # conv weights, suffix artifacts) into the shapes that HF models + # produce, so preprocess_weights only needs to handle HF→ONNX. + state_dict = _normalize_gguf_weights(state_dict) # 8. Run model-specific preprocess_weights (HF → ONNX names) if hasattr(module, "preprocess_weights"): @@ -494,6 +447,10 @@ def build_from_gguf( # 9. Apply weights to ONNX model prefix_map = getattr(module, "weight_prefix_map", None) + if adapter is not None: + from mobius.integrations.gguf._architecture import validate_package_state_dict + + validate_package_state_dict(pkg, state_dict) pkg.apply_weights(state_dict, prefix_map=prefix_map) return pkg @@ -681,7 +638,7 @@ def _normalize_gguf_weights( return result -def _has_quantized_weights(gguf_model, gguf_arch: str) -> bool: +def _has_quantized_weights(gguf_model, gguf_arch: str, *, adapter=None) -> bool: """Return whether a GGUF has mapped weights with a quantized tensor type.""" from gguf import GGMLQuantizationType @@ -696,9 +653,16 @@ def _has_quantized_weights(gguf_model, gguf_arch: str) -> bool: if f64_type is not None: float_types.add(f64_type) - for name, _raw, qtype, _shape in gguf_model.tensor_items_raw(): - hf_name = map_gguf_to_hf_names(name, gguf_arch) - if hf_name is not None and hf_name.endswith(".weight") and qtype not in float_types: + for name, _raw, qtype, shape in gguf_model.tensor_items_raw(): + if adapter is not None: + mapping = adapter.map_tensor(name, shape) + mapped_weight = mapping is not None and any( + target.initializer_name.endswith(".weight") for target in mapping.targets + ) + else: + hf_name = map_gguf_to_hf_names(name, gguf_arch) + mapped_weight = hf_name is not None and hf_name.endswith(".weight") + if mapped_weight and qtype not in float_types: return True return False @@ -984,11 +948,200 @@ def repack_gguf_weight_to_target( ) +def _load_adapter_dequantized_state_dict( + gguf_model, + adapter, +) -> dict: + """Load an adapter-mapped state dict through the float path.""" + import numpy as np + import torch + + from mobius.integrations.gguf._architecture import GGUFMappingAudit + + audit = GGUFMappingAudit() + state_dict: dict[str, torch.Tensor] = {} + for source_name, raw, qtype, source_shape in tqdm.tqdm( + gguf_model.tensor_items_raw(), + desc="Dequantizing tensors", + total=len(gguf_model._tensor_index), + ): + mapping = adapter.map_tensor(source_name, source_shape) + audit.record(source_name, mapping) + if mapping is None or mapping.exclusion is not None: + continue + + array = gguf_model.dequantize_raw_tensor(raw, qtype, source_shape) + if not array.flags.writeable: + array = np.array(array) + source_tensor = torch.from_numpy(array) + for target in mapping.targets: + tensor = ( + source_tensor + if target.source_index is None + else source_tensor[target.source_index] + ) + state_dict[target.state_dict_name] = adapter.transform_tensor( + source_name, + target, + tensor, + ) + + adapter.validate_mapping_audit(audit) + return state_dict + + +def _load_adapter_quantized_state_dict( + gguf_model, + adapter, + module, + config, +) -> dict: + """Load adapter targets with exact affine-block preservation.""" + import numpy as np + import torch + + from mobius.components import QuantizedEmbedding, QuantizedLinear + from mobius.integrations.gguf._architecture import GGUFMappingAudit + from mobius.integrations.gguf._repacker import ( + repack_gguf_tensor, + repack_quant_params, + repack_stacked_gguf_tensor, + ) + + quantized_stems = { + name + for name, child in module.named_modules() + if isinstance(child, QuantizedLinear) + or getattr(child, "_gguf_quantized_linear", False) + } + embedding_stems = { + name for name, child in module.named_modules() if isinstance(child, QuantizedEmbedding) + } + + audit = GGUFMappingAudit() + state_dict: dict[str, torch.Tensor] = {} + repacked_targets = 0 + target_params = (config.quantization.bits, config.quantization.group_size) + + for source_name, raw, qtype, source_shape in tqdm.tqdm( + gguf_model.tensor_items_raw(), + desc="Repacking tensors", + total=len(gguf_model._tensor_index), + ): + mapping = adapter.map_tensor(source_name, source_shape) + audit.record(source_name, mapping) + if mapping is None or mapping.exclusion is not None: + continue + + target_stems = [ + target.initializer_name.removesuffix(".weight") for target in mapping.targets + ] + target_is_quantized = [ + stem in quantized_stems or stem in embedding_stems for stem in target_stems + ] + if any(target_is_quantized) and not all(target_is_quantized): + raise ValueError(f"GGUF source {source_name!r} mixes quantized and float targets") + + if all(target_is_quantized): + if not all( + target.initializer_name.endswith(".weight") for target in mapping.targets + ): + raise ValueError( + f"Quantized GGUF source {source_name!r} has a non-weight target" + ) + qtype_value = qtype.value if hasattr(qtype, "value") else qtype + if repack_quant_params(qtype_value) != target_params: + raise ValueError( + f"GGUF source {source_name!r} has quantization " + f"{getattr(qtype, 'name', qtype)!r}; expected an exact " + f"{target_params[0]}-bit/block-{target_params[1]} repack" + ) + + packed = raw.ravel().view(np.uint8) + if len(source_shape) == 2 and len(mapping.targets) == 1: + repacked = ( + repack_gguf_tensor( + packed, + qtype_value, + source_shape, + ), + ) + elif len(source_shape) == 3: + repacked = repack_stacked_gguf_tensor( + packed, + qtype_value, + source_shape, + ) + source_indices = {target.source_index for target in mapping.targets} + if source_indices != set(range(source_shape[0])): + raise ValueError( + f"Stacked GGUF source {source_name!r} does not map every " + "leading-axis slice exactly once" + ) + else: + raise ValueError( + f"Cannot exactly repack GGUF source {source_name!r} with " + f"shape {source_shape} into {len(mapping.targets)} targets" + ) + + for target in mapping.targets: + index = target.source_index or 0 + target_tensor = adapter.transform_repacked( + source_name, + target, + repacked[index], + ) + state_stem = target.state_dict_name.removesuffix(".weight") + initializer_stem = target.initializer_name.removesuffix(".weight") + if initializer_stem in embedding_stems: + state_dict[f"{state_stem}.qweight"] = torch.from_numpy( + target_tensor.weight.reshape(target_tensor.weight.shape[0], -1) + ) + else: + state_dict[target.state_dict_name] = torch.from_numpy(target_tensor.weight) + state_dict[f"{state_stem}.scales"] = torch.from_numpy(target_tensor.scales) + if target_tensor.zero_points is not None: + state_dict[f"{state_stem}.zero_points"] = torch.from_numpy( + target_tensor.zero_points + ) + repacked_targets += 1 + continue + + array = gguf_model.dequantize_raw_tensor(raw, qtype, source_shape) + if not array.flags.writeable: + array = np.array(array) + source_tensor = torch.from_numpy(array) + for target in mapping.targets: + tensor = ( + source_tensor + if target.source_index is None + else source_tensor[target.source_index] + ) + state_dict[target.state_dict_name] = adapter.transform_tensor( + source_name, + target, + tensor, + ) + + adapter.validate_mapping_audit(audit) + logger.info( + "Loaded %d adapter state-dict entries with %d exact quantized targets", + len(state_dict), + repacked_targets, + ) + return state_dict + + def _load_dequantized_state_dict( gguf_model, gguf_arch: str, + *, + adapter=None, ) -> dict: """Load all tensors dequantized to float (Phase 1 path).""" + if adapter is not None: + return _load_adapter_dequantized_state_dict(gguf_model, adapter) + import numpy as np import torch @@ -1019,6 +1172,8 @@ def _load_quantized_state_dict( gguf_arch: str, module, config, + *, + adapter=None, ) -> dict: """Load tensors, preserving native blocks or normalizing to MatMulNBits. @@ -1033,6 +1188,14 @@ def _load_quantized_state_dict( row-level reverse-permutation that ``process_tensors`` would normally apply. """ + if adapter is not None: + return _load_adapter_quantized_state_dict( + gguf_model, + adapter, + module, + config, + ) + import numpy as np import torch from gguf import GGMLQuantizationType, dequantize diff --git a/src/mobius/integrations/gguf/_builder_test.py b/src/mobius/integrations/gguf/_builder_test.py index 643bdb45f..f5d522845 100644 --- a/src/mobius/integrations/gguf/_builder_test.py +++ b/src/mobius/integrations/gguf/_builder_test.py @@ -877,126 +877,25 @@ def test_dequantize_raw_tensor_matches_get_tensor(self, q4_0_gguf: Path): class TestGGUFPreflightGuards: """Unsupported layouts fail before graph construction or large downloads.""" - def test_nemotron_layout_excludes_combined_mtp_block(self): - from mobius.integrations.gguf._builder import ( - _summarize_nemotron_h_moe_layout, - ) - - # Pinned NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16 schedule: - # 52 backbone layers followed by one combined attention+MoE MTP block. - backbone_schedule = ( - "mamba", - "moe", - "mamba", - "moe", - "mamba", - "attention", - "moe", - "mamba", - "moe", - "mamba", - "moe", - "mamba", - "attention", - "moe", - "mamba", - "moe", - "mamba", - "moe", - "mamba", - "attention", - "moe", - "mamba", - "moe", - "mamba", - "moe", - "mamba", - "attention", - "moe", - "mamba", - "moe", - "mamba", - "moe", - "mamba", - "attention", - "moe", - "mamba", - "moe", - "mamba", - "moe", - "mamba", - "moe", - "mamba", - "attention", - "moe", - "mamba", - "moe", - "mamba", - "moe", - "mamba", - "moe", - "mamba", - "moe", - ) - assert len(backbone_schedule) == 52 - - representative_tensor = { - "mamba": "ssm_in.weight", - "moe": "ffn_up_exps.weight", - "attention": "attn_q.weight", - } - tensor_names = [ - f"blk.{index}.{representative_tensor[layer_type]}" - for index, layer_type in enumerate(backbone_schedule) - ] - tensor_names.extend( - [ - "blk.52.nextn.eh_proj.weight", - "blk.52.attn_q.weight", - "blk.52.ffn_up_exps.weight", - ] - ) - - counts, mtp_blocks, mtp_kinds = _summarize_nemotron_h_moe_layout(tensor_names) - - assert dict(counts) == {"mamba": 23, "moe": 23, "attention": 6} - 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 - - def test_remote_nemotron_h_moe_fails_before_download(self): + def test_remote_supported_adapter_proceeds_to_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.HfApi") as api_type, mock.patch("mobius.integrations.gguf._builder.hf_hub_download") as download, - pytest.raises(NotImplementedError, match="nemotron_h_moe"), ): api_type.return_value.model_info.return_value = SimpleNamespace( gguf={"architecture": "nemotron_h_moe"} ) - _resolve_gguf_path(f"unsloth/nemotron:{filename}") + download.return_value = "cached.gguf" + result = _resolve_gguf_path(f"unsloth/nemotron:{filename}") + assert result == "cached.gguf" api_type.return_value.model_info.assert_called_once_with( "unsloth/nemotron", expand=["gguf"] ) - download.assert_not_called() + download.assert_called_once_with(repo_id="unsloth/nemotron", filename=filename) @pytest.mark.parametrize( "preflight_error", diff --git a/src/mobius/integrations/gguf/_config_mapping.py b/src/mobius/integrations/gguf/_config_mapping.py index 77a508306..a0568dbfa 100644 --- a/src/mobius/integrations/gguf/_config_mapping.py +++ b/src/mobius/integrations/gguf/_config_mapping.py @@ -25,7 +25,7 @@ import contextlib import dataclasses import logging -from typing import Any +from typing import TYPE_CHECKING, Any import numpy as np @@ -37,6 +37,9 @@ _shallow_fields, ) +if TYPE_CHECKING: + from mobius.integrations.gguf._architecture import GGUFArchitectureAdapter + logger = logging.getLogger(__name__) @@ -266,6 +269,8 @@ def _valid_token_id(value: Any) -> int | None: def gguf_to_config( model: Any, # GGUFModel — typed as Any to avoid circular import + *, + adapter: GGUFArchitectureAdapter | None = None, ) -> ArchitectureConfig: """Convert GGUF metadata to an :class:`ArchitectureConfig`. @@ -287,6 +292,26 @@ def gguf_to_config( gguf_arch = model.architecture metadata = model.metadata + if adapter is None: + from mobius.integrations.gguf._architecture import create_architecture_adapter + + adapter = create_architecture_adapter(gguf_arch, model) + if adapter is not None: + adapter.validate_model(source=str(getattr(model, "_path", ""))) + if adapter is not None: + config = adapter.build_config() + logger.info( + "Extracted config through GGUF architecture adapter: " + "arch=%s, model_type=%s, hidden=%d, layers=%d, heads=%d, vocab=%d", + gguf_arch, + adapter.model_type, + config.hidden_size, + config.num_hidden_layers, + config.num_attention_heads, + config.vocab_size, + ) + return config + # Resolve model_type model_type = GGUF_ARCH_TO_MODEL_TYPE.get(gguf_arch, gguf_arch) diff --git a/src/mobius/integrations/gguf/_mmproj_test.py b/src/mobius/integrations/gguf/_mmproj_test.py index 7afed3309..7170b4fbc 100644 --- a/src/mobius/integrations/gguf/_mmproj_test.py +++ b/src/mobius/integrations/gguf/_mmproj_test.py @@ -214,7 +214,7 @@ def clip_mmproj_gguf(tmp_path: Path) -> Path: class TestMultimodalPreflightGuards: - def test_rejects_unsupported_text_architecture_before_config( + def test_validates_text_adapter_before_resolving_mmproj( self, tmp_path: Path, ): @@ -228,7 +228,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="missing metadata 'block_count'"), ): build_gemma4_vlm_from_gguf(text_path, "owner/repo:mmproj.gguf") diff --git a/src/mobius/integrations/gguf/_nemotron_h_moe.py b/src/mobius/integrations/gguf/_nemotron_h_moe.py new file mode 100644 index 000000000..d341efc9f --- /dev/null +++ b/src/mobius/integrations/gguf/_nemotron_h_moe.py @@ -0,0 +1,697 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""GGUF adapter for NVIDIA Nemotron 3.5 Lightning.""" + +from __future__ import annotations + +import math +import re +from collections import Counter +from typing import TYPE_CHECKING, Any + +import numpy as np +import onnx_ir as ir +import torch + +from mobius._configs import NemotronHConfig, QuantizationConfig +from mobius.integrations.gguf._architecture import ( + GGUFArchitectureAdapter, + GGUFMappingAudit, + GGUFTensorMapping, + GGUFTensorTarget, + register_architecture_adapter, +) +from mobius.integrations.gguf._repacker import RepackedTensor +from mobius.integrations.gguf._tensor_processors import _reverse_permute + +if TYPE_CHECKING: + from mobius.integrations.gguf._reader import GGUFModel + +_BLOCK_RE = re.compile(r"^blk\.(\d+)\.(.+)$") + +_PINNED_LAYER_TYPES = ( + "mamba2", + "moe", + "mamba2", + "moe", + "mamba2", + "full_attention", + "moe", + "mamba2", + "moe", + "mamba2", + "moe", + "mamba2", + "full_attention", + "moe", + "mamba2", + "moe", + "mamba2", + "moe", + "mamba2", + "full_attention", + "moe", + "mamba2", + "moe", + "mamba2", + "moe", + "mamba2", + "full_attention", + "moe", + "mamba2", + "moe", + "mamba2", + "moe", + "mamba2", + "full_attention", + "moe", + "mamba2", + "moe", + "mamba2", + "moe", + "mamba2", + "moe", + "mamba2", + "full_attention", + "moe", + "mamba2", + "moe", + "mamba2", + "moe", + "mamba2", + "moe", + "mamba2", + "moe", +) + +_MAMBA_SUFFIXES = frozenset( + { + "attn_norm.weight", + "ssm_a", + "ssm_conv1d.bias", + "ssm_conv1d.weight", + "ssm_d", + "ssm_dt.bias", + "ssm_in.weight", + "ssm_norm.weight", + "ssm_out.weight", + } +) +_MOE_SUFFIXES = frozenset( + { + "attn_norm.weight", + "exp_probs_b.bias", + "ffn_down_exps.weight", + "ffn_down_shexp.weight", + "ffn_gate_inp.weight", + "ffn_up_exps.weight", + "ffn_up_shexp.weight", + } +) +_ATTENTION_SUFFIXES = frozenset( + { + "attn_k.weight", + "attn_norm.weight", + "attn_output.weight", + "attn_q.weight", + "attn_v.weight", + } +) +_Q8_BASE_SUFFIXES = frozenset( + { + "attn_k.weight", + "attn_output.weight", + "attn_q.weight", + "attn_v.weight", + "ffn_down_exps.weight", + "ffn_down_shexp.weight", + "ffn_up_exps.weight", + "ffn_up_shexp.weight", + "ssm_in.weight", + "ssm_out.weight", + } +) +_MTP_SUFFIXES = frozenset( + { + "attn_k.weight", + "attn_norm.weight", + "attn_output.weight", + "attn_q.weight", + "attn_v.weight", + "exp_probs_b.bias", + "ffn_down_exps.weight", + "ffn_down_shexp.weight", + "ffn_gate_inp.weight", + "ffn_up_exps.weight", + "ffn_up_shexp.weight", + "nextn.eh_proj.weight", + "nextn.enorm.weight", + "nextn.hnorm.weight", + "nextn.shared_head_norm.weight", + "post_attention_norm.weight", + } +) +_EXPECTED_SUFFIXES = { + "mamba2": _MAMBA_SUFFIXES, + "moe": _MOE_SUFFIXES, + "full_attention": _ATTENTION_SUFFIXES, +} + + +def _as_int(value: Any, name: str) -> int: + if value is None: + raise ValueError(f"nemotron_h_moe GGUF is missing metadata {name!r}") + return int(value) + + +def _qtype_name(qtype: Any) -> str: + return str(getattr(qtype, "name", qtype)) + + +def _reverse_permute_array(value: np.ndarray, n_head: int) -> np.ndarray: + dim = value.shape[0] // n_head // 2 + return value.reshape(n_head, dim, 2, *value.shape[1:]).swapaxes(1, 2).reshape(value.shape) + + +@register_architecture_adapter +class NemotronHMoEAdapter(GGUFArchitectureAdapter): + """Strict adapter for the pinned Nemotron 3.5 Lightning GGUF layout.""" + + architecture = "nemotron_h_moe" + model_type = "nemotron_h" + + def __init__(self, model: GGUFModel) -> None: + super().__init__(model) + self._block_suffixes = self._collect_block_suffixes(model.tensor_names) + self._mtp_blocks = tuple( + index + for index, suffixes in sorted(self._block_suffixes.items()) + if any(suffix.startswith("nextn.") for suffix in suffixes) + ) + self._layer_types = self._derive_layer_types() + self._config: NemotronHConfig | None = None + + @staticmethod + def _collect_block_suffixes(tensor_names: list[str]) -> dict[int, frozenset[str]]: + result: dict[int, set[str]] = {} + for name in tensor_names: + match = _BLOCK_RE.match(name) + if match is not None: + result.setdefault(int(match.group(1)), set()).add(match.group(2)) + return {index: frozenset(suffixes) for index, suffixes in result.items()} + + def _derive_layer_types(self) -> tuple[str, ...]: + layer_types = [] + for block_index, suffixes in sorted(self._block_suffixes.items()): + if block_index in self._mtp_blocks: + continue + kinds = [] + if any(suffix.startswith("ssm_") for suffix in suffixes): + kinds.append("mamba2") + if any(suffix.startswith(("ffn_", "exp_probs_")) for suffix in suffixes): + kinds.append("moe") + if any( + suffix.startswith(("attn_q.", "attn_k.", "attn_v.", "attn_output.")) + for suffix in suffixes + ): + kinds.append("full_attention") + if len(kinds) != 1: + raise ValueError( + f"Nemotron backbone block {block_index} has ambiguous mixer types: {kinds}" + ) + if block_index != len(layer_types): + raise ValueError( + "Nemotron backbone blocks must be contiguous from zero; " + f"expected {len(layer_types)}, found {block_index}" + ) + layer_types.append(kinds[0]) + return tuple(layer_types) + + def validate_model(self, *, source: str) -> None: + metadata = self.model.metadata + block_count = _as_int(metadata.get(f"{self.architecture}.block_count"), "block_count") + if block_count != 53: + raise ValueError( + f"Expected 53 Nemotron GGUF blocks, got {block_count} in {source!r}" + ) + if self._layer_types != _PINNED_LAYER_TYPES: + counts = Counter(self._layer_types) + raise ValueError( + "Nemotron GGUF backbone schedule differs from the pinned 52-layer " + f"contract: {dict(counts)}" + ) + if self._mtp_blocks != (52,): + raise ValueError( + f"Expected separate combined attention+MoE MTP block 52, got {self._mtp_blocks}" + ) + + for index, layer_type in enumerate(self._layer_types): + actual = self._block_suffixes[index] + expected = _EXPECTED_SUFFIXES[layer_type] + if actual != expected: + raise ValueError( + f"Nemotron block {index} ({layer_type}) tensor inventory mismatch; " + f"missing={sorted(expected - actual)}, extra={sorted(actual - expected)}" + ) + if self._block_suffixes[52] != _MTP_SUFFIXES: + actual = self._block_suffixes[52] + raise ValueError( + "Nemotron MTP block 52 tensor inventory mismatch; " + f"missing={sorted(_MTP_SUFFIXES - actual)}, " + f"extra={sorted(actual - _MTP_SUFFIXES)}" + ) + + records = list(self.model._reader.tensors) + if len(records) != 417: + raise ValueError(f"Expected 417 pinned Nemotron tensors, got {len(records)}") + + unsupported: Counter[str] = Counter() + qtype_counts: Counter[str] = Counter() + base_qtype_counts: Counter[str] = Counter() + mtp_qtype_counts: Counter[str] = Counter() + for record in records: + qtype = _qtype_name(record.tensor_type) + qtype_counts[qtype] += 1 + shape = tuple(int(dim) for dim in reversed(record.shape)) + is_mtp = record.name.startswith("blk.52.") + (mtp_qtype_counts if is_mtp else base_qtype_counts)[qtype] += 1 + if not is_mtp: + match = _BLOCK_RE.match(record.name) + suffix = match.group(2) if match is not None else None + expected_qtype = ( + "Q8_0" + if record.name in {"token_embd.weight", "output.weight"} + or suffix in _Q8_BASE_SUFFIXES + else "F32" + ) + if qtype in {"F32", "F16", "BF16", "Q8_0"} and qtype != expected_qtype: + raise ValueError( + f"Nemotron base tensor {record.name!r} has qtype {qtype}, " + f"expected {expected_qtype}; exact Q8 preservation does not " + "dequantize a quantized source into a float-only target" + ) + if not is_mtp and qtype not in {"F32", "F16", "BF16", "Q8_0"}: + unsupported[qtype] += math.prod(shape) + + if unsupported: + observed = ", ".join( + f"{qtype}={parameters:,} parameters" + for qtype, parameters in sorted(unsupported.items()) + ) + raise NotImplementedError( + "Nemotron 3.5 GGUF conversion currently preserves the validated " + f"Q8_0 production slice only; observed unsupported base types: {observed}. " + "These source formats do not have a validated runtime kernel mapping, " + "so conversion is refused instead of dequantizing and calling it preservation." + ) + + if base_qtype_counts != Counter({"F32": 237, "Q8_0": 164}): + raise ValueError( + "Pinned Nemotron base qtype inventory mismatch: " + f"{dict(sorted(base_qtype_counts.items()))}" + ) + if mtp_qtype_counts != Counter({"Q8_0": 9, "F32": 6, "BF16": 1}): + raise ValueError( + "Pinned Nemotron MTP qtype inventory mismatch: " + f"{dict(sorted(mtp_qtype_counts.items()))}" + ) + if qtype_counts != Counter({"Q8_0": 173, "F32": 243, "BF16": 1}): + raise ValueError( + f"Pinned Nemotron total qtype inventory mismatch: {dict(qtype_counts)}" + ) + + tokenizer_model = metadata.get("tokenizer.ggml.model") + tokenizer_pre = metadata.get("tokenizer.ggml.pre") + if (tokenizer_model, tokenizer_pre) != ("gpt2", "pixtral"): + raise ValueError( + "Nemotron tokenizer contract requires GGUF GPT-2/Pixtral BPE metadata; " + f"got model={tokenizer_model!r}, pre={tokenizer_pre!r}" + ) + special_ids = { + "bos": metadata.get("tokenizer.ggml.bos_token_id"), + "eos": metadata.get("tokenizer.ggml.eos_token_id"), + "padding": metadata.get("tokenizer.ggml.padding_token_id"), + } + if special_ids != {"bos": 1, "eos": 11, "padding": 999}: + raise ValueError( + f"Unexpected pinned Nemotron GGUF special-token ids: {special_ids}" + ) + + def build_config(self) -> NemotronHConfig: + if self._config is not None: + return self._config + + metadata = self.model.metadata + prefix = f"{self.architecture}." + hidden_size = _as_int(metadata.get(prefix + "embedding_length"), "embedding_length") + attention_heads = _as_int( + metadata.get(prefix + "attention.head_count"), "attention.head_count" + ) + head_dim = _as_int( + metadata.get(prefix + "attention.key_length"), "attention.key_length" + ) + kv_values = metadata.get(prefix + "attention.head_count_kv") + if not isinstance(kv_values, list): + raise TypeError("nemotron_h_moe attention.head_count_kv must be a per-block list") + nonzero_kv_heads = {int(value) for value in kv_values if int(value) > 0} + if len(nonzero_kv_heads) != 1: + raise ValueError( + f"Expected one nonzero Nemotron KV-head count, got {sorted(nonzero_kv_heads)}" + ) + + inner_size = _as_int(metadata.get(prefix + "ssm.inner_size"), "ssm.inner_size") + mamba_heads = _as_int( + metadata.get(prefix + "ssm.time_step_rank"), "ssm.time_step_rank" + ) + if inner_size % mamba_heads: + raise ValueError( + f"SSM inner size {inner_size} is not divisible by {mamba_heads} Mamba heads" + ) + + self._config = NemotronHConfig( + vocab_size=_as_int(metadata.get(prefix + "vocab_size"), "vocab_size"), + hidden_size=hidden_size, + intermediate_size=_as_int( + metadata.get(prefix + "expert_feed_forward_length"), + "expert_feed_forward_length", + ), + num_hidden_layers=len(self._layer_types), + num_attention_heads=attention_heads, + num_key_value_heads=nonzero_kv_heads.pop(), + head_dim=head_dim, + hidden_act="relu2", + # The GGUF declares EOS=11 and PAD=999, while the pinned runtime + # contract uses model EOS=2/PAD=0 and accepts 11 as a second + # generation stop. Tokenizer asset padding remains role token 11. + bos_token_id=1, + eos_token_id=2, + pad_token_id=0, + tie_word_embeddings=False, + attn_qkv_bias=False, + attn_o_bias=False, + dtype=ir.DataType.FLOAT, + max_position_embeddings=_as_int( + metadata.get(prefix + "context_length"), "context_length" + ), + layer_types=list(self._layer_types), + rms_norm_eps=float( + metadata.get(prefix + "attention.layer_norm_rms_epsilon", 1e-5) + ), + rope_type=None, + rope_theta=None, + partial_rotary_factor=None, + mlp_bias=False, + num_local_experts=_as_int(metadata.get(prefix + "expert_count"), "expert_count"), + num_experts_per_tok=_as_int( + metadata.get(prefix + "expert_used_count"), "expert_used_count" + ), + moe_intermediate_size=_as_int( + metadata.get(prefix + "expert_feed_forward_length"), + "expert_feed_forward_length", + ), + shared_expert_intermediate_size=_as_int( + metadata.get(prefix + "expert_shared_feed_forward_length"), + "expert_shared_feed_forward_length", + ), + norm_topk_prob=bool(metadata.get(prefix + "expert_weights_norm", True)), + n_group=_as_int(metadata.get(prefix + "expert_group_count"), "expert_group_count"), + topk_group=_as_int( + metadata.get(prefix + "expert_group_used_count"), "expert_group_used_count" + ), + routed_scaling_factor=float(metadata.get(prefix + "expert_weights_scale", 1.0)), + scoring_func="sigmoid", + n_shared_experts=_as_int( + metadata.get(prefix + "expert_shared_count"), "expert_shared_count" + ), + num_nextn_predict_layers=_as_int( + metadata.get(prefix + "nextn_predict_layers"), "nextn_predict_layers" + ), + mamba_n_heads=mamba_heads, + mamba_d_head=inner_size // mamba_heads, + mamba_d_state=_as_int(metadata.get(prefix + "ssm.state_size"), "ssm.state_size"), + mamba_n_groups=_as_int( + metadata.get(prefix + "ssm.group_count"), "ssm.group_count" + ), + mamba_d_conv=_as_int(metadata.get(prefix + "ssm.conv_kernel"), "ssm.conv_kernel"), + mamba_expand=2, + mamba_conv_bias=True, + mamba_proj_bias=False, + mamba_time_step_min=0.001, + mamba_ssm_cache_dtype=ir.DataType.FLOAT, + moe_latent_size=None, + ) + self._config._gguf_model_type = self.model_type + self._config.model_type = self.model_type + return self._config + + def quantization_config(self) -> QuantizationConfig: + return QuantizationConfig( + bits=8, + group_size=32, + quant_method="gguf", + sym=False, + quantize_embeddings=True, + quantize_lm_head=True, + tie_word_embeddings=False, + ) + + def _target( + self, + state_dict_name: str, + initializer_name: str, + *, + source_index: int | None = None, + ) -> GGUFTensorMapping: + return GGUFTensorMapping( + ( + GGUFTensorTarget( + state_dict_name, + initializer_name, + source_index=source_index, + ), + ) + ) + + def _validate_source_shape( + self, + source_name: str, + actual: tuple[int, ...], + expected: tuple[int, ...], + ) -> None: + if actual != expected: + raise ValueError( + f"Nemotron GGUF tensor {source_name!r} has shape {actual}, expected {expected}" + ) + + def map_tensor( + self, + source_name: str, + source_shape: tuple[int, ...], + ) -> GGUFTensorMapping | None: + config = self.build_config() + h = config.hidden_size + q = config.num_attention_heads * config.head_dim + kv = config.num_key_value_heads * config.head_dim + d_inner = config.mamba_n_heads * config.mamba_d_head + conv_dim = d_inner + 2 * config.mamba_n_groups * config.mamba_d_state + experts = config.num_local_experts + moe_inner = config.moe_intermediate_size + shared_inner = config.shared_expert_intermediate_size + assert experts is not None + assert moe_inner is not None + assert shared_inner is not None + + global_mapping = { + "token_embd.weight": ( + "backbone.embeddings.weight", + "model.embed_tokens.weight", + (config.vocab_size, h), + ), + "output_norm.weight": ( + "backbone.norm_f.weight", + "model.norm.weight", + (h,), + ), + "output.weight": ("lm_head.weight", "lm_head.weight", (config.vocab_size, h)), + } + if source_name in global_mapping: + state_name, initializer_name, expected_shape = global_mapping[source_name] + self._validate_source_shape(source_name, source_shape, expected_shape) + return self._target(state_name, initializer_name) + + match = _BLOCK_RE.match(source_name) + if match is None: + return None + block_index = int(match.group(1)) + suffix = match.group(2) + if block_index in self._mtp_blocks: + return GGUFTensorMapping.excluded("auxiliary MTP block outside the decoder graph") + if block_index >= len(self._layer_types): + return None + + state_prefix = f"backbone.layers.{block_index}" + init_prefix = f"model.layers.{block_index}" + if suffix == "attn_norm.weight": + self._validate_source_shape(source_name, source_shape, (h,)) + return self._target( + f"{state_prefix}.norm.weight", + f"{init_prefix}.norm.weight", + ) + + layer_type = self._layer_types[block_index] + if layer_type == "mamba2": + mappings = { + "ssm_in.weight": ( + "in_proj.weight", + (d_inner + conv_dim + config.mamba_n_heads, h), + ), + "ssm_out.weight": ("out_proj.weight", (h, d_inner)), + "ssm_conv1d.weight": ( + "conv1d.weight", + (conv_dim, config.mamba_d_conv), + ), + "ssm_conv1d.bias": ("conv1d.bias", (conv_dim,)), + "ssm_a": ("A_log", (config.mamba_n_heads, 1)), + "ssm_d": ("D", (config.mamba_n_heads, 1)), + "ssm_dt.bias": ("dt_bias", (config.mamba_n_heads,)), + "ssm_norm.weight": ( + "norm.weight", + (config.mamba_n_groups, d_inner // config.mamba_n_groups), + ), + } + mapped = mappings.get(suffix) + if mapped is None: + return None + target_suffix, expected_shape = mapped + self._validate_source_shape(source_name, source_shape, expected_shape) + return self._target( + f"{state_prefix}.mixer.{target_suffix}", + f"{init_prefix}.mamba.{target_suffix}", + ) + + if layer_type == "full_attention": + mappings = { + "attn_q.weight": ("q_proj.weight", (q, h)), + "attn_k.weight": ("k_proj.weight", (kv, h)), + "attn_v.weight": ("v_proj.weight", (kv, h)), + "attn_output.weight": ("o_proj.weight", (h, q)), + } + mapped = mappings.get(suffix) + if mapped is None: + return None + target_suffix, expected_shape = mapped + self._validate_source_shape(source_name, source_shape, expected_shape) + return self._target( + f"{state_prefix}.mixer.{target_suffix}", + f"{init_prefix}.self_attn.{target_suffix}", + ) + + mappings = { + "ffn_gate_inp.weight": ("gate.weight", "gate.weight", (experts, h)), + "exp_probs_b.bias": ( + "gate.e_score_correction_bias", + "gate.e_score_correction_bias", + (experts,), + ), + "ffn_up_shexp.weight": ( + "shared_experts.up_proj.weight", + "shared_experts.up_proj.weight", + (shared_inner, h), + ), + "ffn_down_shexp.weight": ( + "shared_experts.down_proj.weight", + "shared_experts.down_proj.weight", + (h, shared_inner), + ), + } + mapped = mappings.get(suffix) + if mapped is not None: + state_suffix, init_suffix, expected_shape = mapped + self._validate_source_shape(source_name, source_shape, expected_shape) + return self._target( + f"{state_prefix}.mixer.{state_suffix}", + f"{init_prefix}.moe.{init_suffix}", + ) + + if suffix in {"ffn_up_exps.weight", "ffn_down_exps.weight"}: + projection = "up_proj" if suffix.startswith("ffn_up") else "down_proj" + expected_shape = ( + (experts, moe_inner, h) if projection == "up_proj" else (experts, h, moe_inner) + ) + self._validate_source_shape(source_name, source_shape, expected_shape) + targets = tuple( + GGUFTensorTarget( + f"{state_prefix}.mixer.experts.{index}.{projection}.weight", + f"{init_prefix}.moe.experts.{index}.{projection}.weight", + source_index=index, + ) + for index in range(experts) + ) + return GGUFTensorMapping(targets) + return None + + def transform_tensor( + self, + source_name: str, + target: GGUFTensorTarget, + tensor: torch.Tensor, + ) -> torch.Tensor: + del target + if source_name.endswith(".ssm_a"): + tensor = tensor.squeeze(-1) + if not torch.all(tensor < 0): + raise ValueError(f"Nemotron SSM A tensor {source_name!r} must be negative") + return torch.log(-tensor) + if source_name.endswith(".ssm_d"): + return tensor.squeeze(-1) + if source_name.endswith(".ssm_norm.weight"): + return tensor.flatten() + if source_name.endswith(".ssm_conv1d.weight"): + return tensor.unsqueeze(1) + if source_name.endswith(".attn_q.weight"): + return _reverse_permute(tensor, self.build_config().num_attention_heads) + if source_name.endswith(".attn_k.weight"): + return _reverse_permute(tensor, self.build_config().num_key_value_heads) + return tensor + + def transform_repacked( + self, + source_name: str, + target: GGUFTensorTarget, + tensor: RepackedTensor, + ) -> RepackedTensor: + del target + n_head = None + if source_name.endswith(".attn_q.weight"): + n_head = self.build_config().num_attention_heads + elif source_name.endswith(".attn_k.weight"): + n_head = self.build_config().num_key_value_heads + if n_head is None: + return tensor + return RepackedTensor( + weight=_reverse_permute_array(tensor.weight, n_head), + scales=_reverse_permute_array(tensor.scales, n_head), + zero_points=( + None + if tensor.zero_points is None + else _reverse_permute_array(tensor.zero_points, n_head) + ), + block_size=tensor.block_size, + bits=tensor.bits, + ) + + def validate_mapping_audit(self, audit: GGUFMappingAudit) -> None: + super().validate_mapping_audit(audit) + if len(audit.mapped_sources) != 401: + raise ValueError( + f"Expected 401 mapped Nemotron sources, got {len(audit.mapped_sources)}" + ) + if len(audit.excluded_sources) != 16: + raise ValueError( + f"Expected 16 explicit MTP exclusions, got {len(audit.excluded_sources)}" + ) + if len(audit.target_sources) != 6243: + raise ValueError( + f"Expected 6243 logical Nemotron targets, got {len(audit.target_sources)}" + ) diff --git a/src/mobius/integrations/gguf/_repacker.py b/src/mobius/integrations/gguf/_repacker.py index c54df7055..0131d8edd 100644 --- a/src/mobius/integrations/gguf/_repacker.py +++ b/src/mobius/integrations/gguf/_repacker.py @@ -232,6 +232,68 @@ def repack_gguf_tensor( return _repack_q8_0(blocks, n_out, n_blocks_per_row) +def repack_stacked_gguf_tensor( + raw_data: np.ndarray, + gguf_type: int, + shape: tuple[int, ...], +) -> tuple[RepackedTensor, ...]: + """Repack a leading-axis stack of 2-D weights without dequantizing. + + The source block rows must end at each matrix boundary. This is true for + routed-expert tensors whose input dimension is block-aligned and prevents + an expert slice from accidentally consuming bytes from the next expert. + """ + if len(shape) != 3: + raise ValueError(f"Expected stacked 3D shape (E, N, K), got {shape}") + if gguf_type not in _SUPPORTED_TYPES: + raise ValueError( + f"Unsupported GGUF type {gguf_type}. Supported: {sorted(_SUPPORTED_TYPES)}" + ) + + num_slices, n_out, k_in = shape + block_elements = _GGUF_BLOCK_ELEMENTS[gguf_type] + if k_in % block_elements: + raise ValueError( + f"Stacked tensor input dimension {k_in} is not aligned to the " + f"{block_elements}-element GGUF block" + ) + + expected_bytes = num_slices * n_out * (k_in // block_elements) * _BLOCK_BYTES[gguf_type] + packed = raw_data.ravel().view(np.uint8) + if packed.size != expected_bytes: + raise ValueError( + f"Stacked GGUF data size mismatch: got {packed.size} bytes, " + f"expected {expected_bytes} for shape {shape}" + ) + + combined = repack_gguf_tensor( + packed, + gguf_type, + (num_slices * n_out, k_in), + ) + weight = combined.weight.reshape(num_slices, n_out, *combined.weight.shape[1:]) + scales = combined.scales.reshape(num_slices, n_out, *combined.scales.shape[1:]) + zero_points = ( + None + if combined.zero_points is None + else combined.zero_points.reshape( + num_slices, + n_out, + *combined.zero_points.shape[1:], + ) + ) + return tuple( + RepackedTensor( + weight=weight[index], + scales=scales[index], + zero_points=None if zero_points is None else zero_points[index], + block_size=combined.block_size, + bits=combined.bits, + ) + for index in range(num_slices) + ) + + def _reorder_nibbles_gguf_to_ort( gguf_packed: np.ndarray, ) -> np.ndarray: diff --git a/src/mobius/integrations/gguf/_repacker_test.py b/src/mobius/integrations/gguf/_repacker_test.py index d2f6e059e..30186ba86 100644 --- a/src/mobius/integrations/gguf/_repacker_test.py +++ b/src/mobius/integrations/gguf/_repacker_test.py @@ -15,6 +15,7 @@ preserve_native_blocks, repack_dequantized_tensor, repack_gguf_tensor, + repack_stacked_gguf_tensor, ) _Q4_0 = 2 @@ -399,6 +400,105 @@ def test_round_trip_dequantize(self): np.testing.assert_allclose(ort_deq, gguf_deq.ravel(), atol=1e-3) + def test_stacked_experts_execute_as_individual_matmulnbits(self, tmp_path): + import onnx_ir as ir + import onnxruntime as ort + + blocks = [] + source = np.empty((2, 2, 32), dtype=np.float32) + for expert in range(2): + for row in range(2): + values = [((index * (expert + 1) + row * 7) % 31) - 15 for index in range(32)] + scale = 0.125 * (expert + row + 1) + blocks.append(_make_q8_0_block(scale, values)) + source[expert, row] = np.asarray(values, dtype=np.float32) * np.float16(scale) + repacked = repack_stacked_gguf_tensor( + np.concatenate(blocks), + _Q8_0, + (2, 2, 32), + ) + + def _value(name: str, array: np.ndarray) -> ir.Value: + value = ir.Value(name=name) + value.const_value = ir.tensor(array) + value.shape = ir.Shape(array.shape) + value.dtype = value.const_value.dtype + return value + + x = ir.Value( + name="x", + shape=ir.Shape([1, 32]), + type=ir.TensorType(ir.DataType.FLOAT), + ) + nodes = [] + initializers = [] + expert_outputs = [] + for index, tensor in enumerate(repacked): + weight = _value(f"expert_{index}.weight", tensor.weight) + scales = _value( + f"expert_{index}.scales", + tensor.scales.astype(np.float32), + ) + zero_points = _value( + f"expert_{index}.zero_points", + tensor.zero_points, + ) + output = ir.Value( + name=f"expert_{index}.output", + shape=ir.Shape([1, 2]), + type=ir.TensorType(ir.DataType.FLOAT), + ) + nodes.append( + ir.Node( + "com.microsoft", + "MatMulNBits", + inputs=[x, weight, scales, zero_points], + outputs=[output], + attributes=ir.convenience.convert_attributes( + {"K": 32, "N": 2, "bits": 8, "block_size": 32} + ), + ) + ) + initializers.extend([weight, scales, zero_points]) + expert_outputs.append(output) + + output = ir.Value( + name="output", + shape=ir.Shape([1, 4]), + type=ir.TensorType(ir.DataType.FLOAT), + ) + nodes.append( + ir.Node( + "", + "Concat", + inputs=expert_outputs, + outputs=[output], + attributes=ir.convenience.convert_attributes({"axis": -1}), + ) + ) + graph = ir.Graph( + inputs=[x], + outputs=[output], + nodes=nodes, + initializers=initializers, + opset_imports={"": 18, "com.microsoft": 1}, + name="stacked_experts_q8", + ) + path = tmp_path / "stacked_experts_q8.onnx" + ir.save(ir.Model(graph, ir_version=10), path) + + feed = np.linspace(-1.0, 1.0, 32, dtype=np.float32)[None, :] + session = ort.InferenceSession( + str(path), + providers=["CPUExecutionProvider"], + ) + (actual,) = session.run(None, {"x": feed}) + expected = np.concatenate( + [feed @ source[index].T for index in range(2)], + axis=-1, + ) + np.testing.assert_allclose(actual, expected, rtol=1e-5, atol=1e-5) + class TestEdgeCases: def test_unsupported_type_raises(self): diff --git a/src/mobius/integrations/gguf/_tokenizer.py b/src/mobius/integrations/gguf/_tokenizer.py index 49cb650fc..ad7cb03ce 100644 --- a/src/mobius/integrations/gguf/_tokenizer.py +++ b/src/mobius/integrations/gguf/_tokenizer.py @@ -21,19 +21,114 @@ import logging import os from pathlib import Path +from typing import Any _LOGGER = logging.getLogger(__name__) +# Mistral/Pixtral's Unicode-aware split expression. This is the expression +# serialized by NVIDIA's Nemotron 3.5 Lightning tokenizer.json. +_PIXTRAL_GPT2_SPLIT_REGEX = ( + r"[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]*[\p{Ll}\p{Lm}\p{Lo}\p{M}]+" + r"|[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]+[\p{Ll}\p{Lm}\p{Lo}\p{M}]*" + r"|\p{N}| ?[^\s\p{L}\p{N}]+[\r\n/]*|\s*[\r\n]+|\s+(?!\S)|\s+" +) + + +def _is_pixtral_gpt2_profile(metadata: dict[str, Any]) -> bool: + """Whether GGUF metadata needs the strict Pixtral/GPT-2 BPE profile.""" + return ( + metadata.get("tokenizer.ggml.model"), + metadata.get("tokenizer.ggml.pre"), + ) == ("gpt2", "pixtral") + + +def _reconstruct_pixtral_gpt2_tokenizer( + metadata: dict[str, Any], output_dir: str | Path +) -> str: + """Reconstruct the exact ByteLevel BPE tokenizer used by Nemotron Lightning. + + GGUF token types distinguish control tokens from user-defined tokens. + Tokenizers must expose the former as special and the latter as ordinary + added tokens; treating both as special changes chat-template tokenization. + """ + from tokenizers import AddedToken, Regex, Tokenizer, decoders, pre_tokenizers, processors + from tokenizers.models import BPE + + tokens = metadata.get("tokenizer.ggml.tokens") + merges_raw = metadata.get("tokenizer.ggml.merges") + token_types = metadata.get("tokenizer.ggml.token_type") + if not isinstance(tokens, list) or not tokens: + raise ValueError("Pixtral/GPT-2 GGUF tokenizer is missing tokenizer.ggml.tokens.") + if not isinstance(merges_raw, list) or not merges_raw: + raise ValueError("Pixtral/GPT-2 GGUF tokenizer is missing tokenizer.ggml.merges.") + if not isinstance(token_types, list) or len(token_types) != len(tokens): + raise ValueError( + "Pixtral/GPT-2 GGUF tokenizer requires tokenizer.ggml.token_type for every token." + ) + if not all( + isinstance(token_type, int) and token_type in (1, 3, 4) for token_type in token_types + ): + raise ValueError("Pixtral/GPT-2 GGUF tokenizer contains an unsupported token type.") + if not all(isinstance(token, str) for token in tokens): + raise ValueError("Pixtral/GPT-2 GGUF tokenizer contains a non-string token.") + if len(set(tokens)) != len(tokens): + raise ValueError("Pixtral/GPT-2 GGUF tokenizer contains duplicate tokens.") + + merges: list[tuple[str, str]] = [] + for merge in merges_raw: + if not isinstance(merge, str): + raise ValueError( # noqa: TRY004 - malformed GGUF metadata is a value error. + "Pixtral/GPT-2 GGUF tokenizer contains a non-string merge." + ) + parts = merge.split(" ") + if len(parts) != 2 or not all(parts): + raise ValueError(f"Pixtral/GPT-2 GGUF tokenizer has an invalid merge: {merge!r}.") + merges.append((parts[0], parts[1])) + + control_tokens = [ + AddedToken(tokens[index], normalized=False, special=True) + for index, token_type in enumerate(token_types) + if token_type == 3 + ] + user_defined_tokens = [ + AddedToken(tokens[index], normalized=False, special=False) + for index, token_type in enumerate(token_types) + if token_type == 4 + ] + + tokenizer = Tokenizer( + BPE( + vocab={token: index for index, token in enumerate(tokens)}, + merges=merges, + ignore_merges=True, + ) + ) + # Match the Tekken/Pixtral tokenizer: Unicode words are split before byte + # encoding, so the BPE vocab's GPT-2 byte symbols are consumed verbatim. + tokenizer.pre_tokenizer = pre_tokenizers.Sequence( + [ + pre_tokenizers.Split(Regex(_PIXTRAL_GPT2_SPLIT_REGEX), behavior="isolated"), + pre_tokenizers.ByteLevel(add_prefix_space=False, use_regex=False), + ] + ) + tokenizer.decoder = decoders.ByteLevel() + tokenizer.post_processor = processors.ByteLevel(trim_offsets=False) + tokenizer.add_special_tokens(control_tokens) + tokenizer.add_tokens(user_defined_tokens) + + path = os.path.join(str(output_dir), "tokenizer.json") + tokenizer.save(path) + return path + def write_gguf_tokenizer_json(gguf_path: str | Path, output_dir: str | Path) -> str | None: """Write ``tokenizer.json`` for a GGUF-built package. Reconstructs the fast tokenizer from the GGUF ``tokenizer.ggml.*`` metadata - and serializes it to ``/tokenizer.json``. This is best-effort: - it logs a warning and returns ``None`` (never raising) when ``transformers`` - is unavailable or the GGUF's tokenizer model cannot be converted, so the - build is not blocked — the onnx-genai runners can be given a - ``tokenizer.json`` separately. + and serializes it to ``/tokenizer.json``. The verified + GPT-2/Pixtral profile is strict and raises for malformed metadata. Other + profiles retain the historical best-effort behavior and return ``None`` if + no supported converter is available. Args: gguf_path: Path to the ``.gguf`` file whose embedded tokenizer to emit. @@ -44,6 +139,23 @@ def write_gguf_tokenizer_json(gguf_path: str | Path, output_dir: str | Path) -> emitted. """ gguf_path = Path(gguf_path) + # Do not let AutoTokenizer select its generic SentencePiece conversion for + # this known GPT-2 profile. Its reconstruction has strict, verified + # ByteLevel semantics and must report malformed metadata rather than emit a + # tokenizer that looks valid but produces different ids. + try: + from mobius.integrations.gguf._reader import GGUFModel + + gguf_model = GGUFModel(str(gguf_path)) + except (ImportError, IndexError, OSError, RuntimeError, TypeError, ValueError): + # Preserve the historical best-effort behavior for unreadable or + # unsupported GGUF files, including test doubles used by callers. + pass + else: + metadata = gguf_model.metadata + if _is_pixtral_gpt2_profile(metadata): + return _reconstruct_pixtral_gpt2_tokenizer(metadata, output_dir) + try: from transformers import AutoTokenizer except ImportError: @@ -154,7 +266,20 @@ def _reconstruct_tokenizer_from_ggml(gguf_path: Path, output_dir: str | Path) -> return None try: - metadata = GGUFModel(str(gguf_path)).metadata + gguf_model = GGUFModel(str(gguf_path)) + metadata = gguf_model.metadata + except (IndexError, OSError, RuntimeError, TypeError, ValueError) as error: + _LOGGER.warning( + "Failed to read tokenizer metadata from GGUF %r: %s; " + "skipping tokenizer.json emission.", + str(gguf_path), + error, + ) + return None + if _is_pixtral_gpt2_profile(metadata): + return _reconstruct_pixtral_gpt2_tokenizer(metadata, output_dir) + + try: tokens = metadata.get("tokenizer.ggml.tokens") merges_raw = metadata.get("tokenizer.ggml.merges") if not tokens or not merges_raw: diff --git a/src/mobius/integrations/gguf/_tokenizer_test.py b/src/mobius/integrations/gguf/_tokenizer_test.py index 32b58d26c..f4789ccfa 100644 --- a/src/mobius/integrations/gguf/_tokenizer_test.py +++ b/src/mobius/integrations/gguf/_tokenizer_test.py @@ -5,6 +5,8 @@ from __future__ import annotations +import hashlib +import json import os from pathlib import Path from unittest import mock @@ -191,3 +193,243 @@ def test_reconstructs_bpe_tokenizer_with_correct_ids(self, tmp_path): enc = tok.encode("hi") assert enc.ids[0] == 2 # assert tok.decode(enc.ids) == "hi" + + +def _write_gguf_with_pixtral_gpt2_tokenizer(path: Path) -> None: + """Write a tiny GGUF with the Nemotron/Pixtral ByteLevel BPE profile.""" + from gguf import GGUFWriter + + writer = GGUFWriter(str(path), "nemotron_h_moe") + tokens = [ + "", + "", + "", + "<|im_start|>", + "<|im_end|>", + "", + "Hello", + "Ġcafé", + "ä½łå¥½", + "x", + "y", + "xy", + ] + writer.add_tokenizer_model("gpt2") + writer.add_tokenizer_pre("pixtral") + writer.add_token_list(tokens) + writer.add_token_types([3, 3, 3, 3, 3, 4, 1, 1, 1, 1, 1, 1]) + writer.add_token_merges(["x y"]) + writer.write_header_to_file() + writer.write_kv_data_to_file() + writer.write_tensors_to_file() + writer.close() + + +class TestPixtralGpt2Tokenizer: + """Strict ByteLevel reconstruction for Nemotron's Pixtral/GPT-2 metadata.""" + + def test_profile_selection(self): + from mobius.integrations.gguf._tokenizer import _is_pixtral_gpt2_profile + + assert _is_pixtral_gpt2_profile( + { + "tokenizer.ggml.model": "gpt2", + "tokenizer.ggml.pre": "pixtral", + } + ) + assert not _is_pixtral_gpt2_profile({"tokenizer.ggml.model": "gpt2"}) + assert not _is_pixtral_gpt2_profile( + { + "tokenizer.ggml.model": "llama", + "tokenizer.ggml.pre": "pixtral", + } + ) + + def test_reconstructs_bytelevel_ids_and_contract(self, tmp_path: Path): + from tokenizers import Tokenizer + + from mobius.integrations.gguf._tokenizer import _reconstruct_tokenizer_from_ggml + + gguf_path = tmp_path / "nemotron.gguf" + _write_gguf_with_pixtral_gpt2_tokenizer(gguf_path) + out_dir = tmp_path / "out" + out_dir.mkdir() + + result = _reconstruct_tokenizer_from_ggml(gguf_path, out_dir) + + assert result == str(out_dir / "tokenizer.json") + tokenizer = Tokenizer.from_file(result) + assert tokenizer.encode("Hello café").ids == [6, 7] + assert tokenizer.encode("你好").ids == [8] + assert tokenizer.encode("<|im_end|>").ids == [4] + assert tokenizer.encode("").ids == [5] + + serialized = json.loads(Path(result).read_text(encoding="utf-8")) + assert serialized["model"]["type"] == "BPE" + assert serialized["model"]["ignore_merges"] is True + assert serialized["pre_tokenizer"]["type"] == "Sequence" + assert serialized["pre_tokenizer"]["pretokenizers"][1] == { + "type": "ByteLevel", + "add_prefix_space": False, + "trim_offsets": True, + "use_regex": False, + } + assert serialized["decoder"]["type"] == "ByteLevel" + assert serialized["post_processor"] == { + "type": "ByteLevel", + "add_prefix_space": True, + "trim_offsets": False, + "use_regex": True, + } + added_tokens = {token["id"]: token for token in serialized["added_tokens"]} + assert added_tokens[4]["special"] is True + assert added_tokens[5]["special"] is False + + def test_public_writer_bypasses_autotokenizer(self, tmp_path: Path): + from mobius.integrations.gguf import _tokenizer + + gguf_path = tmp_path / "nemotron.gguf" + _write_gguf_with_pixtral_gpt2_tokenizer(gguf_path) + out_dir = tmp_path / "out" + out_dir.mkdir() + + fake_transformers = mock.Mock() + fake_transformers.AutoTokenizer.from_pretrained.side_effect = AssertionError( + "strict GPT-2 profile must not use AutoTokenizer" + ) + with mock.patch.dict("sys.modules", {"transformers": fake_transformers}): + result = _tokenizer.write_gguf_tokenizer_json(gguf_path, out_dir) + + assert result == str(out_dir / "tokenizer.json") + fake_transformers.AutoTokenizer.from_pretrained.assert_not_called() + + def test_rejects_incomplete_strict_metadata(self, tmp_path: Path): + from mobius.integrations.gguf._tokenizer import _reconstruct_pixtral_gpt2_tokenizer + + with pytest.raises(ValueError, match="token_type"): + _reconstruct_pixtral_gpt2_tokenizer( + { + "tokenizer.ggml.tokens": ["a"], + "tokenizer.ggml.merges": ["a b"], + }, + tmp_path, + ) + + def test_private_fallback_preserves_strict_metadata_errors(self, tmp_path: Path): + from mobius.integrations.gguf._tokenizer import _reconstruct_tokenizer_from_ggml + + metadata = { + "tokenizer.ggml.model": "gpt2", + "tokenizer.ggml.pre": "pixtral", + "tokenizer.ggml.tokens": ["a"], + "tokenizer.ggml.merges": ["a b"], + } + with ( + mock.patch( + "mobius.integrations.gguf._reader.GGUFModel", + return_value=mock.Mock(metadata=metadata), + ), + pytest.raises(ValueError, match="token_type"), + ): + _reconstruct_tokenizer_from_ggml(tmp_path / "malformed.gguf", tmp_path) + + def test_pinned_artifact_when_explicitly_requested(self, tmp_path: Path): + """Validate pinned artifact checksums without downloading it in normal tests.""" + artifact = os.environ.get("MOBIUS_NEMOTRON_GGUF_TOKENIZER_TEST_PATH") + if artifact is None: + pytest.skip( + "Set MOBIUS_NEMOTRON_GGUF_TOKENIZER_TEST_PATH to run artifact validation." + ) + + from mobius.integrations.gguf._reader import GGUFModel + from mobius.integrations.gguf._tokenizer import _reconstruct_tokenizer_from_ggml + + metadata = GGUFModel(artifact).metadata + assert len(metadata["tokenizer.ggml.tokens"]) == 131072 + assert len(metadata["tokenizer.ggml.merges"]) == 269443 + assert ( + hashlib.sha256("\n".join(metadata["tokenizer.ggml.tokens"]).encode()).hexdigest() + == "4999709474e3c967358c1f1199b6be65fb9055d3eb59e0cd387f9e7077fc40ed" + ) + assert ( + hashlib.sha256("\n".join(metadata["tokenizer.ggml.merges"]).encode()).hexdigest() + == "b1b0165185b1925118c2f7b1e978439b02010c3a420ebfec5c19a093a0d9b4cb" + ) + + result = _reconstruct_tokenizer_from_ggml(artifact, tmp_path) + serialized = json.loads(Path(result).read_text(encoding="utf-8")) + assert serialized["model"]["ignore_merges"] is True + assert serialized["decoder"]["type"] == "ByteLevel" + + @pytest.mark.integration + def test_pinned_artifact_matches_official_tokenizer(self, tmp_path: Path): + artifact = os.environ.get("MOBIUS_NEMOTRON_GGUF_TOKENIZER_TEST_PATH") + official_dir_value = os.environ.get("MOBIUS_NEMOTRON_OFFICIAL_TOKENIZER_DIR") + if artifact is None or official_dir_value is None: + pytest.skip( + "Set MOBIUS_NEMOTRON_GGUF_TOKENIZER_TEST_PATH and " + "MOBIUS_NEMOTRON_OFFICIAL_TOKENIZER_DIR for pinned parity." + ) + + from tokenizers import Tokenizer + + from mobius.integrations.gguf._reader import GGUFModel + from mobius.integrations.gguf._tokenizer import _reconstruct_tokenizer_from_ggml + + official_dir = Path(official_dir_value) + official_path = official_dir / "tokenizer.json" + assert official_path.is_file() + rebuilt_path = _reconstruct_tokenizer_from_ggml(Path(artifact), tmp_path) + assert rebuilt_path is not None + + official = Tokenizer.from_file(str(official_path)) + rebuilt = Tokenizer.from_file(rebuilt_path) + assert official.get_vocab(with_added_tokens=True) == rebuilt.get_vocab( + with_added_tokens=True + ) + samples = ( + "The capital of France is", + " Paris. \nThe capital of Germany", + "Hello, world!", + "café déjà vu — 你好 🌍", + " leading\tspaces\r\nnewlines ", + "<|im_start|>assistant\nx<|im_end|>", + bytes(range(1, 128)).decode("latin1"), + ) + for sample in samples: + official_encoding = official.encode(sample) + rebuilt_encoding = rebuilt.encode(sample) + assert rebuilt_encoding.ids == official_encoding.ids + assert rebuilt.decode( + rebuilt_encoding.ids, skip_special_tokens=False + ) == official.decode(official_encoding.ids, skip_special_tokens=False) + + official_json = json.loads(official_path.read_text(encoding="utf-8")) + rebuilt_json = json.loads(Path(rebuilt_path).read_text(encoding="utf-8")) + official_added = { + (token["id"], token["content"]): token["special"] + for token in official_json["added_tokens"] + } + rebuilt_added = { + (token["id"], token["content"]): token["special"] + for token in rebuilt_json["added_tokens"] + } + assert rebuilt_added == official_added + assert rebuilt_json["pre_tokenizer"] == official_json["pre_tokenizer"] + assert rebuilt_json["decoder"] == official_json["decoder"] + assert rebuilt_json["post_processor"] == official_json["post_processor"] + + metadata = GGUFModel(artifact).metadata + gguf_template = metadata["tokenizer.chat_template"].replace("\r\n", "\n") + official_template = ( + (official_dir / "chat_template.jinja") + .read_text(encoding="utf-8") + .replace("\r\n", "\n") + ) + assert gguf_template != official_template + assert hashlib.sha256(gguf_template.encode()).hexdigest() == ( + "cbb337473ffde036fd4b6e7e7763dcb97c7cd8b4a311cd52d361d2766b00eb7c" + ) + assert hashlib.sha256( + (official_dir / "chat_template.jinja").read_bytes() + ).hexdigest() == ("58933db77d3099b4f78c55a38347a72e1ea05b97d6bd8f38775303dc0194e0a9") diff --git a/src/mobius/models/nemotron_h.py b/src/mobius/models/nemotron_h.py index ff7cc7529..9eb06e7da 100644 --- a/src/mobius/models/nemotron_h.py +++ b/src/mobius/models/nemotron_h.py @@ -43,8 +43,10 @@ Embedding, Linear, Mamba2Block, + QuantizedEmbedding, RMSNorm, create_padding_mask, + make_quantized_linear_factory, ) # --------------------------------------------------------------------------- @@ -52,6 +54,19 @@ # --------------------------------------------------------------------------- +def _quantized_linear_class(config: NemotronHConfig) -> type | None: + quantization = config.quantization + if quantization is None or quantization.quant_method == "none": + return None + zero_point_dtype = config.dtype if quantization.float_zero_point else ir.DataType.UINT8 + return make_quantized_linear_factory( + bits=quantization.bits, + block_size=quantization.group_size, + has_zero_point=not quantization.sym, + zero_point_dtype=zero_point_dtype, + ) + + class NemotronHMambaLayer(nn.Module): """NemotronH Mamba2 layer: RMSNorm → Mamba2Block → residual. @@ -61,7 +76,7 @@ class NemotronHMambaLayer(nn.Module): config: NemotronH architecture config. """ - def __init__(self, config: NemotronHConfig): + def __init__(self, config: NemotronHConfig, linear_class: type | None = None): super().__init__() # d_inner = num_heads * head_dim (not hidden_size * expand) d_inner = config.mamba_n_heads * config.mamba_d_head @@ -81,6 +96,7 @@ def __init__(self, config: NemotronHConfig): # group of heads_per_group * head_dim dimensions. norm_group_size=d_inner // config.mamba_n_groups, time_step_min=config.mamba_time_step_min, + linear_class=linear_class, ) self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) @@ -121,9 +137,9 @@ class NemotronHAttentionLayer(nn.Module): config: NemotronH architecture config. """ - def __init__(self, config: NemotronHConfig): + def __init__(self, config: NemotronHConfig, linear_class: type | None = None): super().__init__() - self.self_attn = Attention(config) + self.self_attn = Attention(config, linear_class=linear_class) self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) def forward( @@ -159,13 +175,14 @@ class NemotronHMLPLayer(nn.Module): config: NemotronH architecture config. """ - def __init__(self, config: NemotronHConfig): + def __init__(self, config: NemotronHConfig, linear_class: type | None = None): super().__init__() self.mlp = FCMLP( config.hidden_size, config.intermediate_size, activation=config.hidden_act, bias=config.mlp_bias, + linear_class=linear_class, ) self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) @@ -293,7 +310,7 @@ class NemotronHMoEBlock(nn.Module): HuggingFace reference: ``NemotronHMoE``. """ - def __init__(self, config: NemotronHConfig): + def __init__(self, config: NemotronHConfig, linear_class: type | None = None): super().__init__() assert config.num_local_experts is not None assert config.num_experts_per_tok is not None @@ -322,6 +339,7 @@ def __init__(self, config: NemotronHConfig): config.moe_intermediate_size, activation=config.hidden_act, bias=config.mlp_bias, + linear_class=linear_class, ) for _ in range(num_experts) ] @@ -336,18 +354,20 @@ def __init__(self, config: NemotronHConfig): shared_intermediate, activation=config.hidden_act, bias=config.mlp_bias, + linear_class=linear_class, ) # Optional latent projection (e.g. 120B: 4096 → 1024 → experts # → 1024 → 4096) self._has_latent = config.moe_latent_size is not None if self._has_latent: - self.fc1_latent_proj = Linear( + latent_linear_class = linear_class or Linear + self.fc1_latent_proj = latent_linear_class( config.hidden_size, config.moe_latent_size, bias=config.mlp_bias, ) - self.fc2_latent_proj = Linear( + self.fc2_latent_proj = latent_linear_class( config.moe_latent_size, config.hidden_size, bias=config.mlp_bias, @@ -404,9 +424,9 @@ class NemotronHMoELayer(nn.Module): config: NemotronH architecture config. """ - def __init__(self, config: NemotronHConfig): + def __init__(self, config: NemotronHConfig, linear_class: type | None = None): super().__init__() - self.moe = NemotronHMoEBlock(config) + self.moe = NemotronHMoEBlock(config, linear_class=linear_class) self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) def forward( @@ -450,22 +470,34 @@ class _NemotronHTextModel(nn.Module): def __init__(self, config: NemotronHConfig): super().__init__() self._dtype = config.dtype - self.embed_tokens = Embedding( - config.vocab_size, config.hidden_size, config.pad_token_id - ) + linear_class = _quantized_linear_class(config) + quantization = config.quantization + if quantization is not None and quantization.quantize_embeddings: + self.embed_tokens = QuantizedEmbedding( + config.vocab_size, + config.hidden_size, + bits=quantization.bits, + block_size=quantization.group_size, + has_zero_point=not quantization.sym, + padding_idx=config.pad_token_id, + ) + else: + self.embed_tokens = Embedding( + config.vocab_size, config.hidden_size, config.pad_token_id + ) layer_types = config.layer_types or [] self.layers = nn.ModuleList([]) for i in range(config.num_hidden_layers): ltype = layer_types[i] if i < len(layer_types) else "full_attention" if ltype == "mamba2": - self.layers.append(NemotronHMambaLayer(config)) + self.layers.append(NemotronHMambaLayer(config, linear_class=linear_class)) elif ltype == "mlp": - self.layers.append(NemotronHMLPLayer(config)) + self.layers.append(NemotronHMLPLayer(config, linear_class=linear_class)) elif ltype == "moe": - self.layers.append(NemotronHMoELayer(config)) + self.layers.append(NemotronHMoELayer(config, linear_class=linear_class)) else: - self.layers.append(NemotronHAttentionLayer(config)) + self.layers.append(NemotronHAttentionLayer(config, linear_class=linear_class)) self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) @@ -530,8 +562,16 @@ def __init__(self, config: NemotronHConfig): ) self.config = config self.model = _NemotronHTextModel(config) - self.lm_head = Linear(config.hidden_size, config.vocab_size, bias=False) - if config.tie_word_embeddings: + quantization = config.quantization + if quantization is not None and quantization.quantize_lm_head: + linear_class = _quantized_linear_class(config) + assert linear_class is not None + self.lm_head = linear_class(config.hidden_size, config.vocab_size, bias=False) + else: + self.lm_head = Linear(config.hidden_size, config.vocab_size, bias=False) + if config.tie_word_embeddings and not ( + quantization is not None and quantization.quantize_embeddings + ): self.lm_head.weight = self.model.embed_tokens.weight def forward( diff --git a/src/mobius/models/nemotron_h_test.py b/src/mobius/models/nemotron_h_test.py new file mode 100644 index 000000000..079a0d8c2 --- /dev/null +++ b/src/mobius/models/nemotron_h_test.py @@ -0,0 +1,122 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +from __future__ import annotations + +from collections import Counter + +import pytest +import torch + +from mobius._builder import build_from_module +from mobius._configs import NemotronHConfig, QuantizationConfig +from mobius._registry import registry +from mobius.integrations.gguf._architecture import validate_package_state_dict +from mobius.integrations.transformers import _default_task_for_model + + +def _config(*, quantization: QuantizationConfig | None = None) -> NemotronHConfig: + return NemotronHConfig( + hidden_size=64, + intermediate_size=32, + num_hidden_layers=3, + num_attention_heads=2, + num_key_value_heads=1, + head_dim=32, + vocab_size=256, + max_position_embeddings=128, + hidden_act="relu2", + pad_token_id=0, + layer_types=["mamba2", "moe", "full_attention"], + num_local_experts=2, + num_experts_per_tok=1, + moe_intermediate_size=32, + shared_expert_intermediate_size=64, + mamba_n_heads=2, + mamba_d_head=32, + mamba_d_state=16, + mamba_n_groups=2, + mamba_d_conv=4, + quantization=quantization, + ) + + +def _build_package(config: NemotronHConfig): + module = registry.get("nemotron_h")(config) + return build_from_module( + module, + config, + _default_task_for_model("nemotron_h"), + ) + + +def _build_graph(config: NemotronHConfig): + return _build_package(config)["model"].graph + + +def test_float_nemotron_retains_safetensors_projection_contract() -> None: + graph = _build_graph(_config()) + ops = Counter((node.domain, node.op_type) for node in graph.all_nodes()) + + assert ops["com.microsoft", "MatMulNBits"] == 0 + assert ops["com.microsoft", "GatherBlockQuantized"] == 0 + assert len(graph.initializers) == 34 + assert graph.initializers["model.layers.0.mamba.in_proj.weight"].shape == [194, 64] + assert graph.initializers["model.layers.1.moe.experts.0.up_proj.weight"].shape == [32, 64] + assert graph.initializers["model.layers.2.self_attn.q_proj.weight"].shape == [64, 64] + assert graph.initializers["model.embed_tokens.weight"].shape == [256, 64] + assert graph.initializers["lm_head.weight"].shape == [256, 64] + assert not any(name.endswith((".scales", ".zero_points")) for name in graph.initializers) + + +def test_quantized_nemotron_wires_every_projection_family() -> None: + config = _config( + quantization=QuantizationConfig( + bits=8, + group_size=32, + quant_method="gguf", + sym=False, + quantize_embeddings=True, + quantize_lm_head=True, + ) + ) + graph = _build_graph(config) + ops = Counter((node.domain, node.op_type) for node in graph.all_nodes()) + + assert ops["com.microsoft", "MatMulNBits"] == 13 + assert ops["com.microsoft", "GatherBlockQuantized"] == 1 + assert len(graph.initializers) == 62 + assert graph.initializers["model.layers.0.mamba.in_proj.weight"].shape == [ + 194, + 2, + 32, + ] + assert graph.initializers["model.layers.1.moe.experts.0.up_proj.weight"].shape == [ + 32, + 2, + 32, + ] + assert graph.initializers["model.layers.2.self_attn.q_proj.weight"].shape == [ + 64, + 2, + 32, + ] + assert graph.initializers["model.embed_tokens.qweight"].shape == [256, 64] + assert graph.initializers["lm_head.weight"].shape == [256, 2, 32] + + +def test_gguf_state_dict_requires_exact_initializer_coverage() -> None: + package = _build_package(_config()) + required = { + name: torch.empty(0) + for name, value in package["model"].graph.initializers.items() + if value.const_value is None + } + + validate_package_state_dict(package, required) + missing = dict(required) + missing.pop(next(iter(missing))) + with pytest.raises(ValueError, match="1 missing"): + validate_package_state_dict(package, missing) + with pytest.raises(ValueError, match="1 unexpected"): + validate_package_state_dict(package, {**required, "unexpected.weight": torch.empty(0)})