From 15618c2cde76ee358f1c7ad14e05d6a0a9a5c86b Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 7 Oct 2026 14:11:34 -0700 Subject: [PATCH 01/15] Document pinned Nemotron 3.5 Lightning export preflight Add immutable BF16/NVFP4 metadata and exporter-source recipes with a reproducible metadata-only PowerShell preflight. Document unvalidated BF16 graph-only syntax, explicit ModelOpt rejection, MTP and recurrent-state packaging gaps, and evidence limits without claiming successful exports. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> --- docs/index.md | 1 + docs/nemotron-lightning-export.md | 166 ++++++++++++++++++++++++ examples/nemotron_lightning/recipe.json | 122 +++++++++++++++++ 3 files changed, 289 insertions(+) create mode 100644 docs/nemotron-lightning-export.md create mode 100644 examples/nemotron_lightning/recipe.json diff --git a/docs/index.md b/docs/index.md index aae56e1ff..4b67be6be 100644 --- a/docs/index.md +++ b/docs/index.md @@ -53,4 +53,5 @@ ai-model-support-strategy :caption: Research research/testing-strategy-analysis +nemotron-lightning-export ``` diff --git a/docs/nemotron-lightning-export.md b/docs/nemotron-lightning-export.md new file mode 100644 index 000000000..1bd16f705 --- /dev/null +++ b/docs/nemotron-lightning-export.md @@ -0,0 +1,166 @@ +# Nemotron 3.5 Lightning: pinned export preflight + +**Metadata/source preflight only (2026-10-07); both full exports `NOT_RUN`.** +The [machine-readable recipe](../examples/nemotron_lightning/recipe.json) +records immutable inputs and compatibility limits, not a tested export or +validated toolchain. It adds no model implementation or NVFP4 support. + +The original preflight host was Windows ARM64 with a Qualcomm Adreno X1-85, +no NVIDIA/CUDA GPU or `nvidia-smi`, Python 3.12.10, 68,171,038,720 bytes RAM +and approximately 818 GB free disk. Mobius, PyTorch, Transformers, ORT and +ORT-GenAI were absent. Only public metadata and exporter sources were +retrieved: no weight shards, exporter installation, remote model code +execution, GPU work, ONNX export or generation. + +## Immutable inputs + +| Input | BF16 | NVFP4 | +| --- | --- | --- | +| NVIDIA model | `nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16` | `nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-NVFP4` | +| HF revision | `a9904d24bcc1d289a1950fa9d2b978c47cf903b9` | `bee7596271d1495f6992ae224aefde4410e816b8` | +| Index tensor bytes | 65,842,365,568 | 21,559,589,596 | +| Unique index shards | 14 | 52 | +| Config context limit | 262,144 | 1,048,576 | +| Source quantization | Unquantized BF16 | ModelOpt mixed FP8/NVFP4 | +| Classification | Architecture source present; unvalidated | Explicitly unsupported by pinned exporter | + +Both configs declare `NemotronHForCausalLM` / `nemotron_h`, 52 layers +(23 Mamba, 23 MoE, 6 attention), 128 routed experts, top-6 routing, and one +NextN prediction layer. The 3B active-parameter label is **not** a +weight-residency requirement. Index sizes are metadata declarations, not +measured download sizes, output sizes, peak RAM or VRAM. Host RAM barely +exceeds BF16 tensor bytes before loader/graph/OS overhead; this does not +establish that a build fits. + +The BF16 model card advertises up to 1M context, but its pinned config says +262,144. Do not silently change it. The recipe pins raw-byte SHA-256 hashes +for both configs/indexes and six inspected Mobius files at +`dc1798465f77cb2b7070e4d422479c2f890ee24b`. This is a source-evidence pin, +**not** an approved installation/dependency lock. Compatibility claims were +rechecked at `26a712b210a0a7cd895347f601680e82bbf5e452`: five files are +unchanged; `auto_export.py` differs only in unrelated Gemma4 audio settings. +The recorded source hashes remain those of the original inspected revision. + +## Reproduce metadata/source verification + +Run from the Mobius repository root in PowerShell. Set `$Evidence` to a +new private directory outside the checkout. This downloads only four JSON +metadata files and six exporter source files, never checkpoint shards. +It does not import/install Mobius or execute the downloaded source. +HTTP errors, hash mismatches and changed metadata stop the preflight. +Slashes in URL components and recipe source paths are not local paths. + +```powershell +$ErrorActionPreference = 'Stop' +$Evidence = 'C:\private\nemotron-lightning-preflight' +if (Test-Path $Evidence) { throw 'Use a new evidence directory' } +New-Item -ItemType Directory -Path $Evidence | Out-Null +$Recipe = Get-Content .\examples\nemotron_lightning\recipe.json -Raw | ConvertFrom-Json +foreach ($Model in $Recipe.models) { + $Variant = ($Model.model_id -split '-')[-1] + foreach ($File in $Model.metadata) { + $Dest = Join-Path $Evidence "$Variant-$($File.path)" + $Url = "https://huggingface.co/$($Model.model_id)/resolve/$($Model.revision)/$($File.path)" + Invoke-WebRequest -Uri $Url -OutFile $Dest + if ((Get-FileHash $Dest -Algorithm SHA256).Hash.ToLowerInvariant() -ne $File.sha256) { + throw "Metadata hash mismatch: $Variant $($File.path)" + } + } + $Config = Get-Content (Join-Path $Evidence "$Variant-config.json") -Raw | ConvertFrom-Json + $Index = Get-Content (Join-Path $Evidence "$Variant-model.safetensors.index.json") -Raw | ConvertFrom-Json + $Shards = @($Index.weight_map.PSObject.Properties.Value | Sort-Object -Unique).Count + if ($Index.metadata.total_size -ne $Model.checkpoint_tensor_bytes -or $Shards -ne $Model.checkpoint_shards) { + throw "Checkpoint inventory mismatch: $Variant" + } + if ($Config.model_type -ne $Model.model_type -or + $Config.architectures.Count -ne 1 -or $Config.architectures[0] -ne $Model.architecture -or + $Config.num_hidden_layers -ne $Model.num_hidden_layers -or + $Config.max_position_embeddings -ne $Model.max_position_embeddings -or + $Config.num_nextn_predict_layers -ne $Model.num_nextn_predict_layers -or + $Config.n_routed_experts -ne $Model.num_routed_experts -or + $Config.num_experts_per_tok -ne $Model.num_experts_per_tok) { + throw "Architecture/context mismatch: $Variant" + } + foreach ($Layer in $Model.layer_counts.PSObject.Properties) { + $Count = @($Config.layers_block_type | Where-Object { $_ -eq $Layer.Name }).Count + if ($Count -ne $Layer.Value) { throw "Layer inventory mismatch: $Variant $($Layer.Name)" } + } + "$Variant : metadata verified; export NOT_RUN; $($Model.compatibility)" +} +foreach ($File in $Recipe.exporter.source_evidence) { + $Dest = Join-Path $Evidence ($File.path.Replace('/', '_')) + $Url = "https://raw.githubusercontent.com/onnxruntime/mobius/$($Recipe.exporter.revision)/$($File.path)" + Invoke-WebRequest -Uri $Url -OutFile $Dest + if ((Get-FileHash $Dest -Algorithm SHA256).Hash.ToLowerInvariant() -ne $File.sha256) { + throw "Exporter source hash mismatch: $($File.path)" + } +} +'All ten pinned files verified; exports NOT_RUN' +``` + +This verifies inspected bytes and metadata, **not** executable compatibility. +Config/index hashes do not authenticate every shard or tokenizer asset. + +## Source compatibility boundaries + +These links identify the original inspected revision; the recheck above +does not change the scope of evidence. + +| Source | Finding and limit | +| --- | --- | +| [`_registry.py`](https://github.com/onnxruntime/mobius/blob/dc1798465f77cb2b7070e4d422479c2f890ee24b/src/mobius/_registry.py) and [`_configs/_base.py`](https://github.com/onnxruntime/mobius/blob/dc1798465f77cb2b7070e4d422479c2f890ee24b/src/mobius/_configs/_base.py) | Nemotron-H registration and `NemotronHConfig.from_transformers` translate Mamba/attention/MoE layer names. This is not Lightning checkpoint parity. | +| [`models/nemotron_h.py`](https://github.com/onnxruntime/mobius/blob/dc1798465f77cb2b7070e4d422479c2f890ee24b/src/mobius/models/nemotron_h.py) | Hybrid layers, weight renaming and stacked-expert splitting exist, but no MTP implementation in this module. MoE evaluates all experts and masks contributions, not demonstrated sparse top-6 efficiency. Its source documents ReLU2 as incompatible with the described ORT 1.29 MoE/QMoE activation set; substituting ReLU is not equivalent. | +| [`_configs/_quantization.py`](https://github.com/onnxruntime/mobius/blob/dc1798465f77cb2b7070e4d422479c2f890ee24b/src/mobius/_configs/_quantization.py) | `QuantizationConfig.from_value` explicitly raises `NotImplementedError` for ModelOpt NVFP4/FP8. Reconstruction helpers do not wire checkpoint loading or native routed-expert NVFP4 emission. | +| [`integrations/ort_genai/auto_export.py`](https://github.com/onnxruntime/mobius/blob/dc1798465f77cb2b7070e4d422479c2f890ee24b/src/mobius/integrations/ort_genai/auto_export.py) | Generic recurrent packaging expects paired `conv_state` / `recurrent_state`; Nemotron-H describes `conv_state` / `ssm_state`. Actual graph names/shapes and runtime allocation are unverified, not evidence of a loadable package. | +| [`__main__.py`](https://github.com/onnxruntime/mobius/blob/dc1798465f77cb2b7070e4d422479c2f890ee24b/src/mobius/__main__.py) | The candidate below uses source-backed CLI syntax. `--revision` is only valid with `--model`, not local `--config`. | + +## BF16 candidate procedure: NOT_RUN + +Before any full export, obtain a suitable development host, explicit resource +budget, immutable dependency lock and CUDA/ORT version pairing. No package +installation, model download or GPU allocation is part of this preflight. +The publisher's 80 GB-class GPU deployment target is not an export-memory +measurement. Stage a complete read-only snapshot at the pinned HF revision: +config, tokenizer and every indexed weight shard. Verify shard identities, +tokenizer assets, disk capacity and loader peak memory separately. Do not +enable `trust_remote_code` or fetch mutable `main`. + +On that separately prepared host, the candidate **graph-only** command is: + +```powershell +mobius build --config C:\private\models\nemotron-lightning-bf16 ` + --output C:\private\exports\nemotron-lightning-bf16 ` + --ep cuda --dtype bf16 --external-data onnx +``` + +Use a fresh output directory; retain stdout, stderr, exit status, exact tool +versions and resource measurements. This command has **not** been executed +and may fail on Lightning-specific weights/semantics. It intentionally +omits `--runtime ort-genai`: an exported graph is not a loadable runtime +bundle. No FP16 conversion, MTP removal or generic decoder fallback is +implied by this recipe. + +Before claiming success, account for every checkpoint tensor, including +MTP; validate graph/external-data references; compare source/export +full logits and deterministic multi-step cached decoding. A target-only +variant that omits MTP needs explicit scope, a separate name and validation, +not a full-checkpoint/MTP support claim. Verify Mamba convolution/SSM +state and attention cache initialization, update and reset, plus actual +CUDA node placement. Only afterward evaluate downstream ORT-GenAI +packaging, tokenizer/chat-template fidelity and generation. Record memory, +quality, generation settings and tested context limits, not model-card +throughput claims. + +## NVFP4: stop before weights + +The pinned ModelOpt config describes mixed FP8 Mamba projections and FP4 +expert weights with group size 16. The exporter rejects this method: +**there is no supported NVFP4 build command in this recipe**. Do not +download 52 shards merely to reproduce that known source-level blocker. + +Supporting it requires mixed-format/scales-aware loading, faithful ReLU2 +expert execution, runtime kernel/build capability checks and real-weight +parity. Blackwell native FP4 and W4A16 reconstruction are distinct paths; +NVIDIA's vLLM support does not establish ORT support. BF16 reconstruction +or ORT INT4 / `MatMulNBits` quantization creates a different artifact and +must not be labeled a preserved NVFP4 export. diff --git a/examples/nemotron_lightning/recipe.json b/examples/nemotron_lightning/recipe.json new file mode 100644 index 000000000..a6f0eb159 --- /dev/null +++ b/examples/nemotron_lightning/recipe.json @@ -0,0 +1,122 @@ +{ + "schema_version": 1, + "kind": "development_export_recipe", + "checked_on": "2026-10-07", + "status": "PREFLIGHT_ONLY", + "export_status": "NOT_RUN", + "certified": false, + "exporter": { + "repository": "https://github.com/onnxruntime/mobius", + "revision": "dc1798465f77cb2b7070e4d422479c2f890ee24b", + "compatibility_rechecked_at": "26a712b210a0a7cd895347f601680e82bbf5e452", + "package": "mobius-onnx", + "toolchain_validated": false, + "source_evidence": [ + { + "path": "src/mobius/__main__.py", + "sha256": "dccc246cc17992e810820af7e57572a8f3a73fbefc845d5bfe861e4ce80875f5" + }, + { + "path": "src/mobius/_registry.py", + "sha256": "ea502bf2fdb7cd4b73ab5c75d4c8769918f312712112fd2f011652277ac45f37" + }, + { + "path": "src/mobius/_configs/_base.py", + "sha256": "149516ae2be475e3dba51a30a03b16d95d31799391eff2f97ca2666599d209d6" + }, + { + "path": "src/mobius/_configs/_quantization.py", + "sha256": "d4bc4b0ac10c01d1058fb7d45750a8a5f6ea9b73c30319459f1319838e337d27" + }, + { + "path": "src/mobius/models/nemotron_h.py", + "sha256": "75b38c8e52bc7a61f0984b134902b6ed7e9cd779879a62bd8f2879db668d31fe" + }, + { + "path": "src/mobius/integrations/ort_genai/auto_export.py", + "sha256": "f44685481cd6ae3f36a74e62556ca3ab50a3a43f83e6bc8bfc6257a97a44c9c3" + } + ] + }, + "models": [ + { + "model_id": "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16", + "revision": "a9904d24bcc1d289a1950fa9d2b978c47cf903b9", + "metadata": [ + { + "path": "config.json", + "sha256": "a3827a0f5e311547b40943dc081e3ff2f8a277466e8c1a3df2291e8db8a7617c" + }, + { + "path": "model.safetensors.index.json", + "sha256": "67f21da80ce245a3e24967f54eef3fa10e67a63eba16ef19a2d0569f06103f50" + } + ], + "architecture": "NemotronHForCausalLM", + "model_type": "nemotron_h", + "layer_counts": {"mamba": 23, "moe": 23, "attention": 6}, + "num_hidden_layers": 52, + "num_routed_experts": 128, + "num_experts_per_tok": 6, + "num_nextn_predict_layers": 1, + "max_position_embeddings": 262144, + "checkpoint_tensor_bytes": 65842365568, + "checkpoint_shards": 14, + "compatibility": "ARCHITECTURE_SOURCE_PRESENT_NOT_VALIDATED", + "output_precision": "bf16", + "proposed_graph_export_argv": [ + "mobius", "build", "--config", "", + "--output", "", + "--ep", "cuda", "--dtype", "bf16", "--external-data", "onnx" + ], + "command_status": "NOT_RUN", + "runtime_packaging_argv": null, + "blockers": [ + "No NVIDIA CUDA GPU or installed exporter on the preflight host.", + "A suitable host, dependency lock and resource budget have not been validated.", + "Full Lightning checkpoint weight coverage and source/export parity are untested.", + "Nemotron-H MTP export is not implemented in the inspected model module.", + "Mamba conv_state/ssm_state packaging is not verified against ORT-GenAI." + ] + }, + { + "model_id": "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-NVFP4", + "revision": "bee7596271d1495f6992ae224aefde4410e816b8", + "metadata": [ + { + "path": "config.json", + "sha256": "f1d98b530846087dc08b574a219713a94f945bf6583dc7230a19ebf1e8c50933" + }, + { + "path": "model.safetensors.index.json", + "sha256": "3c3bc7efa8d658c2e909a0b9020eb0f72064e6647de348856af4dee9895bead9" + } + ], + "architecture": "NemotronHForCausalLM", + "model_type": "nemotron_h", + "layer_counts": {"mamba": 23, "moe": 23, "attention": 6}, + "num_hidden_layers": 52, + "num_routed_experts": 128, + "num_experts_per_tok": 6, + "num_nextn_predict_layers": 1, + "max_position_embeddings": 1048576, + "checkpoint_tensor_bytes": 21559589596, + "checkpoint_shards": 52, + "quant_method": "modelopt", + "weight_formats": ["FP8", "NVFP4"], + "nvfp4_group_size": 16, + "compatibility": "UNSUPPORTED_BY_PINNED_EXPORTER", + "output_precision": null, + "proposed_graph_export_argv": null, + "command_status": "NOT_RUN", + "runtime_packaging_argv": null, + "blockers": [ + "QuantizationConfig.from_value explicitly raises NotImplementedError for ModelOpt NVFP4/FP8.", + "ModelOpt checkpoint loading and native routed-expert NVFP4 emission are not wired.", + "ReLU2 expert semantics require separate validation; a generic QMoE is not equivalent.", + "BF16 reconstruction or ORT INT4 quantization would be different artifacts, not preserved NVFP4.", + "The BF16 host, checkpoint parity, MTP and runtime packaging blockers also apply." + ] + } + ] +} From e4560d1b9ca47b25632963891c74077396c4da26 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 7 Oct 2026 16:16:45 -0700 Subject: [PATCH 02/15] Add strict Nemotron ModelOpt BF16 reconstruction and compatibility gates Implement explicit, scale-aware mixed FP8/NVFP4 reconstruction with bounded-row temporaries, exact per-expert source accounting, lazy header revalidation and truthful provenance. Validate activation and KV-cache scales without claiming their quantization is preserved. Keep native formats and unsupported NextN/MTP contracts blocked; no tensors are silently discarded. Correct FP32 router/accumulation and preserve the declared SSM cache dtype. Add offline numeric, public-builder/CLI, serialization, source-accounting and cache regressions, and update the pinned recipes to identify the remaining full-Lightning MTP gate. Full checkpoints, CUDA/native kernels, generation and runtime packaging are not claimed. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> --- docs/nemotron-lightning-export.md | 143 ++++++-- examples/nemotron_lightning/recipe.json | 43 ++- src/mobius/_configs/_base.py | 13 + src/mobius/_configs/_quantization.py | 17 +- src/mobius/integrations/_weight_loading.py | 89 ++++- src/mobius/integrations/modelopt/__init__.py | 10 +- src/mobius/integrations/modelopt/_config.py | 163 +++++++++ src/mobius/integrations/modelopt/_dequant.py | 16 +- src/mobius/integrations/modelopt/_weights.py | 75 +++++ .../integrations/modelopt/_weights_test.py | 258 +++++++++++++++ .../integrations/transformers/_builder.py | 75 ++++- .../transformers/_nemotron_h_weights.py | 113 +++++++ .../transformers/_nemotron_h_weights_test.py | 308 ++++++++++++++++++ src/mobius/models/_nemotron_h_test.py | 240 ++++++++++++++ src/mobius/models/nemotron_h.py | 34 +- src/mobius/tasks/_cache_utils.py | 2 +- 16 files changed, 1519 insertions(+), 80 deletions(-) create mode 100644 src/mobius/integrations/modelopt/_config.py create mode 100644 src/mobius/integrations/modelopt/_weights.py create mode 100644 src/mobius/integrations/modelopt/_weights_test.py create mode 100644 src/mobius/integrations/transformers/_nemotron_h_weights.py create mode 100644 src/mobius/integrations/transformers/_nemotron_h_weights_test.py create mode 100644 src/mobius/models/_nemotron_h_test.py diff --git a/docs/nemotron-lightning-export.md b/docs/nemotron-lightning-export.md index 1bd16f705..e4aa7be4f 100644 --- a/docs/nemotron-lightning-export.md +++ b/docs/nemotron-lightning-export.md @@ -1,16 +1,21 @@ -# Nemotron 3.5 Lightning: pinned export preflight +# Nemotron 3.5 Lightning: reconstruction and compatibility gates -**Metadata/source preflight only (2026-10-07); both full exports `NOT_RUN`.** +**CPU-tested exporter improvements; both full Lightning exports `NOT_RUN`.** The [machine-readable recipe](../examples/nemotron_lightning/recipe.json) -records immutable inputs and compatibility limits, not a tested export or -validated toolchain. It adds no model implementation or NVFP4 support. +records immutable inputs and compatibility limits. Native NVFP4 remains +unsupported. Explicit dense BF16 reconstruction is implemented and tested +with miniature mixed-format checkpoints, **not** the complete 30B models. +The pinned checkpoints' NextN/MTP contract remains a blocking gate. The original preflight host was Windows ARM64 with a Qualcomm Adreno X1-85, no NVIDIA/CUDA GPU or `nvidia-smi`, Python 3.12.10, 68,171,038,720 bytes RAM and approximately 818 GB free disk. Mobius, PyTorch, Transformers, ORT and ORT-GenAI were absent. Only public metadata and exporter sources were retrieved: no weight shards, exporter installation, remote model code -execution, GPU work, ONNX export or generation. +execution, GPU work, ONNX export or generation during that initial preflight. +Subsequent code validation used an isolated x64 CPU Python environment on +the same ARM64 machine. Only metadata and bounded safetensors headers were +retrieved from the real checkpoints; no tensor payloads were downloaded. ## Immutable inputs @@ -22,7 +27,7 @@ execution, GPU work, ONNX export or generation. | Unique index shards | 14 | 52 | | Config context limit | 262,144 | 1,048,576 | | Source quantization | Unquantized BF16 | ModelOpt mixed FP8/NVFP4 | -| Classification | Architecture source present; unvalidated | Explicitly unsupported by pinned exporter | +| Current full-checkpoint classification | Blocked on NextN/MTP | Native unsupported; dense reconstruction blocked on NextN/MTP | Both configs declare `NemotronHForCausalLM` / `nemotron_h`, 52 layers (23 Mamba, 23 MoE, 6 attention), 128 routed experts, top-6 routing, and one @@ -40,6 +45,64 @@ for both configs/indexes and six inspected Mobius files at rechecked at `26a712b210a0a7cd895347f601680e82bbf5e452`: five files are unchanged; `auto_export.py` differs only in unrelated Gemma4 audio settings. The recorded source hashes remain those of the original inspected revision. +The historical source table below does not describe the new reconstruction +implementation. + +## Implemented exporter behavior + +The current code adds an explicit Nemotron-H dense reconstruction route: +`build(..., keep_quantized=False, dtype="bf16")`, corresponding to +`--dequantize --dtype bf16`. Both options are required; implicit dtype, +FP16, native requests and unsupported layouts fail before weight loading. +This is **not a working full-Lightning command while the MTP gate remains**. + +- `integrations/modelopt/_config.py` validates exact FP8/NVFP4 module groups, + static block-16 format, redundant `quantized_layers` declarations and FP8 + activation/KV-cache schemes. Unknown formats and inconsistent inventories + fail rather than falling back to affine integer INT4. +- `integrations/modelopt/_weights.py` reuses the numeric helpers to reconstruct + one projection at a time, with 256-row temporary buffers and one final BF16 + rounding. It validates physical/logical shapes, scale types/layouts, positive + finite scalar scales, nonnegative finite block scales and nonfinite results. + Zero block scales reconstruct zero blocks. Lazy loading rechecks + shard headers before materialization. +- `integrations/transformers/_nemotron_h_weights.py` maps per-expert sources + exactly, rejecting collisions, extra/unclassified tensors, missing weights + and missing scales. FP8 activation and KV-cache scales are explicitly + validated/accounted for, **not simulated**. Export metadata and the loading + report identify the reconstructed representation and lost quantization. +- `models/nemotron_h.py` casts router weights/bias to FP32 and accumulates routed + outputs in FP32 before the shared-expert path, retaining sigmoid selection, + unbiased routing weights and squared ReLU. All-expert graph unrolling remains; + this is not sparse top-6 performance or native QMoE support. +- `NemotronHConfig` accepts Transformers' `linear_attention` vocabulary and + preserves the checkpoint's FP32 SSM-cache dtype. The graph still uses + `conv_state`/`ssm_state`; no ORT-GenAI recurrent ABI conversion is added. + +The complete pinned ModelOpt config declares **46 FP8 projections, 5,935 +NVFP4 projections, 46 FP8 input scales and 12 FP8 KV-cache scales**. A bounded +header probe verifies KV-cache scales are FP32 `[1]`; NVFP4 projection global +scales are FP32 scalars and block scales are E4M3 `[N,K/16]`. BF16's index has +6,243 non-MTP tensor names mapping uniquely and **270 `mtp.*` tensors**; +the NVFP4 index has the same MTP inventory. Name accounting and selected +headers do not validate all tensor values or demonstrate full-model parity. + +**Dense BF16 reconstruction removes native FP8/NVFP4 storage, FP8 activation +quantization and FP8 KV-cache quantization.** It is neither native NVFP4 +preservation nor re-quantized ORT INT4, and does not promise equivalence to +the original quantized execution. + +### NextN/MTP: fail closed + +Model construction rejects nonzero `num_nextn_predict_layers`; eager and +streaming weight paths also reject `mtp.*` tensors even if a config omits that +field. No automatic target-only fallback or discarded auxiliary tensors are +allowed. The [inspected HuggingFace reference](https://github.com/huggingface/transformers/blob/a96730c8c97b8efbf35bbaf7f5da33ec99231a49/src/transformers/models/nemotron_h/modeling_nemotron_h.py) +ignores MTP; the [inspected vLLM MTP implementation](https://github.com/vllm-project/vllm/blob/ef63c23d35acccbba8e014fc88da19d541da38a9/vllm/model_executor/models/nemotron_h_mtp.py) +appears to construct shared experts absent from the pinned MTP indexes. +That discrepancy is not resolved by inventing missing weights. +A separately authorized, explicitly named target-only variant or a validated +faithful MTP graph is required before these complete checkpoints can build. ## Reproduce metadata/source verification @@ -114,31 +177,26 @@ does not change the scope of evidence. | [`integrations/ort_genai/auto_export.py`](https://github.com/onnxruntime/mobius/blob/dc1798465f77cb2b7070e4d422479c2f890ee24b/src/mobius/integrations/ort_genai/auto_export.py) | Generic recurrent packaging expects paired `conv_state` / `recurrent_state`; Nemotron-H describes `conv_state` / `ssm_state`. Actual graph names/shapes and runtime allocation are unverified, not evidence of a loadable package. | | [`__main__.py`](https://github.com/onnxruntime/mobius/blob/dc1798465f77cb2b7070e4d422479c2f890ee24b/src/mobius/__main__.py) | The candidate below uses source-backed CLI syntax. `--revision` is only valid with `--model`, not local `--config`. | -## BF16 candidate procedure: NOT_RUN +## Full BF16 procedure: blocked, NOT_RUN Before any full export, obtain a suitable development host, explicit resource budget, immutable dependency lock and CUDA/ORT version pairing. No package installation, model download or GPU allocation is part of this preflight. The publisher's 80 GB-class GPU deployment target is not an export-memory -measurement. Stage a complete read-only snapshot at the pinned HF revision: +measurement. Only after the MTP gate is resolved, stage a complete read-only +snapshot at the pinned HF revision: config, tokenizer and every indexed weight shard. Verify shard identities, tokenizer assets, disk capacity and loader peak memory separately. Do not enable `trust_remote_code` or fetch mutable `main`. -On that separately prepared host, the candidate **graph-only** command is: - -```powershell -mobius build --config C:\private\models\nemotron-lightning-bf16 ` - --output C:\private\exports\nemotron-lightning-bf16 ` - --ep cuda --dtype bf16 --external-data onnx -``` - -Use a fresh output directory; retain stdout, stderr, exit status, exact tool -versions and resource measurements. This command has **not** been executed -and may fail on Lightning-specific weights/semantics. It intentionally -omits `--runtime ort-genai`: an exported graph is not a loadable runtime -bundle. No FP16 conversion, MTP removal or generic decoder fallback is -implied by this recipe. +The earlier graph-only candidate is no longer presented as runnable: +the pinned config declares one NextN layer, and current construction explicitly +rejects that unsupported contract before loading shards. The recipe therefore +contains no approved full-checkpoint build command. Do not edit the config to +zero NextN or remove MTP tensors to bypass this gate. Any future command must +use a fresh output directory and retain logs, exit status, exact versions and +resource measurements; it must omit `--runtime ort-genai` until that separate +runtime contract is validated. Before claiming success, account for every checkpoint tensor, including MTP; validate graph/external-data references; compare source/export @@ -151,16 +209,41 @@ packaging, tokenizer/chat-template fidelity and generation. Record memory, quality, generation settings and tested context limits, not model-card throughput claims. -## NVFP4: stop before weights +## NVFP4: native blocked; dense reconstruction is distinct The pinned ModelOpt config describes mixed FP8 Mamba projections and FP4 -expert weights with group size 16. The exporter rejects this method: -**there is no supported NVFP4 build command in this recipe**. Do not -download 52 shards merely to reproduce that known source-level blocker. - -Supporting it requires mixed-format/scales-aware loading, faithful ReLU2 -expert execution, runtime kernel/build capability checks and real-weight -parity. Blackwell native FP4 and W4A16 reconstruction are distinct paths; +expert weights with group size 16, plus static FP8 activations and KV cache. +The generic/native exporter still rejects this method. The opt-in dense +reconstruction route above handles these physical formats but remains blocked +on full Lightning MTP accounting. **There is no supported native NVFP4 or +full-Lightning build command in this recipe.** Do not download 52 shards merely +to reproduce that known contract blocker. + +Native preservation additionally requires faithful ReLU2/routing kernel +semantics, runtime capability checks and real-weight parity. +Blackwell native FP4 and dense BF16 reconstruction are distinct paths; NVIDIA's vLLM support does not establish ORT support. BF16 reconstruction or ORT INT4 / `MatMulNBits` quantization creates a different artifact and must not be labeled a preserved NVFP4 export. + +## Offline regression validation + +The added suites use reduced configurations, independently specified routing +and E2M1/E4M3 values, and local synthetic safetensors files. They exercise public +builder/CLI reconstruction, save/reload, every layer/projection family, source +accounting, malformed/changing headers, dtype/cache contracts and MTP refusal. +Existing CPU random-weight HuggingFace Nemotron parity tests also run without +model downloads. These are not full-checkpoint L4/L5 evidence. + +```powershell +python -m pytest src\mobius\integrations\modelopt ` + src\mobius\models\_nemotron_h_test.py ` + src\mobius\integrations\transformers\_nemotron_h_weights_test.py -q --tb=short +python -m pytest tests\synthetic_parity_test.py -k nemotron_h -q -ra --tb=short +``` + +The validated isolated CPU versions were PyTorch `2.14.1+cpu` (no CUDA build), +Transformers `5.18.0`, ONNXScript `0.7.2`, ONNX IR `1.0.0` and ORT `1.30.0`. +These are execution evidence, not a recommended GPU dependency lock. +Full 30B export/parity/generation, BF16 CUDA kernels, native NVFP4, +MTP execution, resource fit and ORT-GenAI packaging remain unrun/unvalidated. diff --git a/examples/nemotron_lightning/recipe.json b/examples/nemotron_lightning/recipe.json index a6f0eb159..a6acf245d 100644 --- a/examples/nemotron_lightning/recipe.json +++ b/examples/nemotron_lightning/recipe.json @@ -2,7 +2,7 @@ "schema_version": 1, "kind": "development_export_recipe", "checked_on": "2026-10-07", - "status": "PREFLIGHT_ONLY", + "status": "CPU_REGRESSION_TESTED_FULL_CHECKPOINT_BLOCKED", "export_status": "NOT_RUN", "certified": false, "exporter": { @@ -62,20 +62,18 @@ "max_position_embeddings": 262144, "checkpoint_tensor_bytes": 65842365568, "checkpoint_shards": 14, - "compatibility": "ARCHITECTURE_SOURCE_PRESENT_NOT_VALIDATED", + "compatibility": "NEXTN_CONTRACT_BLOCKED", "output_precision": "bf16", - "proposed_graph_export_argv": [ - "mobius", "build", "--config", "", - "--output", "", - "--ep", "cuda", "--dtype", "bf16", "--external-data", "onnx" - ], - "command_status": "NOT_RUN", + "checkpoint_mtp_tensors": 270, + "checkpoint_non_mtp_tensors": 6243, + "proposed_graph_export_argv": null, + "command_status": "BLOCKED_MTP", "runtime_packaging_argv": null, "blockers": [ - "No NVIDIA CUDA GPU or installed exporter on the preflight host.", + "No NVIDIA CUDA GPU; executed validation uses an isolated x64 CPU environment.", "A suitable host, dependency lock and resource budget have not been validated.", "Full Lightning checkpoint weight coverage and source/export parity are untested.", - "Nemotron-H MTP export is not implemented in the inspected model module.", + "NextN/MTP construction and mtp.* tensors explicitly fail closed; no target-only fallback.", "Mamba conv_state/ssm_state packaging is not verified against ORT-GenAI." ] }, @@ -105,15 +103,30 @@ "quant_method": "modelopt", "weight_formats": ["FP8", "NVFP4"], "nvfp4_group_size": 16, - "compatibility": "UNSUPPORTED_BY_PINNED_EXPORTER", + "compatibility": "NATIVE_UNSUPPORTED_DENSE_RECONSTRUCTION_NEXTN_BLOCKED", "output_precision": null, "proposed_graph_export_argv": null, - "command_status": "NOT_RUN", + "command_status": "BLOCKED_MTP_AND_NATIVE_FORMAT", + "checkpoint_mtp_tensors": 270, + "source_format_counts": {"FP8": 46, "NVFP4": 5935}, + "source_fp8_input_scales": 46, + "source_fp8_kv_cache_scales": 12, + "dense_reconstruction": { + "representation": "dense_bf16", + "api_options": {"keep_quantized": false, "dtype": "bf16"}, + "cli_options": ["--dequantize", "--dtype", "bf16"], + "validation": "SYNTHETIC_CPU_BUILDER_CLI_SERIALIZATION_NUMERIC_TESTS", + "full_checkpoint_export": "NOT_RUN", + "native_weight_format_preserved": false, + "activation_quantization_preserved": false, + "kv_cache_quantization_preserved": false, + "requires_validated_nextn_contract": true + }, "runtime_packaging_argv": null, "blockers": [ - "QuantizationConfig.from_value explicitly raises NotImplementedError for ModelOpt NVFP4/FP8.", - "ModelOpt checkpoint loading and native routed-expert NVFP4 emission are not wired.", - "ReLU2 expert semantics require separate validation; a generic QMoE is not equivalent.", + "Generic/native ModelOpt NVFP4/FP8 remains unsupported; explicit reconstruction is separately dispatched.", + "Mixed-format loading/reconstruction is CPU-tested, not native routed-expert NVFP4 emission.", + "Portable ReLU2/sigmoid routing is CPU-tested; a generic QMoE is not equivalent.", "BF16 reconstruction or ORT INT4 quantization would be different artifacts, not preserved NVFP4.", "The BF16 host, checkpoint parity, MTP and runtime packaging blockers also apply." ] diff --git a/src/mobius/_configs/_base.py b/src/mobius/_configs/_base.py index 61a7e15ee..acebe31aa 100644 --- a/src/mobius/_configs/_base.py +++ b/src/mobius/_configs/_base.py @@ -4232,10 +4232,21 @@ class NemotronHConfig(ArchitectureConfig): mamba_proj_bias: bool = False mamba_time_step_min: float = 0.001 moe_latent_size: int | None = None + mamba_ssm_dtype: ir.DataType | None = None @classmethod def from_transformers(cls, config, parent_config=None) -> NemotronHConfig: base = ArchitectureConfig.from_transformers(config, parent_config) + if getattr(config, "n_group", 1) != 1 or getattr(config, "topk_group", 1) != 1: + raise NotImplementedError("Nemotron-H grouped expert selection is not implemented") + cache_dtype = getattr(config, "mamba_ssm_cache_dtype", None) + ssm_dtype = _resolve_dtype_value(cache_dtype) + if cache_dtype is not None and ssm_dtype not in { + ir.DataType.FLOAT, + ir.DataType.FLOAT16, + ir.DataType.BFLOAT16, + }: + raise ValueError(f"Unsupported Nemotron-H SSM cache dtype {cache_dtype!r}") # Get layer types from layers_block_type or hybrid_override_pattern layers_block_type = getattr(config, "layers_block_type", None) @@ -4253,6 +4264,7 @@ def from_transformers(cls, config, parent_config=None) -> NemotronHConfig: # Convert HF names to mobius names type_map = { "mamba": "mamba2", + "linear_attention": "mamba2", "attention": "full_attention", "moe": "moe", } @@ -4305,6 +4317,7 @@ def from_transformers(cls, config, parent_config=None) -> NemotronHConfig: mamba_conv_bias=getattr(config, "use_conv_bias", True), mamba_proj_bias=getattr(config, "mamba_proj_bias", False), mamba_time_step_min=getattr(config, "time_step_min", 0.001), + mamba_ssm_dtype=ssm_dtype, moe_latent_size=getattr(config, "moe_latent_size", None), shared_expert_intermediate_size=shared_expert_intermediate_size, ) diff --git a/src/mobius/_configs/_quantization.py b/src/mobius/_configs/_quantization.py index defb49ffc..961fb4ec0 100644 --- a/src/mobius/_configs/_quantization.py +++ b/src/mobius/_configs/_quantization.py @@ -149,20 +149,19 @@ def from_value( # silently mis-dequantize. Checked before the ``quant_method == "none"`` # early-return because ModelOpt may name its scheme only via # ``quant_algo``/``quant_cfg`` (no ``quant_method``). Fail loudly: the - # reconstruction math is available in :mod:`mobius.integrations.modelopt`, - # but the full weight-load integration + native routed-expert NVFP4 QMoE - # emission (CUDA/Blackwell, ``onnxruntime_USE_FP4_QMOE=ON``) are not wired. + # Nemotron-H's explicit BF16 reconstruction is dispatched separately by + # its checkpoint builder. This generic/native INT4 path must stay blocked. from mobius.integrations.modelopt import is_modelopt_quant_config if is_modelopt_quant_config(qc): raise NotImplementedError( "NVIDIA ModelOpt NVFP4/FP8 checkpoints are not yet fully " - "supported for export. The weight-reconstruction core lives in " - "mobius.integrations.modelopt (dequantize_nvfp4 / dequantize_fp8); " - "remaining work is checkpoint weight-loading and native " - "routed-expert NVFP4 QMoE emission (CUDA/Blackwell, " - "onnxruntime_USE_FP4_QMOE=ON). Export the unquantized (bf16) " - "checkpoint instead, or quantize the bf16 export via Olive." + "supported by the generic/native quantized exporter. Nemotron-H " + "supports explicit dense BF16 weight reconstruction through " + "build(..., keep_quantized=False, dtype='bf16') / " + "--dequantize --dtype bf16, subject to its NextN/MTP contract. " + "This does not preserve native NVFP4/FP8 execution. Other model " + "architectures remain unsupported." ) # compressed-tensors is a container format, not an INT4 method. Its # ``bits``/``group_size`` values live inside ordered config groups and diff --git a/src/mobius/integrations/_weight_loading.py b/src/mobius/integrations/_weight_loading.py index 9b3372eb6..acac583ee 100644 --- a/src/mobius/integrations/_weight_loading.py +++ b/src/mobius/integrations/_weight_loading.py @@ -98,9 +98,12 @@ class StreamingWeightSource: """One checkpoint tensor bound to one dense ONNX initializer.""" source_name: str - mode: Literal["direct", "fp8_scalar", "fp8_block_128"] = "direct" + mode: Literal[ + "direct", "fp8_scalar", "fp8_block_128", "modelopt_fp8", "modelopt_nvfp4" + ] = "direct" scale_name: str | None = None expected_scale: float | None = None + global_scale_name: str | None = None @dataclasses.dataclass(frozen=True) @@ -690,6 +693,20 @@ def _materialize_preprocessed_source( """Read and reconstruct one classified source tensor.""" source_dtype = key_index[source.source_name][2] tensor = _load_indexed_tensor(source.source_name, key_index) + if source.mode in {"modelopt_fp8", "modelopt_nvfp4"}: + from mobius.integrations.modelopt._weights import reconstruct_modelopt_weight + + assert source.scale_name is not None + scale = _load_indexed_tensor(source.scale_name, key_index) + global_scale = ( + _load_indexed_tensor(source.global_scale_name, key_index) + if source.global_scale_name is not None + else None + ) + tensor = reconstruct_modelopt_weight( + tensor, scale, global_scale, name=source.source_name + ) + return tensor if target_dtype == torch.bfloat16 else tensor.to(target_dtype) if source.mode == "fp8_block_128": assert source.scale_name is not None scale = _load_indexed_tensor(source.scale_name, key_index) @@ -912,6 +929,57 @@ def validate_source( if source.source_name not in key_index: raise ValueError(f"Streaming source '{source.source_name}' does not exist") _source_path, source_shape, source_dtype = key_index[source.source_name] + if source.mode in {"modelopt_fp8", "modelopt_nvfp4"}: + from mobius.integrations.modelopt._weights import validate_modelopt_scale + + if source.mode == "modelopt_fp8" and source.global_scale_name is not None: + raise ValueError("ModelOpt FP8 must not declare an NVFP4 global scale") + + expected_dtype = "U8" if source.mode == "modelopt_nvfp4" else "F8_E4M3" + if len(source_shape) != 2 or source_dtype != expected_dtype: + raise ValueError( + f"ModelOpt source '{source.source_name}' requires rank-2 " + f"{expected_dtype}, got {source_dtype}/{source_shape}" + ) + if source.scale_name is None or source.scale_name not in key_index: + raise ValueError(f"Missing ModelOpt weight scale for '{source.source_name}'") + scale_path, scale_shape, scale_dtype = key_index[source.scale_name] + if source.mode == "modelopt_nvfp4": + logical_width = source_shape[1] * 2 + if ( + logical_width % 16 + or scale_shape != [source_shape[0], logical_width // 16] + or scale_dtype != "F8_E4M3" + ): + raise ValueError(f"Invalid NVFP4 block scales for '{source.source_name}'") + if ( + source.global_scale_name is None + or source.global_scale_name not in key_index + ): + raise ValueError(f"Missing NVFP4 global scale for '{source.source_name}'") + global_path, global_shape, global_dtype = key_index[source.global_scale_name] + if global_shape not in ([], [1]) or global_dtype != "F32": + raise ValueError(f"Invalid NVFP4 global scale for '{source.source_name}'") + with safe_open(global_path, framework="pt") as handle: + validate_modelopt_scale( + handle.get_tensor(source.global_scale_name), + source.global_scale_name, + scalar=True, + ) + consumed.add(source.global_scale_name) + elif scale_shape not in ([], [1]) or scale_dtype != "F32": + raise ValueError(f"Invalid ModelOpt FP8 scale for '{source.source_name}'") + elif source.global_scale_name is not None: + raise ValueError( + f"FP8 source '{source.source_name}' must not have global scale" + ) + with safe_open(scale_path, framework="pt") as handle: + validate_modelopt_scale( + handle.get_tensor(source.scale_name), + source.scale_name, + scalar=source.mode == "modelopt_fp8", + ) + consumed.add(source.scale_name) if source.mode in {"fp8_block_128", "fp8_scalar"}: if source_dtype not in {"F8_E4M3", "F8_E5M2"}: raise ValueError( @@ -978,7 +1046,14 @@ def validate_source( ) scale_bytes = math.prod(scale_shape) * scale_element_bytes consumed.add(source.source_name) - return source_shape, source_bytes, scale_bytes + if source.global_scale_name is not None: + scale_bytes += 4 + logical_shape = ( + [source_shape[0], source_shape[1] * 2] + if source.mode == "modelopt_nvfp4" + else source_shape + ) + return logical_shape, source_bytes, scale_bytes for target_name, source in plan.targets.items(): initializer = model.graph.initializers.get(target_name) @@ -1036,7 +1111,15 @@ def validate_source( largest_source_tensor_bytes = max(largest_source_tensor_bytes, source_bytes) largest_reconstruction_working_set_bytes = max( largest_reconstruction_working_set_bytes, - source_bytes + bf16_bytes + cast_bytes + scale_bytes, + source_bytes + + bf16_bytes + + cast_bytes + + scale_bytes + + ( + min(256, source_shape[0]) * source_shape[1] * 40 + if source.mode in {"modelopt_fp8", "modelopt_nvfp4"} + else 0 + ), ) else: if len(expected_shape) != 3 or len(source.experts) != expected_shape[0]: diff --git a/src/mobius/integrations/modelopt/__init__.py b/src/mobius/integrations/modelopt/__init__.py index a7a2c0992..0c6d400dd 100644 --- a/src/mobius/integrations/modelopt/__init__.py +++ b/src/mobius/integrations/modelopt/__init__.py @@ -11,11 +11,11 @@ scales and a per-tensor FP32 global scale (``weight_scale_2``). - FP8 (``E4M3``): per-tensor scaled float8 weights. -The dequantization functions reconstruct BF16 weights so the standard mobius -build path (plain ``Linear`` / ``bf16`` graph) can consume ModelOpt checkpoints -without a native FP8/NVFP4 kernel. Native routed-expert NVFP4 QMoE emission -(the CUDA-only, Blackwell ``QMoE`` ``quant_type="nvfp4"`` op) is intentionally -out of scope here — see the module docstring in :mod:`._dequant`. +Nemotron-H's explicit ``keep_quantized=False, dtype="bf16"`` builder route +validates mixed-format inventories and reconstructs dense BF16 weights lazily. +It validates and accounts for activation and KV-cache scales, but does not +simulate their quantization. Native NVFP4/FP8 emission is unsupported; other +architectures remain blocked. NextN/MTP requires a separate validated contract. """ from __future__ import annotations diff --git a/src/mobius/integrations/modelopt/_config.py b/src/mobius/integrations/modelopt/_config.py new file mode 100644 index 000000000..3beb93b1b --- /dev/null +++ b/src/mobius/integrations/modelopt/_config.py @@ -0,0 +1,163 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Strict mixed FP8/NVFP4 weight-only reconstruction plans for ModelOpt exports.""" + +from __future__ import annotations + +import dataclasses +from collections.abc import Mapping +from typing import Literal + + +@dataclasses.dataclass(frozen=True) +class ModelOptConfig: + """Exact source-module formats, not an affine INT4 quantization config. + + Only explicit module targets with scalar-scaled E4M3 or block-16 E2M1 + weights are accepted. Input FP8 scales are validated and accounted for, + but dense reconstruction does not preserve activation quantization. + """ + + modules: Mapping[str, Literal["modelopt_fp8", "modelopt_nvfp4"]] + activation_modules: frozenset[str] + kv_cache_fp8: bool = False + + @classmethod + def parse(cls, value: object) -> ModelOptConfig: + """Reject unknown layouts instead of inferring semantics from bit width.""" + if not isinstance(value, Mapping) or value.get("quant_method") != "modelopt": + raise ValueError("ModelOpt reconstruction requires quant_method='modelopt'") + unknown = set(value) - { + "quant_method", + "config_groups", + "ignore", + "quant_algo", + "quantized_layers", + "kv_cache_scheme", + "producer", + } + if unknown: + raise NotImplementedError( + f"Unsupported ModelOpt configuration fields: {sorted(unknown)}" + ) + groups = value.get("config_groups") + if not isinstance(groups, Mapping) or not groups: + raise ValueError("ModelOpt reconstruction requires explicit config_groups") + modules: dict[str, Literal["modelopt_fp8", "modelopt_nvfp4"]] = {} + activation_modules: set[str] = set() + for group_name, group in groups.items(): + if not isinstance(group, Mapping): + raise TypeError(f"Malformed ModelOpt group {group_name!r}") + unknown = set(group) - { + "weights", + "input_activations", + "output_activations", + "targets", + } + if unknown: + raise NotImplementedError( + f"Unsupported ModelOpt group fields: {sorted(unknown)}" + ) + weights = group.get("weights") + if not isinstance(weights, Mapping) or weights.get("dynamic") is not False: + raise ValueError(f"ModelOpt weights in {group_name!r} must be static") + if weights == {"dynamic": False, "num_bits": 8, "type": "float"}: + mode: Literal["modelopt_fp8", "modelopt_nvfp4"] = "modelopt_fp8" + elif weights == { + "dynamic": False, + "num_bits": 4, + "type": "float", + "group_size": 16, + }: + mode = "modelopt_nvfp4" + else: + raise NotImplementedError( + f"Unsupported ModelOpt weight layout in {group_name!r}: {weights!r}" + ) + activation = group.get("input_activations") + if activation is not None and ( + mode != "modelopt_fp8" + or activation != {"dynamic": False, "num_bits": 8, "type": "float"} + or not isinstance(activation, Mapping) + or activation.get("dynamic") is not False + ): + raise NotImplementedError( + f"Unsupported ModelOpt activation layout in {group_name!r}" + ) + if group.get("output_activations") is not None: + raise NotImplementedError( + "ModelOpt output activation quantization unsupported" + ) + targets = group.get("targets") + if not isinstance(targets, list) or not targets: + raise ValueError(f"ModelOpt group {group_name!r} needs explicit targets") + for target in targets: + if ( + not isinstance(target, str) + or not target + or any( + not (part.isidentifier() or part.isdecimal()) + for part in target.split(".") + ) + ): + raise ValueError( + f"ModelOpt target must be an exact module path: {target!r}" + ) + if target in modules: + raise ValueError(f"Duplicate ModelOpt target {target!r}") + modules[target] = mode + if activation is not None: + activation_modules.add(target) + ignored = value.get("ignore", []) + if not isinstance(ignored, list) or any(not isinstance(x, str) for x in ignored): + raise ValueError("ModelOpt ignore must contain module paths") + if set(ignored) & modules.keys(): + raise ValueError("ModelOpt targets must not also be ignored") + algo = value.get("quant_algo") + if algo is not None and ( + algo != "MIXED_PRECISION" + or set(modules.values()) != {"modelopt_fp8", "modelopt_nvfp4"} + ): + raise NotImplementedError(f"Unsupported ModelOpt quant_algo {algo!r}") + legacy = value.get("quantized_layers") + if legacy is not None: + expected = { + name: ( + {"quant_algo": "FP8"} + if mode == "modelopt_fp8" + else {"quant_algo": "W4A16_NVFP4", "group_size": 16} + ) + for name, mode in modules.items() + } + if legacy != expected: + raise ValueError("ModelOpt quantized_layers and config_groups disagree") + cache_scheme = value.get("kv_cache_scheme") + if cache_scheme is not None and ( + not isinstance(cache_scheme, Mapping) + or cache_scheme != {"dynamic": False, "num_bits": 8, "type": "float"} + or cache_scheme.get("dynamic") is not False + ): + raise NotImplementedError( + "Only static E4M3 ModelOpt KV-cache scales are supported" + ) + producer = value.get("producer") + if producer is not None and ( + not isinstance(producer, Mapping) + or set(producer) != {"name", "version"} + or producer["name"] != "modelopt" + or not isinstance(producer["version"], str) + or not producer["version"] + ): + raise ValueError("Malformed ModelOpt producer metadata") + return cls(modules, frozenset(activation_modules), cache_scheme is not None) + + @property + def source_format(self) -> str: + """Describe the declared numeric formats rather than guessing from model name.""" + modes = set(self.modules.values()) + if modes == {"modelopt_fp8", "modelopt_nvfp4"}: + return "modelopt-mixed-fp8-nvfp4" + if modes == {"modelopt_fp8"}: + return "modelopt-fp8" + return "modelopt-nvfp4" diff --git a/src/mobius/integrations/modelopt/_dequant.py b/src/mobius/integrations/modelopt/_dequant.py index e87b9f4f1..39a31c711 100644 --- a/src/mobius/integrations/modelopt/_dequant.py +++ b/src/mobius/integrations/modelopt/_dequant.py @@ -4,21 +4,19 @@ """NVFP4 / FP8 weight reconstruction for NVIDIA ModelOpt checkpoints. NVIDIA TensorRT Model Optimizer (ModelOpt) exports mixed-precision checkpoints -where different modules use different numeric formats. For Qwen3.6 the layout is: +where different modules use different numeric formats. The supported numeric layout is: - **Routed MoE experts, shared expert, lm_head**: ``W4A16_NVFP4`` — block-16 E2M1 (fp4) 4-bit weights, FP8-E4M3 per-block scales, and a per-tensor FP32 global scale stored as ``weight_scale_2``. -- **Attention / linear-attention projections**: ``FP8`` (E4M3) — float8 weights +- **Selected dense projections**: ``FP8`` (E4M3) — float8 weights with a per-tensor ``weight_scale``. -ONNX Runtime has no FP8 attention GEMM, so — following ORT GenAI's ModelOpt -loader — the dense FP8 projections and the NVFP4 shared-expert / lm_head are -**dequantized back to BF16**, which reconstructs them exactly and lets the -standard (BF16) build path consume the checkpoint. Only the *routed* MoE experts -are consumed natively by the CUDA QMoE ``quant_type="nvfp4"`` op; that native -emission is Blackwell/``onnxruntime_USE_FP4_QMOE=ON``-only and is intentionally -NOT implemented here (it cannot be built or verified without that runtime). +The Nemotron-H loader can explicitly reconstruct ALL selected projections, +including routed experts, to dense BF16 using these helpers. The final BF16 +rounding and removal of activation quantization do not preserve native +quantized execution. There is no native NVFP4/QMoE emission here, nor an +automatic fallback to dense weights or affine integer INT4. The functions below are the numeric core of the loader and are format-faithful to ModelOpt's on-disk encoding (verified in ``_dequant_test.py``). diff --git a/src/mobius/integrations/modelopt/_weights.py b/src/mobius/integrations/modelopt/_weights.py new file mode 100644 index 000000000..2dceb31a0 --- /dev/null +++ b/src/mobius/integrations/modelopt/_weights.py @@ -0,0 +1,75 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Bounded-row, explicitly dense reconstruction of ModelOpt checkpoint weights.""" + +from __future__ import annotations + +import math + +import numpy as np +import torch + +from mobius.integrations.modelopt._dequant import dequantize_fp8, dequantize_nvfp4 + +_ROWS_PER_CHUNK = 256 + + +def validate_modelopt_scale(scale: torch.Tensor, name: str, *, scalar: bool) -> None: + """Validate scale values without broadcasting malformed rows or scalars.""" + if scalar and (scale.numel() != 1 or scale.ndim > 1): + raise ValueError(f"ModelOpt scalar scale {name!r} must have shape [] or [1]") + values = scale.to(torch.float32) + invalid = (values <= 0).any() if scalar else (values < 0).any() + if not torch.isfinite(values).all() or invalid: + requirement = "positive" if scalar else "nonnegative" + raise ValueError(f"ModelOpt scale {name!r} must be finite and {requirement}") + + +def reconstruct_modelopt_weight( + weight: torch.Tensor, + scale: torch.Tensor, + global_scale: torch.Tensor | None, + *, + name: str, +) -> torch.Tensor: + """Decode E4M3 or packed E2M1 to BF16, never native FP4 or affine INT4. + + The destination is allocated once; numeric-helper temporaries are bounded + to 256 rows. Scalar scaling happens in FP32 before one BF16 rounding. + """ + validate_modelopt_scale(scale, name + ".weight_scale", scalar=global_scale is None) + if weight.ndim != 2: + raise ValueError(f"ModelOpt weight {name!r} must be rank 2") + rows, width = weight.shape + if global_scale is not None: + if global_scale.dtype != torch.float32: + raise TypeError("NVFP4 global scale must use FP32") + validate_modelopt_scale(global_scale, name + ".weight_scale_2", scalar=True) + if weight.dtype != torch.uint8 or scale.dtype != torch.float8_e4m3fn: + raise TypeError("NVFP4 requires uint8 weights and E4M3 block scales") + width *= 2 + if width % 16 or tuple(scale.shape) != (rows, width // 16): + raise ValueError(f"NVFP4 block scale shape mismatch for {name!r}") + factor = float(global_scale.item()) + else: + if scale.dtype != torch.float32: + raise TypeError("ModelOpt FP8 scalar scale must use FP32") + if weight.dtype != torch.float8_e4m3fn: + raise TypeError("ModelOpt FP8 weights must use E4M3") + factor = float(scale.item()) + if not math.isfinite(factor) or factor <= 0: + raise ValueError(f"Invalid ModelOpt scalar scale for {name!r}") + result = torch.empty((rows, width), dtype=torch.bfloat16) + for start in range(0, rows, _ROWS_PER_CHUNK): + stop = min(start + _ROWS_PER_CHUNK, rows) + codes = weight[start:stop].contiguous().view(torch.uint8).numpy() + if global_scale is not None: + blocks = scale[start:stop].contiguous().view(torch.uint8).numpy() + dense = dequantize_nvfp4(codes, blocks, factor) + else: + dense = dequantize_fp8(codes, factor) + if not np.isfinite(dense.astype(np.float32)).all(): + raise ValueError(f"ModelOpt reconstruction overflow/nonfinite values in {name!r}") + result[start:stop] = torch.from_numpy(dense.view(np.uint16)).view(torch.bfloat16) + return result diff --git a/src/mobius/integrations/modelopt/_weights_test.py b/src/mobius/integrations/modelopt/_weights_test.py new file mode 100644 index 000000000..c86fa8054 --- /dev/null +++ b/src/mobius/integrations/modelopt/_weights_test.py @@ -0,0 +1,258 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Format-faithful mixed ModelOpt reconstruction and fail-closed loader regressions.""" + +from __future__ import annotations + +import numpy as np +import onnx_ir as ir +import pytest +import safetensors.torch +import torch + +from mobius.integrations._weight_loading import ( + StreamingWeightPlan, + StreamingWeightSource, + stream_preprocessed_safetensors_to_model, +) +from mobius.integrations.modelopt._config import ModelOptConfig +from mobius.integrations.modelopt._weights import reconstruct_modelopt_weight + + +def _mixed_config(): + return { + "quant_method": "modelopt", + "config_groups": { + "fp8": { + "weights": {"dynamic": False, "num_bits": 8, "type": "float"}, + "input_activations": {"dynamic": False, "num_bits": 8, "type": "float"}, + "targets": ["backbone.layers.0.mixer.in_proj"], + }, + "fp4": { + "weights": { + "dynamic": False, + "num_bits": 4, + "type": "float", + "group_size": 16, + }, + "targets": ["backbone.layers.1.mixer.experts.0.up_proj", "lm_head"], + }, + }, + } + + +def test_exact_mixed_group_layout(): + config = ModelOptConfig.parse(_mixed_config()) + assert config.modules["backbone.layers.0.mixer.in_proj"] == "modelopt_fp8" + assert config.modules["lm_head"] == "modelopt_nvfp4" + assert config.activation_modules == frozenset({"backbone.layers.0.mixer.in_proj"}) + assert config.source_format == "modelopt-mixed-fp8-nvfp4" + + +@pytest.mark.parametrize("unknown", ["format", "strategy"]) +def test_unknown_modelopt_fields_are_not_silently_ignored(unknown): + raw = _mixed_config() + raw[unknown] = "unsupported" + with pytest.raises(NotImplementedError, match="configuration fields"): + ModelOptConfig.parse(raw) + raw = _mixed_config() + raw["config_groups"]["fp4"][unknown] = "unsupported" + with pytest.raises(NotImplementedError, match="group fields"): + ModelOptConfig.parse(raw) + + +def test_redundant_format_and_cache_declarations_must_agree(): + raw = _mixed_config() + raw["quantized_layers"] = {"lm_head": {"quant_algo": "FP8"}} + with pytest.raises(ValueError, match="disagree"): + ModelOptConfig.parse(raw) + raw = _mixed_config() + raw["kv_cache_scheme"] = {"dynamic": True, "num_bits": 8, "type": "float"} + with pytest.raises(NotImplementedError, match="KV-cache"): + ModelOptConfig.parse(raw) + + +@pytest.mark.parametrize( + "change", ["group-size", "dynamic", "integer", "duplicate", "wildcard"] +) +def test_reject_unknown_or_ambiguous_group_layout(change): + raw = _mixed_config() + group = raw["config_groups"]["fp4"] + if change == "group-size": + group["weights"]["group_size"] = 32 + elif change == "dynamic": + group["weights"]["dynamic"] = True + elif change == "integer": + group["weights"]["type"] = "int" + elif change == "duplicate": + group["targets"].append("lm_head") + else: + group["targets"] = ["backbone.layers.*"] + with pytest.raises((ValueError, NotImplementedError)): + ModelOptConfig.parse(raw) + + +def test_fp4_all_codes_distinct_rows_blocks_and_chunk_boundary(): + # E2M1 codes 0..15 exercise signs, nibble order and zero. 257 rows cross + # the bounded-row materialization boundary; scales differ in every block. + codes = torch.arange(16, dtype=torch.uint8).repeat(257, 2) + packed = codes[:, 0::2] | (codes[:, 1::2] << 4) + scales = torch.tensor([0.5, 2.0]).repeat(257, 1).to(torch.float8_e4m3fn) + scales[256] = torch.tensor([1.0, 4.0]).to(torch.float8_e4m3fn) + output = reconstruct_modelopt_weight( + packed, scales, torch.tensor(0.25), name="experts.0.up_proj" + ) + magnitude = torch.tensor([0, 0.5, 1, 1.5, 2, 3, 4, 6]) + expected = magnitude[(codes & 7).long()] + expected = torch.where((codes & 8).bool(), -expected, expected) + expected *= scales.float().repeat_interleave(16, dim=1) * 0.25 + assert output.dtype == torch.bfloat16 + torch.testing.assert_close(output, expected.bfloat16(), rtol=0, atol=0) + assert torch.equal(torch.signbit(output), torch.signbit(expected)) + + +def test_fp8_single_rounding_with_nonunit_scalar(): + weight = torch.tensor([[1.0, -2.0, 1.5, 0.25]]).to(torch.float8_e4m3fn) + scale = torch.tensor(0.123456, dtype=torch.float32) + output = reconstruct_modelopt_weight(weight, scale, None, name="mamba.in_proj") + torch.testing.assert_close(output, (weight.float() * scale).bfloat16(), rtol=0, atol=0) + + +def test_zero_nvfp4_block_scale_reconstructs_zero_block(): + packed = torch.full((1, 16), 0x92, dtype=torch.uint8) + blocks = torch.tensor([[0.0, 2.0]]).to(torch.float8_e4m3fn) + output = reconstruct_modelopt_weight(packed, blocks, torch.tensor(0.5), name="expert") + assert torch.equal(output[:, :16], torch.zeros(1, 16, dtype=torch.bfloat16)) + torch.testing.assert_close( + output[:, 16:], + torch.tensor([[1.0, -0.5] * 8], dtype=torch.bfloat16), + rtol=0, + atol=0, + ) + + +@pytest.mark.parametrize("bad", [0.0, -1.0, float("nan"), float("inf")]) +def test_reject_nonpositive_nonfinite_scale(bad): + with pytest.raises(ValueError, match="finite and positive"): + reconstruct_modelopt_weight( + torch.ones(1, 16).to(torch.float8_e4m3fn), + torch.tensor(bad), + None, + name="mamba", + ) + + +def test_reject_scale_row_broadcast(): + with pytest.raises(ValueError, match="shape mismatch"): + reconstruct_modelopt_weight( + torch.full((2, 8), 0x22, dtype=torch.uint8), + torch.ones(1, 1).to(torch.float8_e4m3fn), + torch.tensor(1.0), + name="expert", + ) + + +def _stream_fixture(tmp_path, *, mutate=None): + tensors = { + "fp8.weight": torch.tensor([[1.0, -2.0] * 8]).to(torch.float8_e4m3fn), + "fp8.weight_scale": torch.tensor(0.25), + "fp4.weight": torch.full((2, 16), 0x92, dtype=torch.uint8), + "fp4.weight_scale": torch.tensor([[0.5, 2.0], [1.0, 4.0]]).to(torch.float8_e4m3fn), + "fp4.weight_scale_2": torch.tensor(0.5), + } + if mutate is not None: + mutate(tensors) + path = tmp_path / "model.safetensors" + safetensors.torch.save_file(tensors, path) + initializers = { + name: ir.Value( + name=name, shape=ir.Shape(shape), type=ir.TensorType(ir.DataType.BFLOAT16) + ) + for name, shape in {"fp8.weight": [1, 16], "fp4.weight": [2, 32]}.items() + } + model = ir.Model( + ir.Graph([], [], nodes=[], initializers=initializers.values()), ir_version=11 + ) + plan = StreamingWeightPlan( + { + "fp8.weight": StreamingWeightSource( + "fp8.weight", "modelopt_fp8", "fp8.weight_scale" + ), + "fp4.weight": StreamingWeightSource( + "fp4.weight", + "modelopt_nvfp4", + "fp4.weight_scale", + global_scale_name="fp4.weight_scale_2", + ), + } + ) + return tensors, path, model, plan + + +def test_streaming_mixed_weights_materialize_and_report_dense(tmp_path): + tensors, _path, model, plan = _stream_fixture(tmp_path) + report = stream_preprocessed_safetensors_to_model(model, str(tmp_path), lambda *_: plan) + assert report["output_weight_format"] == "dense" + assert report["native_fp8"] is False + for name in ("fp8.weight", "fp4.weight"): + expected = reconstruct_modelopt_weight( + tensors[name], + tensors[name.replace(".weight", ".weight_scale")], + tensors.get(name.replace(".weight", ".weight_scale_2")), + name=name, + ) + actual = model.graph.initializers[name].const_value.numpy().astype(np.float32) + np.testing.assert_array_equal(actual, expected.float().numpy()) + # Exercise serialization and reload, not just a success-shaped lazy binding. + out = tmp_path / "dense.onnx" + ir.save(model, out) + loaded = ir.load(out) + assert all(x.dtype == ir.DataType.BFLOAT16 for x in loaded.graph.initializers.values()) + + +@pytest.mark.parametrize( + "change", ["missing-scale", "row-broadcast", "integer-scale", "extra"] +) +def test_streaming_preflight_rejects_before_any_binding(tmp_path, change): + def mutate(tensors): + if change == "missing-scale": + del tensors["fp4.weight_scale_2"] + elif change == "row-broadcast": + tensors["fp4.weight_scale"] = torch.ones(1, 2).to(torch.float8_e4m3fn) + elif change == "integer-scale": + tensors["fp8.weight_scale"] = torch.tensor(1, dtype=torch.int64) + else: + tensors["unknown.weight"] = torch.ones(1) + + _tensors, _path, model, plan = _stream_fixture(tmp_path, mutate=mutate) + with pytest.raises(ValueError): + stream_preprocessed_safetensors_to_model(model, str(tmp_path), lambda *_: plan) + assert all(x.const_value is None for x in model.graph.initializers.values()) + + +def test_lazy_reconstruction_rechecks_source_header(tmp_path): + tensors, path, model, plan = _stream_fixture(tmp_path) + stream_preprocessed_safetensors_to_model(model, str(tmp_path), lambda *_: plan) + tensors["fp4.weight_scale_2"] = torch.ones(2) + safetensors.torch.save_file(tensors, path) + with pytest.raises(ValueError, match="changed after indexing"): + model.graph.initializers["fp4.weight"].const_value.numpy() + + +def test_fp8_descriptor_cannot_borrow_an_nvfp4_global_scale(tmp_path): + _tensors, _path, model, original = _stream_fixture(tmp_path) + plan = StreamingWeightPlan( + { + **original.targets, + "fp8.weight": StreamingWeightSource( + "fp8.weight", + "modelopt_fp8", + "fp8.weight_scale", + global_scale_name="fp4.weight_scale_2", + ), + } + ) + with pytest.raises(ValueError, match="must not declare"): + stream_preprocessed_safetensors_to_model(model, str(tmp_path), lambda *_: plan) + assert all(x.const_value is None for x in model.graph.initializers.values()) diff --git a/src/mobius/integrations/transformers/_builder.py b/src/mobius/integrations/transformers/_builder.py index 3dbe632f6..f5f0e92da 100644 --- a/src/mobius/integrations/transformers/_builder.py +++ b/src/mobius/integrations/transformers/_builder.py @@ -5,6 +5,7 @@ from __future__ import annotations +import copy import dataclasses import logging import pathlib @@ -31,6 +32,7 @@ CompressedTensorsConfig, stream_compressed_tensors_to_package, ) +from mobius.integrations.modelopt._config import ModelOptConfig from mobius.tasks import ModelTask from mobius.weights import adapt_model_weights @@ -45,6 +47,40 @@ ) +def _prepare_nemotron_modelopt_source( + hf_config, + parent_config, + *, + keep_quantized: bool, + dtype: str | ir.DataType | None, +): + """Admit only explicit BF16 weight reconstruction; native ModelOpt stays blocked.""" + from mobius.integrations.modelopt import is_modelopt_quant_config + + raw = getattr(hf_config, "quantization_config", None) + if raw is not None and hasattr(raw, "to_dict"): + raw = raw.to_dict() + if getattr(hf_config, "model_type", None) != "nemotron_h" or not is_modelopt_quant_config( + raw + ): + return hf_config, parent_config, None + if keep_quantized: + raise NotImplementedError( + "Native Nemotron-H ModelOpt NVFP4/FP8 preservation is unsupported. " + "Explicit dense reconstruction requires --dequantize --dtype bf16 " + "(keep_quantized=False, dtype='bf16'); no automatic fallback is available." + ) + if dtype is None or resolve_dtype(dtype) != ir.DataType.BFLOAT16: + raise ValueError("ModelOpt dense reconstruction requires explicit dtype='bf16'") + source = ModelOptConfig.parse(raw) + dense_config = copy.copy(hf_config) + dense_config.quantization_config = None + dense_parent = dense_config if parent_config is hf_config else copy.copy(parent_config) + if dense_parent is not dense_config: + raise NotImplementedError("ModelOpt Nemotron-H composite configs are unsupported") + return dense_config, dense_parent, source + + def _uses_affine_checkpoint_loader(config: object) -> bool: """Whether config requires the generic Olive/GPTQ/AWQ packed loader.""" component_quantization = getattr(config, "component_quantization", None) @@ -517,6 +553,9 @@ def build_transformers_model( ) hf_config, parent_config, model_type = _select_primary_config(hf_config) + hf_config, parent_config, modelopt_source = _prepare_nemotron_modelopt_source( + hf_config, parent_config, keep_quantized=keep_quantized, dtype=dtype + ) compressed_tensors_config = CompressedTensorsConfig.from_hf_config(parent_config) source_model_type = model_type @@ -675,6 +714,17 @@ def build_transformers_model( ) for name, model in package.items(): model.graph.name = f"{graph_source_name}/{name}" + if modelopt_source is not None: + model.metadata_props["mobius.storage_policy"] = ( + "explicit-dense-bf16-reconstruction" + ) + model.metadata_props["mobius.source_weight_format"] = modelopt_source.source_format + model.metadata_props["mobius.source_revision"] = revision or "local-or-unpinned" + model.metadata_props["mobius.weight_loading_status"] = "not_loaded" + model.metadata_props["mobius.source_kv_cache_format"] = ( + "modelopt-fp8" if modelopt_source.kv_cache_fp8 else "floating" + ) + model.metadata_props["mobius.kv_cache_quantization_preserved"] = "false" if model_type in _QWEN4_MODEL_TYPES | { "vibevoice", "vibevoice_streaming", @@ -692,7 +742,30 @@ def build_transformers_model( if load_weights: _reject_unsupported_affine_qwen4(model_type, config) - if is_gptoss_mxfp4_source and keep_quantized: + if modelopt_source is not None: + from mobius.integrations.transformers._nemotron_h_weights import ( + build_nemotron_h_streaming_plan, + ) + + if len(package) != 1: + raise ValueError("Nemotron-H dense streaming requires one decoder component") + report = stream_preprocessed_safetensors_to_model( + next(iter(package.values())), + model_id, + lambda index, initializers: build_nemotron_h_streaming_plan( + config, index, initializers, modelopt=modelopt_source + ), + revision=revision, + ) + package.weight_loading_report = report + for model in package.values(): + model.metadata_props["mobius.weight_loading_status"] = "lazy-bound" + logger.warning( + "Explicitly reconstructing ModelOpt weights to dense BF16; native " + "NVFP4/FP8 storage, FP8 activations and FP8 KV-cache quantization " + "are not preserved." + ) + elif is_gptoss_mxfp4_source and keep_quantized: from mobius.integrations.transformers._gptoss_weights import ( stream_gptoss_mxfp4_safetensors_to_package, ) diff --git a/src/mobius/integrations/transformers/_nemotron_h_weights.py b/src/mobius/integrations/transformers/_nemotron_h_weights.py new file mode 100644 index 000000000..982d24dfd --- /dev/null +++ b/src/mobius/integrations/transformers/_nemotron_h_weights.py @@ -0,0 +1,113 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Fail-closed streaming plans for Nemotron-H BF16 and reconstructed ModelOpt weights.""" + +from __future__ import annotations + +from collections.abc import Mapping + +import onnx_ir as ir +import torch + +from mobius._configs import NemotronHConfig +from mobius.integrations._weight_loading import ( + StreamingWeightPlan, + StreamingWeightSource, + _load_indexed_tensor, +) +from mobius.integrations.modelopt._config import ModelOptConfig +from mobius.integrations.modelopt._weights import validate_modelopt_scale +from mobius.models.nemotron_h import _rename_nemotron_h_weight + + +def build_nemotron_h_streaming_plan( + config: NemotronHConfig, + key_index: Mapping[str, tuple[str, list[int], str]], + initializers: Mapping[str, ir.Value], + *, + modelopt: ModelOptConfig | None = None, +) -> StreamingWeightPlan: + """Map exact per-expert checkpoint keys and account for every source tensor.""" + if any(name.startswith("mtp.") for name in key_index): + raise NotImplementedError( + "Nemotron-H NextN/MTP tensors require a validated MTP export contract; " + "they must not be silently discarded." + ) + targets: dict[str, StreamingWeightSource] = {} + constants: dict[str, torch.Tensor] = {} + ignored: dict[str, str] = {} + quantized_modules: set[str] = set() + activation_modules = modelopt.activation_modules if modelopt is not None else frozenset() + cache_scales = set() + if modelopt is not None and modelopt.kv_cache_fp8: + for layer_index, kind in enumerate(config.layer_types or []): + if kind == "full_attention": + for projection, suffix in (("k_proj", "k_scale"), ("v_proj", "v_scale")): + cache_scales.add( + f"backbone.layers.{layer_index}.mixer.{projection}.{suffix}" + ) + for name in sorted(cache_scales): + metadata = key_index.get(name) + if metadata is None or metadata[1] != [1] or metadata[2] != "F32": + raise ValueError(f"Missing/malformed ModelOpt KV-cache scale {name!r}") + scale = _load_indexed_tensor(name, key_index) + validate_modelopt_scale(scale, name, scalar=True) + constants[name] = scale + for source_name in key_index: + if source_name in cache_scales: + continue + if source_name.endswith((".weight_scale", ".weight_scale_2", ".input_scale")): + continue + target_name = _rename_nemotron_h_weight(source_name, config.layer_types or []) + if target_name in targets: + raise ValueError(f"Multiple Nemotron-H sources map to {target_name!r}") + if target_name not in initializers: + raise ValueError(f"Unmapped Nemotron-H checkpoint tensor {source_name!r}") + module_name = source_name.removesuffix(".weight") + mode = modelopt.modules.get(module_name) if modelopt is not None else None + if mode is not None: + if not source_name.endswith(".weight"): + raise ValueError(f"ModelOpt target {module_name!r} is not a projection weight") + quantized_modules.add(module_name) + targets[target_name] = StreamingWeightSource( + source_name, + mode=mode, + scale_name=module_name + ".weight_scale", + global_scale_name=( + module_name + ".weight_scale_2" if mode == "modelopt_nvfp4" else None + ), + ) + if module_name in activation_modules: + input_name = module_name + ".input_scale" + metadata = key_index.get(input_name) + if metadata is None or metadata[1] not in ([], [1]) or metadata[2] != "F32": + raise ValueError(f"Missing/malformed ModelOpt input scale {input_name!r}") + scale = _load_indexed_tensor(input_name, key_index) + validate_modelopt_scale(scale, input_name, scalar=True) + # Explicitly validated, but not used: this route reconstructs + # weights for dense BF16 execution, not FP8 activations. + constants[input_name] = scale + else: + if key_index[source_name][2] not in {"BF16", "F16", "F32"}: + raise ValueError(f"Unclassified quantized Nemotron-H weight {source_name!r}") + targets[target_name] = StreamingWeightSource(source_name) + if modelopt is not None and quantized_modules != set(modelopt.modules): + raise ValueError( + "ModelOpt projection inventory mismatch; missing " + f"{sorted(set(modelopt.modules) - quantized_modules)[:5]}" + ) + return StreamingWeightPlan( + targets, + ignored=ignored, + constants=constants, + report={ + "source_weight_format": modelopt.source_format if modelopt else "floating", + "storage_policy": "explicit-dense-bf16-reconstruction" if modelopt else "dense", + "native_nvfp4": False, + "activation_quantization_preserved": False, + "kv_cache_quantization_preserved": False, + "validated_kv_cache_scale_count": len(cache_scales), + "reconstructed_projection_count": len(quantized_modules), + }, + ) diff --git a/src/mobius/integrations/transformers/_nemotron_h_weights_test.py b/src/mobius/integrations/transformers/_nemotron_h_weights_test.py new file mode 100644 index 000000000..d38474042 --- /dev/null +++ b/src/mobius/integrations/transformers/_nemotron_h_weights_test.py @@ -0,0 +1,308 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Offline public-builder and exact per-expert weight-format regressions.""" + +from __future__ import annotations + +import json +from types import SimpleNamespace + +import onnx_ir as ir +import pytest +import safetensors.torch +import torch + +from mobius import build +from mobius.__main__ import main +from mobius._configs import NemotronHConfig +from mobius.integrations._weight_loading import stream_preprocessed_safetensors_to_model +from mobius.integrations.transformers import _builder +from mobius.integrations.transformers._builder import _prepare_nemotron_modelopt_source +from mobius.integrations.transformers._nemotron_h_weights import ( + build_nemotron_h_streaming_plan, +) + + +def _config(): + # Reduced pinned Lightning topology/format, NOT reduced real checkpoint weights. + return { + "model_type": "nemotron_h", + "architectures": ["NemotronHForCausalLM"], + "dtype": "bfloat16", + "hidden_size": 32, + "vocab_size": 48, + "intermediate_size": 32, + "num_hidden_layers": 3, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "head_dim": 8, + "layers_block_type": ["mamba", "moe", "attention"], + "n_routed_experts": 4, + "num_experts_per_tok": 2, + "moe_intermediate_size": 32, + "moe_shared_expert_intermediate_size": 64, + "mamba_num_heads": 4, + "mamba_head_dim": 8, + "n_groups": 2, + "ssm_state_size": 8, + "mamba_ssm_cache_dtype": "float32", + "layer_norm_epsilon": 1e-5, + "conv_kernel": 4, + "routed_scaling_factor": 2.5, + "num_nextn_predict_layers": 0, + } + + +def _source_name(target): + if target == "model.embed_tokens.weight": + return "backbone.embeddings.weight" + if target == "model.norm.weight": + return "backbone.norm_f.weight" + if target.startswith("model.layers."): + name = target.replace("model.", "backbone.", 1) + for kind in ("mamba", "moe", "self_attn"): + name = name.replace(f".{kind}.", ".mixer.") + return name + return target + + +def _write_mixed_checkpoint(tmp_path): + raw = _config() + config_path = tmp_path / "config.json" + config_path.write_text(json.dumps(raw), encoding="utf-8") + graph = build(str(tmp_path), load_weights=False, dtype="bf16")["model"].graph + tensors = {} + fp8_modules = [] + fp4_modules = [] + expected = {} + for name, initializer in graph.initializers.items(): + if initializer.const_value is not None: + continue + source = _source_name(name) + shape = tuple(initializer.shape) + module = source.removesuffix(".weight") + if ".mixer.in_proj." in source or ".mixer.out_proj." in source: + tensor = torch.full(shape, 1.5).to(torch.float8_e4m3fn) + tensors[source] = tensor + tensors[module + ".weight_scale"] = torch.tensor(0.25) + tensors[module + ".input_scale"] = torch.tensor([0.5]) + fp8_modules.append(module) + expected[name] = torch.full(shape, 0.375, dtype=torch.bfloat16) + elif ( + ".experts." in source or ".shared_experts." in source or source == "lm_head.weight" + ): + tensors[source] = torch.full((shape[0], shape[1] // 2), 0x92, dtype=torch.uint8) + tensors[module + ".weight_scale"] = torch.ones(shape[0], shape[1] // 16).to( + torch.float8_e4m3fn + ) + tensors[module + ".weight_scale_2"] = torch.tensor(0.5) + fp4_modules.append(module) + dense = torch.empty(shape, dtype=torch.bfloat16) + dense[:, 0::2] = 0.5 + dense[:, 1::2] = -0.25 + expected[name] = dense + else: + tensors[source] = torch.full( + shape, + 0.125, + dtype=torch.float32 + if initializer.dtype == ir.DataType.FLOAT + else torch.bfloat16, + ) + expected[name] = tensors[source] + raw["quantization_config"] = { + "quant_method": "modelopt", + "quant_algo": "MIXED_PRECISION", + "producer": {"name": "modelopt", "version": "0.44.0rc5"}, + "kv_cache_scheme": {"dynamic": False, "num_bits": 8, "type": "float"}, + "quantized_layers": { + **{name: {"quant_algo": "FP8"} for name in fp8_modules}, + **{name: {"quant_algo": "W4A16_NVFP4", "group_size": 16} for name in fp4_modules}, + }, + "config_groups": { + "group_0": { + "weights": {"dynamic": False, "num_bits": 8, "type": "float"}, + "input_activations": {"dynamic": False, "num_bits": 8, "type": "float"}, + "targets": fp8_modules, + }, + "group_1": { + "weights": { + "dynamic": False, + "num_bits": 4, + "type": "float", + "group_size": 16, + }, + "targets": fp4_modules, + }, + }, + } + for projection, suffix in (("k_proj", "k_scale"), ("v_proj", "v_scale")): + tensors[f"backbone.layers.2.mixer.{projection}.{suffix}"] = torch.tensor([0.25]) + config_path.write_text(json.dumps(raw), encoding="utf-8") + safetensors.torch.save_file(tensors, tmp_path / "model.safetensors") + return raw, expected + + +def test_public_local_build_reconstructs_all_projection_families(tmp_path, monkeypatch): + _raw, expected = _write_mixed_checkpoint(tmp_path) + monkeypatch.setattr( + _builder, + "_download_weights", + lambda *args, **kwargs: pytest.fail("ModelOpt must never reach eager/affine loader"), + ) + package = build(str(tmp_path), keep_quantized=False, dtype="bf16") + report = package.weight_loading_report + assert report["storage_policy"] == "explicit-dense-bf16-reconstruction" + assert report["native_nvfp4"] is False + assert report["ignored_tensors"] == 0 + assert report["validated_constants"] == 4 # Input + KV scales, not simulated. + assert report["validated_kv_cache_scale_count"] == 2 + assert report["kv_cache_quantization_preserved"] is False + assert report["reconstructed_projection_count"] == 13 + graph = package["model"].graph + assert not any(node.op_type in {"MatMulNBits", "QMoE"} for node in graph) + for name, value in expected.items(): + actual = graph.initializers[name].const_value.numpy().astype("float32") + torch.testing.assert_close(torch.from_numpy(actual), value.float(), rtol=0, atol=0) + output = tmp_path / "export" + package.save(output, external_data="onnx") + assert (output / "model.onnx").is_file() + loaded = ir.load(output / "model.onnx") + assert loaded.metadata_props["mobius.source_weight_format"] == "modelopt-mixed-fp8-nvfp4" + + +def test_floating_per_expert_checkpoint_uses_exact_streaming_mapping(tmp_path): + raw, expected = _write_mixed_checkpoint(tmp_path) + del raw["quantization_config"] + (tmp_path / "config.json").write_text(json.dumps(raw), encoding="utf-8") + safetensors.torch.save_file( + {_source_name(name): tensor for name, tensor in expected.items()}, + tmp_path / "model.safetensors", + ) + package = build(str(tmp_path), load_weights=False, dtype="bf16") + config = NemotronHConfig.from_transformers(SimpleNamespace(**raw)) + report = stream_preprocessed_safetensors_to_model( + package["model"], + str(tmp_path), + lambda index, initializers: build_nemotron_h_streaming_plan( + config, + index, + initializers, + ), + ) + assert report["ignored_tensors"] == 0 + assert report["storage_policy"] == "dense" + assert report["reconstructed_projection_count"] == 0 + for name, tensor in expected.items(): + actual = ( + package["model"].graph.initializers[name].const_value.numpy().astype("float32") + ) + torch.testing.assert_close(torch.from_numpy(actual), tensor.float(), rtol=0, atol=0) + + +def test_cli_explicit_dequantize_writes_dense_not_native_format(tmp_path): + _raw, _expected = _write_mixed_checkpoint(tmp_path) + output = tmp_path / "cli-export" + main( + [ + "build", + "--config", + str(tmp_path), + "--output", + str(output), + "--ep", + "cpu", + "--dtype", + "bf16", + "--dequantize", + "--external-data", + "onnx", + ] + ) + model = ir.load(output / "model.onnx") + assert ( + model.metadata_props["mobius.storage_policy"] == "explicit-dense-bf16-reconstruction" + ) + report = json.loads((output / "weight-loading-report.json").read_text()) + assert report["source_weight_format"] == "modelopt-mixed-fp8-nvfp4" + assert report["native_nvfp4"] is False + + +@pytest.mark.parametrize( + "kwargs,error", + [ + ({}, NotImplementedError), + ({"keep_quantized": False}, ValueError), + ({"keep_quantized": False, "dtype": "f16"}, ValueError), + ], +) +def test_public_native_or_implicit_dtype_requests_fail_before_weight_io( + tmp_path, monkeypatch, kwargs, error +): + raw = _config() + raw["quantization_config"] = {"quant_method": "modelopt"} + (tmp_path / "config.json").write_text(json.dumps(raw), encoding="utf-8") + monkeypatch.setattr( + _builder, + "_download_weights", + lambda *args, **kwargs: pytest.fail("unsupported path attempted weight download"), + ) + with pytest.raises(error): + build(str(tmp_path), **kwargs) + + +def test_reconstruction_policy_does_not_mutate_source_config(): + raw = SimpleNamespace( + model_type="nemotron_h", + quantization_config={ + "quant_method": "modelopt", + "config_groups": { + "group": { + "weights": {"dynamic": False, "num_bits": 8, "type": "float"}, + "targets": ["projection"], + }, + }, + }, + ) + dense, parent, source = _prepare_nemotron_modelopt_source( + raw, raw, keep_quantized=False, dtype="bf16" + ) + assert dense is not raw and parent is dense + assert dense.quantization_config is None + assert raw.quantization_config["quant_method"] == "modelopt" + assert source.modules == {"projection": "modelopt_fp8"} + + +def test_mtp_sources_fail_closed_even_when_config_omits_nextn(tmp_path, monkeypatch): + _raw, _expected = _write_mixed_checkpoint(tmp_path) + path = tmp_path / "model.safetensors" + weights = safetensors.torch.load_file(path) + weights["mtp.layers.0.hnorm.weight"] = torch.ones(32, dtype=torch.bfloat16) + safetensors.torch.save_file(weights, path) + monkeypatch.setattr( + "mobius.integrations.transformers._nemotron_h_weights._load_indexed_tensor", + lambda *args: pytest.fail("MTP source refusal must precede scale/value reads"), + ) + with pytest.raises(NotImplementedError, match="MTP tensors"): + build(str(tmp_path), keep_quantized=False, dtype="bf16") + + +@pytest.mark.parametrize("change", ["missing", "wrong-shape", "nonfinite", "undeclared"]) +def test_cache_scale_contract_fails_before_initializer_binding(tmp_path, change): + raw, _expected = _write_mixed_checkpoint(tmp_path) + path = tmp_path / "model.safetensors" + weights = safetensors.torch.load_file(path) + key = "backbone.layers.2.mixer.k_proj.k_scale" + if change == "missing": + del weights[key] + elif change == "wrong-shape": + weights[key] = torch.ones(2) + elif change == "nonfinite": + weights[key] = torch.tensor([float("nan")]) + else: + del raw["quantization_config"]["kv_cache_scheme"] + (tmp_path / "config.json").write_text(json.dumps(raw), encoding="utf-8") + safetensors.torch.save_file(weights, path) + with pytest.raises(ValueError): + build(str(tmp_path), keep_quantized=False, dtype="bf16") diff --git a/src/mobius/models/_nemotron_h_test.py b/src/mobius/models/_nemotron_h_test.py new file mode 100644 index 000000000..9112affce --- /dev/null +++ b/src/mobius/models/_nemotron_h_test.py @@ -0,0 +1,240 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Lightning config, checkpoint accounting, routing dtype and recurrent ABI regressions.""" + +from __future__ import annotations + +import dataclasses +from types import SimpleNamespace + +import numpy as np +import onnx_ir as ir +import onnxruntime as ort +import pytest +import torch + +from mobius import build_from_module +from mobius._builder import _cast_module_dtype +from mobius._configs import NemotronHConfig +from mobius._testing import create_test_builder, create_test_input +from mobius.integrations._weight_loading import apply_weights +from mobius.models.nemotron_h import ( + NemotronHCausalLMModel, + NemotronHMoEBlock, + _rename_nemotron_h_weight, +) + +# Selected architecture fields from pinned BF16 config +# a9904d24bcc1d289a1950fa9d2b978c47cf903b9 (raw SHA-256 in the recipe). +# NVFP4 uses the same topology, but max_position_embeddings=1048576. +_LIGHTNING = { + "model_type": "nemotron_h", + "hidden_size": 2688, + "intermediate_size": 1856, + "vocab_size": 131072, + "num_attention_heads": 32, + "num_key_value_heads": 2, + "head_dim": 128, + "num_hidden_layers": 52, + "max_position_embeddings": 262144, + "num_nextn_predict_layers": 1, + "layer_norm_epsilon": 1e-5, + "layers_block_type": [ + "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", + ], + "mamba_num_heads": 64, + "mamba_head_dim": 64, + "n_groups": 8, + "ssm_state_size": 128, + "mamba_ssm_cache_dtype": "float32", + "n_routed_experts": 128, + "num_experts_per_tok": 6, + "moe_intermediate_size": 1856, + "moe_shared_expert_intermediate_size": 3712, + "routed_scaling_factor": 2.5, + "norm_topk_prob": True, + "dtype": "bfloat16", +} + + +def _tiny_config(**kwargs): + config = NemotronHConfig( + model_type="nemotron_h", + hidden_size=32, + vocab_size=48, + intermediate_size=32, + num_hidden_layers=3, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=8, + layer_types=["mamba2", "moe", "full_attention"], + mamba_n_heads=4, + mamba_d_head=8, + mamba_d_state=8, + mamba_n_groups=2, + mamba_ssm_dtype=ir.DataType.FLOAT, + hidden_act="relu2", + num_local_experts=4, + num_experts_per_tok=2, + moe_intermediate_size=32, + shared_expert_intermediate_size=64, + routed_scaling_factor=2.5, + rms_norm_eps=1e-5, + ) + return dataclasses.replace(config, **kwargs) + + +@pytest.mark.parametrize("context", [262144, 1048576]) +def test_pinned_lightning_architecture_fields(context): + raw = SimpleNamespace(**{**_LIGHTNING, "max_position_embeddings": context}) + config = NemotronHConfig.from_transformers(raw) + assert len(config.layer_types) == 52 + 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_local_experts == 128 and config.num_experts_per_tok == 6 + assert config.shared_expert_intermediate_size == 3712 + assert config.mamba_n_heads == 64 and config.mamba_d_head == 64 + assert config.mamba_ssm_dtype == ir.DataType.FLOAT + assert config.rms_norm_eps == pytest.approx(1e-5) and config.hidden_act == "relu2" + assert config.max_position_embeddings == context + with pytest.raises(NotImplementedError, match="NextN/MTP"): + NemotronHCausalLMModel(config) + + +@pytest.mark.parametrize("dtype", [ir.DataType.FLOAT16, ir.DataType.BFLOAT16]) +def test_mamba_cache_abi_preserves_configured_float32_ssm(dtype): + config = _tiny_config(dtype=dtype) + package = build_from_module( + NemotronHCausalLMModel(config), config, "hybrid-text-generation" + ) + graph = package["model"].graph + inputs = {x.name: x for x in graph.inputs} + outputs = {x.name: x for x in graph.outputs} + assert inputs["past_key_values.0.conv_state"].dtype == dtype + assert inputs["past_key_values.0.ssm_state"].dtype == ir.DataType.FLOAT + assert list(inputs["past_key_values.0.ssm_state"].shape)[1:] == [4, 8, 8] + assert outputs["present.0.ssm_state"].dtype == ir.DataType.FLOAT + assert inputs["past_key_values.2.key"].dtype == dtype + assert not any("recurrent_state" in name for name in inputs) + # Every statically typed MatMul/Mul has homogeneous input types, including + # FP32 routing and accumulation. ONNX has no implicit mixed-type promotion. + for node in graph: + if node.op_type in {"MatMul", "Mul", "Add"}: + dtypes = {x.dtype for x in node.inputs if x is not None and x.dtype is not None} + assert len(dtypes) <= 1, (node.name, dtypes) + + +def test_expert_routing_relu2_and_sigmoid_matches_independent_reference(tmp_path): + config = _tiny_config() + block = NemotronHMoEBlock(config) + builder, op, graph = create_test_builder() + value = create_test_input(builder, "hidden", [2, 3, 32]) + builder.add_output(block(op, value), "out") + model = ir.Model(graph, ir_version=11) + rng = np.random.default_rng(19) + weights = { + name: torch.from_numpy(rng.normal(0, 0.12, list(x.shape)).astype(np.float32)) + for name, x in graph.initializers.items() + if x.const_value is None + } + # Nonzero correction bias must affect selection, not the mixture weights. + weights["gate.e_score_correction_bias"] = torch.tensor([0.8, -0.4, 0.3, -0.6]) + apply_weights(model, weights) + path = tmp_path / "moe.onnx" + ir.save(model, path) + session = ort.InferenceSession(str(path), providers=["CPUExecutionProvider"]) + hidden = torch.from_numpy(rng.normal(size=(2, 3, 32)).astype(np.float32)) + scores = torch.sigmoid(hidden @ weights["gate.weight"].T) + selected = torch.topk(scores + weights["gate.e_score_correction_bias"], 2, dim=-1).indices + mixture = scores.gather(-1, selected) + mixture = mixture / (mixture.sum(-1, keepdim=True) + 1e-20) * 2.5 + expected = torch.zeros_like(hidden) + for idx in range(4): + up = torch.relu(hidden @ weights[f"experts.{idx}.up_proj.weight"].T).square() + expert = up @ weights[f"experts.{idx}.down_proj.weight"].T + mass = (mixture * (selected == idx)).sum(-1, keepdim=True) + expected += expert * mass + shared = torch.relu(hidden @ weights["shared_experts.up_proj.weight"].T).square() + expected += shared @ weights["shared_experts.down_proj.weight"].T + actual = session.run(None, {"hidden": hidden.numpy()})[0] + np.testing.assert_allclose(actual, expected.numpy(), atol=1e-6, rtol=1e-5) + + +def test_router_correction_bias_remains_float32(): + block = NemotronHMoEBlock(_tiny_config()) + _cast_module_dtype(block, ir.DataType.BFLOAT16) + assert block.gate.e_score_correction_bias.dtype == ir.DataType.FLOAT + assert block.gate.weight.dtype == ir.DataType.BFLOAT16 + + +def test_mtp_and_duplicate_source_weights_are_not_silently_dropped(): + module = NemotronHCausalLMModel(_tiny_config()) + with pytest.raises(NotImplementedError, match="MTP tensors"): + module.preprocess_weights({"mtp.layers.0.hnorm.weight": torch.ones(32)}) + with pytest.raises(ValueError, match="Duplicate"): + module.preprocess_weights( + { + "backbone.embeddings.weight": torch.ones(48, 32), + "model.embeddings.weight": torch.ones(48, 32), + } + ) + with pytest.raises(ValueError, match="out of range"): + _rename_nemotron_h_weight("backbone.layers.3.mixer.q_proj.weight", ["mamba2"]) + + +def test_reject_unimplemented_grouped_expert_routing(): + with pytest.raises(NotImplementedError, match="grouped expert"): + NemotronHConfig.from_transformers(SimpleNamespace(**{**_LIGHTNING, "n_group": 2})) diff --git a/src/mobius/models/nemotron_h.py b/src/mobius/models/nemotron_h.py index e16b1587e..120c22982 100644 --- a/src/mobius/models/nemotron_h.py +++ b/src/mobius/models/nemotron_h.py @@ -233,6 +233,7 @@ def __init__( self.weight = nn.Parameter([num_experts, hidden_size]) # Correction bias for expert selection (loaded from checkpoint) self.e_score_correction_bias = nn.Parameter([num_experts]) + setattr(self.e_score_correction_bias, "_keep_float32", True) def forward(self, op: OpBuilder, hidden_states: ir.Value): # Cast to float32 for numerical stability (eps=1e-20 underflows @@ -240,14 +241,14 @@ def forward(self, op: OpBuilder, hidden_states: ir.Value): # in NemotronHTopkRouter.forward and never casts back. hidden_states = op.Cast(hidden_states, to=1) # FLOAT32 - weight_t = op.Transpose(self.weight, perm=[1, 0]) + weight_t = op.Transpose(op.Cast(self.weight, to=1), perm=[1, 0]) router_logits = op.MatMul(hidden_states, weight_t) # Sigmoid probabilities (these become the final routing weights) probs = op.Sigmoid(router_logits) # Add correction bias for expert selection only - choice_scores = op.Add(probs, self.e_score_correction_bias) + choice_scores = op.Add(probs, op.Cast(self.e_score_correction_bias, to=1)) # Select top-k experts based on biased scores k = op.Constant(value_ints=[self.top_k]) @@ -262,9 +263,8 @@ def forward(self, op: OpBuilder, hidden_states: ir.Value): if self.routed_scaling_factor != 1.0: # noqa: RUF069 routing_weights = op.Mul(routing_weights, self.routed_scaling_factor) - # Keep routing_weights in float32 (matching HF which never casts back). - # The expert dispatch multiplies these with expert outputs, and ONNX - # type promotion handles the mixed-dtype matmul naturally. + # Keep routing weights in FP32; dispatch explicitly upcasts expert + # outputs because ONNX does not implicitly promote mixed input types. return routing_weights, selected_experts @@ -375,12 +375,15 @@ def forward(self, op: OpBuilder, hidden_states: ir.Value): weighted = op.Mul(routing_weights, match_float) # Sum matched routing weights across top_k dim → per-token weight weight = op.ReduceSum(weighted, [-1], keepdims=True) - contribution = op.Mul(expert_output, weight) + # HF accumulates routed contributions in FP32, then rounds once + # back to the expert input dtype before the shared-expert addition. + contribution = op.Mul(op.Cast(expert_output, to=1), weight) if result is None: result = contribution else: result = op.Add(result, contribution) + result = op.CastLike(result, hidden_states) # Optional latent projection back to hidden_size if self._has_latent: result = self.fc2_latent_proj(op, result) @@ -523,6 +526,11 @@ class NemotronHCausalLMModel(nn.Module): def __init__(self, config: NemotronHConfig): super().__init__() self.config = config + if config.num_nextn_predict_layers: + raise NotImplementedError( + "Nemotron-H NextN/MTP export is not implemented; auxiliary tensors " + "must not be silently skipped by a target-only graph." + ) self.model = _NemotronHTextModel(config) self.lm_head = Linear(config.hidden_size, config.vocab_size, bias=False) if config.tie_word_embeddings: @@ -580,6 +588,10 @@ def preprocess_weights( new_state_dict: dict[str, torch.Tensor] = {} for key, value in state_dict.items(): + if key.startswith("mtp."): + raise NotImplementedError( + "Nemotron-H MTP tensors cannot be silently discarded" + ) new_key = _rename_nemotron_h_weight(key, layer_types) # Split stacked 3D expert tensors into per-expert 2D weights. # HF stores experts.up_proj as (num_experts, inter, input) and @@ -592,8 +604,12 @@ def preprocess_weights( expert_key = ( new_key.rsplit("experts.", 1)[0] + f"experts.{i}.{suffix}.weight" ) + if expert_key in new_state_dict: + raise ValueError(f"Duplicate Nemotron-H weight target {expert_key!r}") new_state_dict[expert_key] = expert_weight else: + if new_key in new_state_dict: + raise ValueError(f"Duplicate Nemotron-H weight target {new_key!r}") new_state_dict[new_key] = value return new_state_dict @@ -645,7 +661,11 @@ def _rename_nemotron_h_weight(key: str, layer_types: list[str]) -> str: if m: layer_idx = int(m.group(1)) rest = m.group(2) - ltype = layer_types[layer_idx] if layer_idx < len(layer_types) else "full_attention" + if layer_idx >= len(layer_types): + raise ValueError(f"Nemotron-H checkpoint layer index out of range: {key!r}") + ltype = layer_types[layer_idx] + if ltype not in {"mamba2", "full_attention", "moe", "mlp"}: + raise ValueError(f"Unknown Nemotron-H layer type {ltype!r} for {key!r}") if rest.startswith("mixer."): mixer_rest = rest[len("mixer.") :] diff --git a/src/mobius/tasks/_cache_utils.py b/src/mobius/tasks/_cache_utils.py index da4162c4c..2e901a9f6 100644 --- a/src/mobius/tasks/_cache_utils.py +++ b/src/mobius/tasks/_cache_utils.py @@ -362,7 +362,7 @@ def _make_hybrid_cache_inputs( ) ssm_state = builder.input( f"{prefix}.{i}.ssm_state", - dtype=dtype, + dtype=getattr(config, "mamba_ssm_dtype", None) or dtype, shape=[batch, mamba2_n_heads, mamba2_d_state, mamba2_d_head], ) pairs.append((conv_state, ssm_state)) From 2e74e227c31ed8d658309582d40d460804e37693 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 7 Oct 2026 16:26:22 -0700 Subject: [PATCH 03/15] Follow FP32 parameter marker convention and clean up regression setup Address automated lint/code-quality findings by using the established direct dtype-marker assignment with its intentional dynamic-attribute type annotation, and avoid binding unused checkpoint-fixture return values. Numerical behavior and represented code relationships are unchanged. Targeted model/public-builder regressions and the B010 lint gate pass. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> --- .../integrations/transformers/_nemotron_h_weights_test.py | 8 ++++---- src/mobius/models/nemotron_h.py | 2 +- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/src/mobius/integrations/transformers/_nemotron_h_weights_test.py b/src/mobius/integrations/transformers/_nemotron_h_weights_test.py index d38474042..df8b2e175 100644 --- a/src/mobius/integrations/transformers/_nemotron_h_weights_test.py +++ b/src/mobius/integrations/transformers/_nemotron_h_weights_test.py @@ -145,7 +145,7 @@ def _write_mixed_checkpoint(tmp_path): def test_public_local_build_reconstructs_all_projection_families(tmp_path, monkeypatch): - _raw, expected = _write_mixed_checkpoint(tmp_path) + _, expected = _write_mixed_checkpoint(tmp_path) monkeypatch.setattr( _builder, "_download_weights", @@ -202,7 +202,7 @@ def test_floating_per_expert_checkpoint_uses_exact_streaming_mapping(tmp_path): def test_cli_explicit_dequantize_writes_dense_not_native_format(tmp_path): - _raw, _expected = _write_mixed_checkpoint(tmp_path) + _write_mixed_checkpoint(tmp_path) output = tmp_path / "cli-export" main( [ @@ -275,7 +275,7 @@ def test_reconstruction_policy_does_not_mutate_source_config(): def test_mtp_sources_fail_closed_even_when_config_omits_nextn(tmp_path, monkeypatch): - _raw, _expected = _write_mixed_checkpoint(tmp_path) + _write_mixed_checkpoint(tmp_path) path = tmp_path / "model.safetensors" weights = safetensors.torch.load_file(path) weights["mtp.layers.0.hnorm.weight"] = torch.ones(32, dtype=torch.bfloat16) @@ -290,7 +290,7 @@ def test_mtp_sources_fail_closed_even_when_config_omits_nextn(tmp_path, monkeypa @pytest.mark.parametrize("change", ["missing", "wrong-shape", "nonfinite", "undeclared"]) def test_cache_scale_contract_fails_before_initializer_binding(tmp_path, change): - raw, _expected = _write_mixed_checkpoint(tmp_path) + raw, _ = _write_mixed_checkpoint(tmp_path) path = tmp_path / "model.safetensors" weights = safetensors.torch.load_file(path) key = "backbone.layers.2.mixer.k_proj.k_scale" diff --git a/src/mobius/models/nemotron_h.py b/src/mobius/models/nemotron_h.py index 120c22982..a4a9205a0 100644 --- a/src/mobius/models/nemotron_h.py +++ b/src/mobius/models/nemotron_h.py @@ -233,7 +233,7 @@ def __init__( self.weight = nn.Parameter([num_experts, hidden_size]) # Correction bias for expert selection (loaded from checkpoint) self.e_score_correction_bias = nn.Parameter([num_experts]) - setattr(self.e_score_correction_bias, "_keep_float32", True) + self.e_score_correction_bias._keep_float32 = True # type: ignore[attr-defined] def forward(self, op: OpBuilder, hidden_states: ir.Value): # Cast to float32 for numerical stability (eps=1e-20 underflows From 734331cfc28f96c98292b27c0de87a33a1af5999 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 7 Oct 2026 16:37:45 -0700 Subject: [PATCH 04/15] Preserve Lightning's FP32 router source precision Keep router parameters FP32 before BF16 dtype conversion, not only during graph computation. The pinned NVFP4 checkpoint stores its unquantized router in FP32, while the BF16 variant stores it in BF16. Independent source fixtures and a non-BF16-representable value catch premature rounding across streaming and serialization paths. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> --- docs/nemotron-lightning-export.md | 6 +++++- .../transformers/_nemotron_h_weights_test.py | 14 +++++++++----- src/mobius/models/_nemotron_h_test.py | 4 ++-- src/mobius/models/nemotron_h.py | 1 + 4 files changed, 17 insertions(+), 8 deletions(-) diff --git a/docs/nemotron-lightning-export.md b/docs/nemotron-lightning-export.md index e4aa7be4f..a33af9ec1 100644 --- a/docs/nemotron-lightning-export.md +++ b/docs/nemotron-lightning-export.md @@ -73,7 +73,11 @@ This is **not a working full-Lightning command while the MTP gate remains**. report identify the reconstructed representation and lost quantization. - `models/nemotron_h.py` casts router weights/bias to FP32 and accumulates routed outputs in FP32 before the shared-expert path, retaining sigmoid selection, - unbiased routing weights and squared ReLU. All-expert graph unrolling remains; + unbiased routing weights and squared ReLU. It also retains FP32 router + parameters during dtype conversion: the pinned NVFP4 router weight is FP32, + unlike BF16's BF16 router weight; correction biases are FP32 in both. + Widening BF16 sources is exact, while rounding NVFP4's FP32 router to BF16 + would lose unquantized source precision. All-expert graph unrolling remains; this is not sparse top-6 performance or native QMoE support. - `NemotronHConfig` accepts Transformers' `linear_attention` vocabulary and preserves the checkpoint's FP32 SSM-cache dtype. The graph still uses diff --git a/src/mobius/integrations/transformers/_nemotron_h_weights_test.py b/src/mobius/integrations/transformers/_nemotron_h_weights_test.py index df8b2e175..7c9f12567 100644 --- a/src/mobius/integrations/transformers/_nemotron_h_weights_test.py +++ b/src/mobius/integrations/transformers/_nemotron_h_weights_test.py @@ -105,10 +105,10 @@ def _write_mixed_checkpoint(tmp_path): else: tensors[source] = torch.full( shape, - 0.125, - dtype=torch.float32 - if initializer.dtype == ir.DataType.FLOAT - else torch.bfloat16, + 0.123456 if source.endswith(".gate.weight") else 0.125, + # Pinned NVFP4 keeps both router tensors FP32; this source + # contract must not be inferred from the target graph's dtype. + dtype=torch.float32 if ".gate." in source else torch.bfloat16, ) expected[name] = tensors[source] raw["quantization_config"] = { @@ -172,8 +172,12 @@ def test_public_local_build_reconstructs_all_projection_families(tmp_path, monke assert loaded.metadata_props["mobius.source_weight_format"] == "modelopt-mixed-fp8-nvfp4" -def test_floating_per_expert_checkpoint_uses_exact_streaming_mapping(tmp_path): +@pytest.mark.parametrize("router_dtype", [torch.float32, torch.bfloat16]) +def test_floating_per_expert_checkpoint_uses_exact_streaming_mapping(tmp_path, router_dtype): raw, expected = _write_mixed_checkpoint(tmp_path) + expected["model.layers.1.moe.gate.weight"] = expected["model.layers.1.moe.gate.weight"].to( + router_dtype + ) del raw["quantization_config"] (tmp_path / "config.json").write_text(json.dumps(raw), encoding="utf-8") safetensors.torch.save_file( diff --git a/src/mobius/models/_nemotron_h_test.py b/src/mobius/models/_nemotron_h_test.py index 9112affce..d6a1a7598 100644 --- a/src/mobius/models/_nemotron_h_test.py +++ b/src/mobius/models/_nemotron_h_test.py @@ -213,11 +213,11 @@ def test_expert_routing_relu2_and_sigmoid_matches_independent_reference(tmp_path np.testing.assert_allclose(actual, expected.numpy(), atol=1e-6, rtol=1e-5) -def test_router_correction_bias_remains_float32(): +def test_router_parameters_remain_float32(): block = NemotronHMoEBlock(_tiny_config()) _cast_module_dtype(block, ir.DataType.BFLOAT16) assert block.gate.e_score_correction_bias.dtype == ir.DataType.FLOAT - assert block.gate.weight.dtype == ir.DataType.BFLOAT16 + assert block.gate.weight.dtype == ir.DataType.FLOAT def test_mtp_and_duplicate_source_weights_are_not_silently_dropped(): diff --git a/src/mobius/models/nemotron_h.py b/src/mobius/models/nemotron_h.py index a4a9205a0..27e41d814 100644 --- a/src/mobius/models/nemotron_h.py +++ b/src/mobius/models/nemotron_h.py @@ -231,6 +231,7 @@ def __init__( self.norm_topk_prob = norm_topk_prob self.routed_scaling_factor = routed_scaling_factor self.weight = nn.Parameter([num_experts, hidden_size]) + self.weight._keep_float32 = True # type: ignore[attr-defined] # Correction bias for expert selection (loaded from checkpoint) self.e_score_correction_bias = nn.Parameter([num_experts]) self.e_score_correction_bias._keep_float32 = True # type: ignore[attr-defined] From 1e5a2e78e59940ec9192c7defcf4645387d7bc7e Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 7 Oct 2026 17:07:26 -0700 Subject: [PATCH 05/15] Refresh Nemotron graph census after FP32 routing fixes CI's exact census regression exposed stale graph-only evidence: FP32 expert accumulation raises the raw/optimized counts to measured 43180/40109 nodes. Retain strict graph, initializer, MatMul and state census assertions and regenerate the catalog/matrix. Distinguish this local graph refresh from historical ORT 1.29 runtime inspection; no runtime support is claimed. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> --- docs/gguf-capability-catalog.md | 2 +- .../integrations/gguf/_quant_capabilities_test.py | 4 ++-- .../integrations/gguf/_runtime_blocker_evidence.py | 11 +++++++---- .../integrations/gguf/_runtime_evidence_test.py | 4 ++-- testdata/evidence/gguf_quantization_capabilities.json | 6 +++--- 5 files changed, 15 insertions(+), 12 deletions(-) diff --git a/docs/gguf-capability-catalog.md b/docs/gguf-capability-catalog.md index f0830c51b..58978c56b 100644 --- a/docs/gguf-capability-catalog.md +++ b/docs/gguf-capability-catalog.md @@ -55,7 +55,7 @@ Runtime support above is independent from tokenizer materialization support belo | Evidence ID | Pinned candidate | Bounded result | Withheld runtime claims | |---|---|---|---| -| `nemotron-h-moe-30b-iq2-xxs-runtime-blocker` | `bartowski/nvidia_Nemotron-3-Nano-30B-A3B-GGUF@1fc64d5b160654ec892df2708aa893b0e96e6491`
`nvidia_Nemotron-3-Nano-30B-A3B-IQ2_XXS.gguf`
18,010,755,296 B
`f3da710c046ce7cc6ff28a9b5f1a9153ac72e3f60603e51c7bb679d80716b58a` | config/tokenizer `nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16@bf77c3174f68ad409e1c2aa60daeb46e32d1c606`; GGUF tokenizer metadata `6089bcaf08b3fe0d49379ca7e85bd3c93e8705bac6130636425c159212971225`; result=blocked; 401 tensors / 31,577,940,288 parameters; ORT 1.29.0 / CPUExecutionProvider; ort-genai 0.15.2; The smallest model GGUF in the pinned 29-file repository revision is 18,010,755,296 bytes, above the 16 GiB bounded-artifact policy; explicit float16/float32 weights require 63,155,880,576/126,311,761,152 bytes.
The normal optimized CPU export has 37,142 nodes (40,167 before Mobius optimization) and 6,028 MatMul nodes because its truthful ReLU2 routed experts remain an ONNX loop. ORT 1.29 MoE/QMoE exposes ReLU but not ReLU2. QMoE has separate router_probs/router_weights for correction-biased selection with unbiased sigmoid mixing, while shared experts and optional latent projections can surround it; those are not fused-op blockers. The pinned 30B has no latent projection. QMoE does not support GGUF IQ2_XXS storage, and Mobius has no proven IQ2_XXS packer/kernel path.
ORT GenAI 0.15.2 discovers sparse/nonconsecutive KV and conv/recurrent slots, but derives recurrent_state names while this export uses ssm_state, does not beam-reorder recurrent state, and rejects nonzero recurrent-state rewind.
The GGUF tokenizer declares pre=pixtral, whose compiled llama.cpp behavior is not serialized; exact ORT tokenizer materialization is unavailable. | cached decode and deterministic generation, full-logit parity, package and report roundtrip, state replay, rollback, and reorder, tensor value closure | +| `nemotron-h-moe-30b-iq2-xxs-runtime-blocker` | `bartowski/nvidia_Nemotron-3-Nano-30B-A3B-GGUF@1fc64d5b160654ec892df2708aa893b0e96e6491`
`nvidia_Nemotron-3-Nano-30B-A3B-IQ2_XXS.gguf`
18,010,755,296 B
`f3da710c046ce7cc6ff28a9b5f1a9153ac72e3f60603e51c7bb679d80716b58a` | config/tokenizer `nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16@bf77c3174f68ad409e1c2aa60daeb46e32d1c606`; GGUF tokenizer metadata `6089bcaf08b3fe0d49379ca7e85bd3c93e8705bac6130636425c159212971225`; result=blocked; 401 tensors / 31,577,940,288 parameters; ORT 1.29.0 / CPUExecutionProvider; ort-genai 0.15.2; The smallest model GGUF in the pinned 29-file repository revision is 18,010,755,296 bytes, above the 16 GiB bounded-artifact policy; explicit float16/float32 weights require 63,155,880,576/126,311,761,152 bytes.
The normal optimized CPU export has 40,109 nodes (43,180 before Mobius optimization) and 6,028 MatMul nodes because its truthful ReLU2 routed experts remain an ONNX loop with FP32 routing/accumulation. This graph-only census was refreshed locally with ORT 1.30.0 installed; the pinned ORT 1.29 runtime inspection below is historical, not a new runtime result. ORT 1.29 MoE/QMoE exposes ReLU but not ReLU2. QMoE has separate router_probs/router_weights for correction-biased selection with unbiased sigmoid mixing, while shared experts and optional latent projections can surround it; those are not fused-op blockers. The pinned 30B has no latent projection. QMoE does not support GGUF IQ2_XXS storage, and Mobius has no proven IQ2_XXS packer/kernel path.
ORT GenAI 0.15.2 discovers sparse/nonconsecutive KV and conv/recurrent slots, but derives recurrent_state names while this export uses ssm_state, does not beam-reorder recurrent state, and rejects nonzero recurrent-state rewind.
The GGUF tokenizer declares pre=pixtral, whose compiled llama.cpp behavior is not serialized; exact ORT tokenizer materialization is unavailable. | cached decode and deterministic generation, full-logit parity, package and report roundtrip, state replay, rollback, and reorder, tensor value closure | | `glm-dsa-iq1-s-artifact-budget-blocker` | `unsloth/GLM-5.2-GGUF@abc55e72527792c6e77069c99b4cb7de16fa9f23`
`UD-IQ1_S/GLM-5.2-UD-IQ1_S-00001-of-00006.gguf` 9,423,744 B `46b6148389219ae45167cb8124fbb18ef7d432daf619b4faf9e06ea80d3f4777`
`UD-IQ1_S/GLM-5.2-UD-IQ1_S-00002-of-00006.gguf` 49,208,128,256 B `f2180207285e04fcaa5b8c53ba6e77ad5cc58666b6e7c6b04a5eded3fe8bef09`
`UD-IQ1_S/GLM-5.2-UD-IQ1_S-00003-of-00006.gguf` 49,684,417,024 B `b1c0c5a302cc8d5d9ea0bcd4467c01db72c26839f820f7e882079582ea0a8d2b`
`UD-IQ1_S/GLM-5.2-UD-IQ1_S-00004-of-00006.gguf` 49,396,052,864 B `a6a42da6975e29f89866dcde2956e9e50e6ea26635fb5063b74f3973f4f863b6`
`UD-IQ1_S/GLM-5.2-UD-IQ1_S-00005-of-00006.gguf` 49,246,275,936 B `a4a9851a50db533f21ef824e5d8038f04e6782e7d602d18e5fdd6643f68ccccb`
`UD-IQ1_S/GLM-5.2-UD-IQ1_S-00006-of-00006.gguf` 19,171,063,136 B `3b767f55df64e0432d52fcf1a14eb47a1ef3bbc91339e2ae220f38602237d7d7`
total 216,715,360,960 B | blocked; The smallest mainline GLM-5.2 GGUF is a 216,715,360,960-byte split set whose trunk also declares an unsupported routed DSA/MLA MTP block. | real-weight full-logit parity, runtime packaging, deterministic generation | | `minimax-m2-iq1-s-artifact-budget-blocker` | `mradermacher/MiniMax-M2-i1-GGUF@2d4f9b1a86d32ce4dfc47db312c8d6fcae8d7b37`
`MiniMax-M2.i1-IQ1_S.gguf` 46,514,882,176 B `7bae986e3cd380c28c6177d612fce1d52373241a0dfa13a6fca25de79abf15fb`
total 46,514,882,176 B | blocked; The smallest immutable public MiniMax-M2 GGUF is 46,514,882,176 bytes; real-weight parity cannot enter the 16 GiB bounded evidence set. | real-weight full-logit parity, runtime packaging, deterministic generation | | `mistral4-iq1-m-artifact-budget-blocker` | `unsloth/Mistral-Small-4-119B-2603-GGUF@bd93c721735aa32c035c0f19e738cb3371fd56ff`
`Mistral-Small-4-119B-2603-UD-IQ1_M.gguf` 32,306,941,632 B `40fcdee4869110938638c6b8bac253f442b196518d4623f4afdf4b885cd961c7`
total 32,306,941,632 B | blocked; The smallest immutable public Mistral4 GGUF is 32,306,941,632 bytes; real-weight parity cannot enter the 16 GiB bounded evidence set. | real-weight full-logit parity, runtime packaging, deterministic generation | diff --git a/src/mobius/integrations/gguf/_quant_capabilities_test.py b/src/mobius/integrations/gguf/_quant_capabilities_test.py index a2d6e24c5..e60fd5e8c 100644 --- a/src/mobius/integrations/gguf/_quant_capabilities_test.py +++ b/src/mobius/integrations/gguf/_quant_capabilities_test.py @@ -42,8 +42,8 @@ def test_runtime_blocker_candidate_is_metadata_only_and_not_budgeted_as_support( assert record["size"] > policy["max_selected_artifact_bytes"] assert record["size"] not in {artifact["size"] for artifact in selected_artifacts} assert record["tokenizer"]["revision"] == "bf77c3174f68ad409e1c2aa60daeb46e32d1c606" - assert record["graph"]["pre_optimization_node_count"] == 40_167 - assert record["graph"]["node_count"] == 37_142 + assert record["graph"]["pre_optimization_node_count"] == 43_180 + assert record["graph"]["node_count"] == 40_109 assert record["graph"]["initializer_count"] == 6_255 assert record["graph"]["matmul_count"] == 6_028 assert record["graph"]["state_slots"]["mamba2.ssm_state"] == 23 diff --git a/src/mobius/integrations/gguf/_runtime_blocker_evidence.py b/src/mobius/integrations/gguf/_runtime_blocker_evidence.py index 8727cb6ed..5d1f294a2 100644 --- a/src/mobius/integrations/gguf/_runtime_blocker_evidence.py +++ b/src/mobius/integrations/gguf/_runtime_blocker_evidence.py @@ -200,8 +200,8 @@ def __post_init__(self) -> None: expert_count=128, experts_per_token=6, layer_counts=(("full_attention", 6), ("mamba2", 23), ("moe", 23)), - pre_optimization_graph_node_count=40_167, - graph_node_count=37_142, + pre_optimization_graph_node_count=43_180, + graph_node_count=40_109, graph_initializer_count=6_255, graph_matmul_count=6_028, state_slots=( @@ -222,9 +222,12 @@ def __post_init__(self) -> None: "float16/float32 weights require 63,155,880,576/126,311,761,152 bytes." ), ( - "The normal optimized CPU export has 37,142 nodes (40,167 before Mobius " + "The normal optimized CPU export has 40,109 nodes (43,180 before Mobius " "optimization) and 6,028 MatMul nodes because its truthful ReLU2 routed " - "experts remain an ONNX loop. ORT 1.29 MoE/QMoE exposes ReLU but not ReLU2. " + "experts remain an ONNX loop with FP32 routing/accumulation. " + "This graph-only census was refreshed locally with ORT 1.30.0 installed; " + "the pinned ORT 1.29 runtime inspection below is historical, not a new " + "runtime result. ORT 1.29 MoE/QMoE exposes ReLU but not ReLU2. " "QMoE has separate router_probs/router_weights for correction-biased " "selection with unbiased sigmoid mixing, while shared experts and optional " "latent projections can surround it; those are not fused-op blockers. The " diff --git a/src/mobius/integrations/gguf/_runtime_evidence_test.py b/src/mobius/integrations/gguf/_runtime_evidence_test.py index 4f0c96ec7..160c51481 100644 --- a/src/mobius/integrations/gguf/_runtime_evidence_test.py +++ b/src/mobius/integrations/gguf/_runtime_evidence_test.py @@ -215,8 +215,8 @@ def test_nemotron_h_runtime_blocker_is_pinned_without_support_claim() -> None: assert record.runtime_schema_issue.endswith("/issues/605") assert _runtime_evidence.runtime_evidence(record.evidence_id) is None assert "full-logit parity" in record.withheld_checks - assert record.graph_node_count == 37_142 - assert record.pre_optimization_graph_node_count == 40_167 + assert record.graph_node_count == 40_109 + assert record.pre_optimization_graph_node_count == 43_180 assert "separate router_probs/router_weights" in record.blockers[1] assert "not fused-op blockers" in record.blockers[1] assert "has no latent projection" in record.blockers[1] diff --git a/testdata/evidence/gguf_quantization_capabilities.json b/testdata/evidence/gguf_quantization_capabilities.json index 8efaf1e53..bad8ebf7a 100644 --- a/testdata/evidence/gguf_quantization_capabilities.json +++ b/testdata/evidence/gguf_quantization_capabilities.json @@ -2842,7 +2842,7 @@ "architecture": "nemotron_h_moe", "blockers": [ "The smallest model GGUF in the pinned 29-file repository revision is 18,010,755,296 bytes, above the 16 GiB bounded-artifact policy; explicit float16/float32 weights require 63,155,880,576/126,311,761,152 bytes.", - "The normal optimized CPU export has 37,142 nodes (40,167 before Mobius optimization) and 6,028 MatMul nodes because its truthful ReLU2 routed experts remain an ONNX loop. ORT 1.29 MoE/QMoE exposes ReLU but not ReLU2. QMoE has separate router_probs/router_weights for correction-biased selection with unbiased sigmoid mixing, while shared experts and optional latent projections can surround it; those are not fused-op blockers. The pinned 30B has no latent projection. QMoE does not support GGUF IQ2_XXS storage, and Mobius has no proven IQ2_XXS packer/kernel path.", + "The normal optimized CPU export has 40,109 nodes (43,180 before Mobius optimization) and 6,028 MatMul nodes because its truthful ReLU2 routed experts remain an ONNX loop with FP32 routing/accumulation. This graph-only census was refreshed locally with ORT 1.30.0 installed; the pinned ORT 1.29 runtime inspection below is historical, not a new runtime result. ORT 1.29 MoE/QMoE exposes ReLU but not ReLU2. QMoE has separate router_probs/router_weights for correction-biased selection with unbiased sigmoid mixing, while shared experts and optional latent projections can surround it; those are not fused-op blockers. The pinned 30B has no latent projection. QMoE does not support GGUF IQ2_XXS storage, and Mobius has no proven IQ2_XXS packer/kernel path.", "ORT GenAI 0.15.2 discovers sparse/nonconsecutive KV and conv/recurrent slots, but derives recurrent_state names while this export uses ssm_state, does not beam-reorder recurrent state, and rejects nonzero recurrent-state rewind.", "The GGUF tokenizer declares pre=pixtral, whose compiled llama.cpp behavior is not serialized; exact ORT tokenizer materialization is unavailable." ], @@ -2863,8 +2863,8 @@ "graph": { "initializer_count": 6255, "matmul_count": 6028, - "node_count": 37142, - "pre_optimization_node_count": 40167, + "node_count": 40109, + "pre_optimization_node_count": 43180, "state_slots": { "attention.key": 6, "attention.value": 6, From 671d00020338e3c864f81ce28ff34ef9c6b5f1ab Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 7 Oct 2026 17:15:37 -0700 Subject: [PATCH 06/15] Add explicit Nemotron target-only export with named MTP omissions Wire the user-approved target_decoder_only API and --target-decoder-only CLI option through both pinned floating and opt-in reconstructed ModelOpt paths. Retain the source NextN config, default refusal and native NVFP4 block; persist variant provenance and every omitted mtp.* name, source shape/dtype and reason. Keep all non-MTP inventory/shape/scale failures strict and distinguish graph-only uninspected inventories. Cover API/CLI/default/native/errors/save/reload and update pinned unexecuted recipe candidates. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> --- docs/nemotron-lightning-export.md | 72 +++++--- examples/nemotron_lightning/recipe.json | 39 ++-- src/mobius/__main__.py | 32 +++- .../integrations/transformers/_builder.py | 60 ++++++- .../transformers/_nemotron_h_weights.py | 29 ++- .../transformers/_nemotron_h_weights_test.py | 169 ++++++++++++++++++ src/mobius/models/_nemotron_h_test.py | 12 ++ src/mobius/models/nemotron_h.py | 8 +- 8 files changed, 370 insertions(+), 51 deletions(-) diff --git a/docs/nemotron-lightning-export.md b/docs/nemotron-lightning-export.md index a33af9ec1..d5fa9d408 100644 --- a/docs/nemotron-lightning-export.md +++ b/docs/nemotron-lightning-export.md @@ -5,7 +5,8 @@ The [machine-readable recipe](../examples/nemotron_lightning/recipe.json) records immutable inputs and compatibility limits. Native NVFP4 remains unsupported. Explicit dense BF16 reconstruction is implemented and tested with miniature mixed-format checkpoints, **not** the complete 30B models. -The pinned checkpoints' NextN/MTP contract remains a blocking gate. +An explicitly selected **target-decoder-only** variant can omit NextN/MTP +with complete named reporting. Faithful MTP export remains unsupported. The original preflight host was Windows ARM64 with a Qualcomm Adreno X1-85, no NVIDIA/CUDA GPU or `nvidia-smi`, Python 3.12.10, 68,171,038,720 bytes RAM @@ -27,7 +28,7 @@ retrieved from the real checkpoints; no tensor payloads were downloaded. | Unique index shards | 14 | 52 | | Config context limit | 262,144 | 1,048,576 | | Source quantization | Unquantized BF16 | ModelOpt mixed FP8/NVFP4 | -| Current full-checkpoint classification | Blocked on NextN/MTP | Native unsupported; dense reconstruction blocked on NextN/MTP | +| Current classification | Explicit target-decoder-only; full MTP unsupported | Explicit reconstructed-BF16 target-decoder-only; native/full MTP unsupported | Both configs declare `NemotronHForCausalLM` / `nemotron_h`, 52 layers (23 Mamba, 23 MoE, 6 attention), 128 routed experts, top-6 routing, and one @@ -54,7 +55,8 @@ The current code adds an explicit Nemotron-H dense reconstruction route: `build(..., keep_quantized=False, dtype="bf16")`, corresponding to `--dequantize --dtype bf16`. Both options are required; implicit dtype, FP16, native requests and unsupported layouts fail before weight loading. -This is **not a working full-Lightning command while the MTP gate remains**. +For the pinned Lightning sources, also select `target_decoder_only=True` +or `--target-decoder-only`. Neither opt-in implies the other. - `integrations/modelopt/_config.py` validates exact FP8/NVFP4 module groups, static block-16 format, redundant `quantized_layers` declarations and FP8 @@ -96,17 +98,30 @@ quantization and FP8 KV-cache quantization.** It is neither native NVFP4 preservation nor re-quantized ORT INT4, and does not promise equivalence to the original quantized execution. -### NextN/MTP: fail closed +### NextN/MTP: explicit target-decoder-only, default fail closed Model construction rejects nonzero `num_nextn_predict_layers`; eager and streaming weight paths also reject `mtp.*` tensors even if a config omits that -field. No automatic target-only fallback or discarded auxiliary tensors are -allowed. The [inspected HuggingFace reference](https://github.com/huggingface/transformers/blob/a96730c8c97b8efbf35bbaf7f5da33ec99231a49/src/transformers/models/nemotron_h/modeling_nemotron_h.py) +field, **unless the public builder's separate target-only option is explicit**. +The approved `target_decoder_only=True` / `--target-decoder-only` variant +preserves source `num_nextn_predict_layers`, uses strict streaming for both +floating per-expert and reconstructed weights, and excludes only `mtp.*` keys. +Every omitted name, source shape/dtype and reason is persisted in +`weight-loading-report.json` and ONNX metadata. Unexpected non-MTP tensors, +missing target weights and scale errors still fail. A source declaring NextN +but containing no MTP inventory fails when loading. The eager preprocessor +still rejects MTP; it is not an unreported omission path. + +Graph names include `/target-decoder-only/`, metadata records +`mobius.export_variant=target-decoder-only` and `mobius.mtp_preserved=false`, +and logs warn about omissions. With `--no-weights`, inventory status is +`not_inspected`: no omitted-tensor coverage is claimed from config alone. +The [inspected HuggingFace reference](https://github.com/huggingface/transformers/blob/a96730c8c97b8efbf35bbaf7f5da33ec99231a49/src/transformers/models/nemotron_h/modeling_nemotron_h.py) ignores MTP; the [inspected vLLM MTP implementation](https://github.com/vllm-project/vllm/blob/ef63c23d35acccbba8e014fc88da19d541da38a9/vllm/model_executor/models/nemotron_h_mtp.py) appears to construct shared experts absent from the pinned MTP indexes. That discrepancy is not resolved by inventing missing weights. -A separately authorized, explicitly named target-only variant or a validated -faithful MTP graph is required before these complete checkpoints can build. +The target-only variant is not a resolution of that faithful MTP discrepancy, +nor a claim to preserve all checkpoint behavior. ## Reproduce metadata/source verification @@ -181,26 +196,35 @@ does not change the scope of evidence. | [`integrations/ort_genai/auto_export.py`](https://github.com/onnxruntime/mobius/blob/dc1798465f77cb2b7070e4d422479c2f890ee24b/src/mobius/integrations/ort_genai/auto_export.py) | Generic recurrent packaging expects paired `conv_state` / `recurrent_state`; Nemotron-H describes `conv_state` / `ssm_state`. Actual graph names/shapes and runtime allocation are unverified, not evidence of a loadable package. | | [`__main__.py`](https://github.com/onnxruntime/mobius/blob/dc1798465f77cb2b7070e4d422479c2f890ee24b/src/mobius/__main__.py) | The candidate below uses source-backed CLI syntax. `--revision` is only valid with `--model`, not local `--config`. | -## Full BF16 procedure: blocked, NOT_RUN +## Target-decoder-only procedures: implemented, full-size NOT_RUN Before any full export, obtain a suitable development host, explicit resource budget, immutable dependency lock and CUDA/ORT version pairing. No package installation, model download or GPU allocation is part of this preflight. The publisher's 80 GB-class GPU deployment target is not an export-memory -measurement. Only after the MTP gate is resolved, stage a complete read-only +measurement. With explicit target-decoder-only scope, stage a complete read-only snapshot at the pinned HF revision: config, tokenizer and every indexed weight shard. Verify shard identities, tokenizer assets, disk capacity and loader peak memory separately. Do not enable `trust_remote_code` or fetch mutable `main`. -The earlier graph-only candidate is no longer presented as runnable: -the pinned config declares one NextN layer, and current construction explicitly -rejects that unsupported contract before loading shards. The recipe therefore -contains no approved full-checkpoint build command. Do not edit the config to -zero NextN or remove MTP tensors to bypass this gate. Any future command must -use a fresh output directory and retain logs, exit status, exact versions and -resource measurements; it must omit `--runtime ort-genai` until that separate -runtime contract is validated. +The following commands are **unexecuted full-size target-only candidates**, +not approved GPU work or successful exports. The local snapshot paths must +contain the verified pinned variants respectively. Do not edit config to zero +NextN or remove tensors before loading: the opt-in loader reports all omissions. + +```powershell +mobius build --config C:\private\lightning-bf16-pinned ` + --output C:\private\lightning-target-decoder-only-bf16 ` + --ep cuda --dtype bf16 --target-decoder-only --external-data onnx +mobius build --config C:\private\lightning-nvfp4-pinned ` + --output C:\private\lightning-target-decoder-only-reconstructed-bf16 ` + --ep cuda --dtype bf16 --target-decoder-only --dequantize --external-data onnx +``` + +Use fresh output directories and retain logs, exit status, exact versions and +resource measurements. Omit `--runtime ort-genai` until that separate runtime +contract is validated. `--revision` is only valid with `--model`, not `--config`. Before claiming success, account for every checkpoint tensor, including MTP; validate graph/external-data references; compare source/export @@ -218,10 +242,11 @@ throughput claims. The pinned ModelOpt config describes mixed FP8 Mamba projections and FP4 expert weights with group size 16, plus static FP8 activations and KV cache. The generic/native exporter still rejects this method. The opt-in dense -reconstruction route above handles these physical formats but remains blocked -on full Lightning MTP accounting. **There is no supported native NVFP4 or -full-Lightning build command in this recipe.** Do not download 52 shards merely -to reproduce that known contract blocker. +reconstruction route handles these physical formats, and the independent +target-only option accounts for omitted MTP tensors. **There is no native NVFP4 +or faithful full-MTP build command in this recipe.** The candidate above +produces a differently represented, target-only artifact; do not download 52 +shards merely to reproduce the known native/full-MTP blockers. Native preservation additionally requires faithful ReLU2/routing kernel semantics, runtime capability checks and real-weight parity. @@ -235,7 +260,8 @@ must not be labeled a preserved NVFP4 export. The added suites use reduced configurations, independently specified routing and E2M1/E4M3 values, and local synthetic safetensors files. They exercise public builder/CLI reconstruction, save/reload, every layer/projection family, source -accounting, malformed/changing headers, dtype/cache contracts and MTP refusal. +accounting, malformed/changing headers, dtype/cache contracts, default MTP +refusal and explicit omission/variant provenance. Existing CPU random-weight HuggingFace Nemotron parity tests also run without model downloads. These are not full-checkpoint L4/L5 evidence. diff --git a/examples/nemotron_lightning/recipe.json b/examples/nemotron_lightning/recipe.json index a6acf245d..8a7e41c22 100644 --- a/examples/nemotron_lightning/recipe.json +++ b/examples/nemotron_lightning/recipe.json @@ -2,7 +2,7 @@ "schema_version": 1, "kind": "development_export_recipe", "checked_on": "2026-10-07", - "status": "CPU_REGRESSION_TESTED_FULL_CHECKPOINT_BLOCKED", + "status": "CPU_REGRESSION_TESTED_TARGET_DECODER_ONLY_FULL_SIZE_NOT_RUN", "export_status": "NOT_RUN", "certified": false, "exporter": { @@ -62,18 +62,25 @@ "max_position_embeddings": 262144, "checkpoint_tensor_bytes": 65842365568, "checkpoint_shards": 14, - "compatibility": "NEXTN_CONTRACT_BLOCKED", + "compatibility": "EXPLICIT_TARGET_DECODER_ONLY_FULL_MTP_UNSUPPORTED", "output_precision": "bf16", "checkpoint_mtp_tensors": 270, "checkpoint_non_mtp_tensors": 6243, - "proposed_graph_export_argv": null, - "command_status": "BLOCKED_MTP", + "export_variant": "target-decoder-only", + "mtp_preserved": false, + "proposed_graph_export_argv": [ + "mobius", "build", "--config", "C:\\private\\lightning-bf16-pinned", + "--output", "C:\\private\\lightning-target-decoder-only-bf16", + "--ep", "cuda", "--dtype", "bf16", "--target-decoder-only", + "--external-data", "onnx" + ], + "command_status": "TARGET_DECODER_ONLY_NOT_RUN", "runtime_packaging_argv": null, "blockers": [ "No NVIDIA CUDA GPU; executed validation uses an isolated x64 CPU environment.", "A suitable host, dependency lock and resource budget have not been validated.", "Full Lightning checkpoint weight coverage and source/export parity are untested.", - "NextN/MTP construction and mtp.* tensors explicitly fail closed; no target-only fallback.", + "Faithful NextN/MTP is unsupported; the explicit target-only variant names all omitted mtp.* tensors. No default fallback.", "Mamba conv_state/ssm_state packaging is not verified against ORT-GenAI." ] }, @@ -103,24 +110,32 @@ "quant_method": "modelopt", "weight_formats": ["FP8", "NVFP4"], "nvfp4_group_size": 16, - "compatibility": "NATIVE_UNSUPPORTED_DENSE_RECONSTRUCTION_NEXTN_BLOCKED", + "compatibility": "NATIVE_UNSUPPORTED_EXPLICIT_RECONSTRUCTED_BF16_TARGET_ONLY", "output_precision": null, - "proposed_graph_export_argv": null, - "command_status": "BLOCKED_MTP_AND_NATIVE_FORMAT", + "export_variant": "target-decoder-only", + "mtp_preserved": false, + "proposed_graph_export_argv": [ + "mobius", "build", "--config", "C:\\private\\lightning-nvfp4-pinned", + "--output", "C:\\private\\lightning-target-decoder-only-reconstructed-bf16", + "--ep", "cuda", "--dtype", "bf16", "--target-decoder-only", "--dequantize", + "--external-data", "onnx" + ], + "command_status": "RECONSTRUCTED_BF16_TARGET_DECODER_ONLY_NOT_RUN", "checkpoint_mtp_tensors": 270, "source_format_counts": {"FP8": 46, "NVFP4": 5935}, "source_fp8_input_scales": 46, "source_fp8_kv_cache_scales": 12, "dense_reconstruction": { "representation": "dense_bf16", - "api_options": {"keep_quantized": false, "dtype": "bf16"}, - "cli_options": ["--dequantize", "--dtype", "bf16"], + "api_options": {"keep_quantized": false, "dtype": "bf16", "target_decoder_only": true}, + "cli_options": ["--dequantize", "--dtype", "bf16", "--target-decoder-only"], "validation": "SYNTHETIC_CPU_BUILDER_CLI_SERIALIZATION_NUMERIC_TESTS", "full_checkpoint_export": "NOT_RUN", "native_weight_format_preserved": false, "activation_quantization_preserved": false, "kv_cache_quantization_preserved": false, - "requires_validated_nextn_contract": true + "mtp_preserved": false, + "omission_reporting": "every mtp.* name, source shape/dtype and explicit reason" }, "runtime_packaging_argv": null, "blockers": [ @@ -128,7 +143,7 @@ "Mixed-format loading/reconstruction is CPU-tested, not native routed-expert NVFP4 emission.", "Portable ReLU2/sigmoid routing is CPU-tested; a generic QMoE is not equivalent.", "BF16 reconstruction or ORT INT4 quantization would be different artifacts, not preserved NVFP4.", - "The BF16 host, checkpoint parity, MTP and runtime packaging blockers also apply." + "The BF16 host, checkpoint parity, faithful MTP and runtime packaging blockers also apply." ] } ] diff --git a/src/mobius/__main__.py b/src/mobius/__main__.py index f9e4d4834..8835d9ac3 100644 --- a/src/mobius/__main__.py +++ b/src/mobius/__main__.py @@ -229,6 +229,9 @@ def _resolve_static_cache_task(model_type: str) -> ModelTask: load_weights = not args.no_weights keep_quantized = not getattr(args, "dequantize", False) + target_decoder_only = getattr(args, "target_decoder_only", False) + if target_decoder_only and args.component: + raise SystemExit("Error: --target-decoder-only cannot be combined with --component.") task: str | ModelTask | None = args.task # FP8 KV cache: resolve the optional per-layer scale file up front so both @@ -307,7 +310,13 @@ def _resolve_static_cache_task(model_type: str) -> ModelTask: # that flag only applies to transformers decoder exports, so we let the # central build() validation reject a diffusers/unsupported repo rather # than silently exporting a diffusion pipeline and ignoring the flag. - if args.model and not args.config and not args.text_only and not is_personaplex: + if ( + args.model + and not args.config + and not args.text_only + and not target_decoder_only + and not is_personaplex + ): pipeline_index = _load_diffusers_pipeline_index(args.model, revision=revision) if pipeline_index is not None: if input_sampling_rate is not None or bwe_sampling_rate is not None: @@ -342,6 +351,8 @@ def _resolve_static_cache_task(model_type: str) -> ModelTask: # 'owner/repo:model.nemo'). Routes to the NeMo import path; reuses the # standard build args (--dtype, --ep, --external-data) and save logic. if args.model and args.model.endswith(".nemo"): + if target_decoder_only: + raise SystemExit("Error: --target-decoder-only is only supported for Nemotron-H.") from mobius.integrations.nemo import build_from_nemo print(f"Detected NeMo archive: {args.model}") @@ -381,6 +392,7 @@ def _resolve_static_cache_task(model_type: str) -> ModelTask: keep_quantized=keep_quantized, input_sampling_rate=input_sampling_rate, bwe_sampling_rate=bwe_sampling_rate, + target_decoder_only=target_decoder_only, ) _save_package(pkg, output_dir, args, optimize, component_filter) return @@ -388,6 +400,10 @@ def _resolve_static_cache_task(model_type: str) -> ModelTask: from mobius.models.reuse import _build_reuse, _is_reuse_checkpoint if _is_reuse_checkpoint(config_path): + if target_decoder_only: + raise SystemExit( + "Error: --target-decoder-only is only supported for Nemotron-H." + ) if task not in (None, "speech-enhancement"): from mobius.tasks import SpeechEnhancementTask @@ -447,6 +463,7 @@ def _resolve_static_cache_task(model_type: str) -> ModelTask: keep_quantized=keep_quantized, input_sampling_rate=input_sampling_rate, bwe_sampling_rate=bwe_sampling_rate, + target_decoder_only=target_decoder_only, ) else: model_id_or_path = args.model @@ -482,6 +499,7 @@ def _resolve_static_cache_task(model_type: str) -> ModelTask: keep_quantized=keep_quantized, input_sampling_rate=input_sampling_rate, bwe_sampling_rate=bwe_sampling_rate, + target_decoder_only=target_decoder_only, ) _save_package(pkg, output_dir, args, optimize, component_filter) @@ -1434,12 +1452,22 @@ def build_parser() -> argparse.ArgumentParser: "without it all layers use a unit scale of 1.0." ), ) + build_parser.add_argument( + "--target-decoder-only", + action="store_true", + help=( + "Explicit Nemotron-H variant without NextN/MTP. Names every omitted " + "mtp.* tensor in the loading report; does not preserve MTP. No default " + "fallback and no native NVFP4 support." + ), + ) build_parser.add_argument( "--dequantize", action="store_true", help=( "Explicitly reconstruct supported compressed-tensors and GPT-OSS " - "MXFP4 weights as dense floating point. Dense GPT-OSS reconstruction " + "MXFP4 weights as dense floating point, or Nemotron-H ModelOpt weights " + "with explicit --dtype bf16. Dense GPT-OSS reconstruction " "is eager and memory-intensive; the default preserves MXFP4 with " "bounded native streaming." ), diff --git a/src/mobius/integrations/transformers/_builder.py b/src/mobius/integrations/transformers/_builder.py index f5f0e92da..821bbb7aa 100644 --- a/src/mobius/integrations/transformers/_builder.py +++ b/src/mobius/integrations/transformers/_builder.py @@ -7,6 +7,7 @@ import copy import dataclasses +import json import logging import pathlib from typing import Any @@ -380,6 +381,7 @@ def build_transformers_model( export_paged_attention: bool = False, input_sampling_rate: int | None = None, bwe_sampling_rate: int | None = None, + target_decoder_only: bool = False, ) -> ModelPackage: """Build a model package from a Transformers checkpoint. @@ -395,6 +397,10 @@ def build_transformers_model( ``keep_quantized`` preserves supported compressed-tensors checkpoints in their native block-weight representation. Set it to ``False`` only to request explicit dense reconstruction. + + ``target_decoder_only`` explicitly exports Nemotron-H without NextN/MTP. + Source config facts remain unchanged; the strict streaming report names every + omitted ``mtp.*`` tensor. This is not faithful MTP or a default fallback. """ if input_sampling_rate is not None and bwe_sampling_rate is not None: raise ValueError("input_sampling_rate and bwe_sampling_rate are mutually exclusive") @@ -420,6 +426,7 @@ def build_transformers_model( "export_paged_attention": export_paged_attention, "input_sampling_rate": input_sampling_rate is not None, "bwe_sampling_rate": bwe_sampling_rate is not None, + "target_decoder_only": target_decoder_only, } selected = sorted(name for name, enabled in unsupported.items() if enabled) if selected: @@ -488,6 +495,8 @@ def build_transformers_model( trust_remote_code=trust_remote_code, ) if hf_config is None or (loaded_from_raw_json and hf_config.model_type not in registry): + if target_decoder_only: + raise ValueError("target_decoder_only is only supported for Nemotron-H") from mobius.models.reuse import _build_reuse, _is_reuse_checkpoint if module_class is None and _is_reuse_checkpoint(model_id, detection_revision): @@ -553,6 +562,8 @@ def build_transformers_model( ) hf_config, parent_config, model_type = _select_primary_config(hf_config) + if target_decoder_only and model_type != "nemotron_h": + raise ValueError("target_decoder_only is only supported for Nemotron-H") hf_config, parent_config, modelopt_source = _prepare_nemotron_modelopt_source( hf_config, parent_config, keep_quantized=keep_quantized, dtype=dtype ) @@ -690,7 +701,15 @@ def build_transformers_model( model_type=model_type, hf_config=parent_config, ) - model_module = module_class(config) + model_module: nn.Module + if target_decoder_only: + from mobius.models.nemotron_h import NemotronHCausalLMModel + + if module_class is not NemotronHCausalLMModel: + raise ValueError("target_decoder_only requires the Nemotron-H model class") + model_module = NemotronHCausalLMModel(config, target_decoder_only=True) + else: + model_module = module_class(config) if dequantize_gptoss_mxfp4: model_module._dequantize_mxfp4_checkpoint = True attach_hf_component_sources( @@ -714,6 +733,15 @@ def build_transformers_model( ) for name, model in package.items(): model.graph.name = f"{graph_source_name}/{name}" + if target_decoder_only: + model.graph.name = f"{graph_source_name}/target-decoder-only/{name}" + model.metadata_props["mobius.export_variant"] = "target-decoder-only" + model.metadata_props["mobius.source_num_nextn_predict_layers"] = str( + config.num_nextn_predict_layers + ) + model.metadata_props["mobius.mtp_preserved"] = "false" + model.metadata_props["mobius.mtp_inventory_status"] = "not_inspected" + model.metadata_props["mobius.source_revision"] = revision or "local-or-unpinned" if modelopt_source is not None: model.metadata_props["mobius.storage_policy"] = ( "explicit-dense-bf16-reconstruction" @@ -742,7 +770,7 @@ def build_transformers_model( if load_weights: _reject_unsupported_affine_qwen4(model_type, config) - if modelopt_source is not None: + if modelopt_source is not None or target_decoder_only: from mobius.integrations.transformers._nemotron_h_weights import ( build_nemotron_h_streaming_plan, ) @@ -753,18 +781,34 @@ def build_transformers_model( next(iter(package.values())), model_id, lambda index, initializers: build_nemotron_h_streaming_plan( - config, index, initializers, modelopt=modelopt_source + config, + index, + initializers, + modelopt=modelopt_source, + target_decoder_only=target_decoder_only, ), revision=revision, ) package.weight_loading_report = report for model in package.values(): model.metadata_props["mobius.weight_loading_status"] = "lazy-bound" - logger.warning( - "Explicitly reconstructing ModelOpt weights to dense BF16; native " - "NVFP4/FP8 storage, FP8 activations and FP8 KV-cache quantization " - "are not preserved." - ) + if target_decoder_only: + model.metadata_props["mobius.mtp_inventory_status"] = "inspected" + model.metadata_props["mobius.omitted_mtp_tensors"] = json.dumps( + report["omitted_mtp_tensors"], sort_keys=True + ) + if target_decoder_only: + logger.warning( + "Explicit target-decoder-only export omits %s named NextN/MTP " + "tensors; see weight-loading-report.json. This is not MTP support.", + report["omitted_mtp_tensor_count"], + ) + if modelopt_source is not None: + logger.warning( + "Explicitly reconstructing ModelOpt weights to dense BF16; native " + "NVFP4/FP8 storage, FP8 activations and FP8 KV-cache quantization " + "are not preserved." + ) elif is_gptoss_mxfp4_source and keep_quantized: from mobius.integrations.transformers._gptoss_weights import ( stream_gptoss_mxfp4_safetensors_to_package, diff --git a/src/mobius/integrations/transformers/_nemotron_h_weights.py b/src/mobius/integrations/transformers/_nemotron_h_weights.py index 982d24dfd..2fd15ee97 100644 --- a/src/mobius/integrations/transformers/_nemotron_h_weights.py +++ b/src/mobius/integrations/transformers/_nemotron_h_weights.py @@ -27,16 +27,24 @@ def build_nemotron_h_streaming_plan( initializers: Mapping[str, ir.Value], *, modelopt: ModelOptConfig | None = None, + target_decoder_only: bool = False, ) -> StreamingWeightPlan: """Map exact per-expert checkpoint keys and account for every source tensor.""" - if any(name.startswith("mtp.") for name in key_index): + mtp_names = sorted(name for name in key_index if name.startswith("mtp.")) + if (config.num_nextn_predict_layers or mtp_names) and not target_decoder_only: raise NotImplementedError( "Nemotron-H NextN/MTP tensors require a validated MTP export contract; " - "they must not be silently discarded." + "they must not be silently discarded. Use explicit target_decoder_only=True " + "for the separately labeled target decoder variant." ) + if target_decoder_only and config.num_nextn_predict_layers and not mtp_names: + raise ValueError("Source config declares NextN/MTP but no mtp.* tensors were found") + omitted = dict.fromkeys( + mtp_names, "explicit target-decoder-only variant excludes NextN/MTP" + ) targets: dict[str, StreamingWeightSource] = {} constants: dict[str, torch.Tensor] = {} - ignored: dict[str, str] = {} + ignored = omitted quantized_modules: set[str] = set() activation_modules = modelopt.activation_modules if modelopt is not None else frozenset() cache_scales = set() @@ -55,6 +63,8 @@ def build_nemotron_h_streaming_plan( validate_modelopt_scale(scale, name, scalar=True) constants[name] = scale for source_name in key_index: + if source_name in omitted: + continue if source_name in cache_scales: continue if source_name.endswith((".weight_scale", ".weight_scale_2", ".input_scale")): @@ -102,6 +112,19 @@ def build_nemotron_h_streaming_plan( ignored=ignored, constants=constants, report={ + "export_variant": "target-decoder-only" if target_decoder_only else "standard", + "source_num_nextn_predict_layers": config.num_nextn_predict_layers, + "mtp_preserved": False, + "mtp_inventory_status": "inspected", + "omitted_mtp_tensor_count": len(omitted), + "omitted_mtp_tensors": { + name: { + "reason": reason, + "shape": key_index[name][1], + "dtype": key_index[name][2], + } + for name, reason in omitted.items() + }, "source_weight_format": modelopt.source_format if modelopt else "floating", "storage_policy": "explicit-dense-bf16-reconstruction" if modelopt else "dense", "native_nvfp4": False, diff --git a/src/mobius/integrations/transformers/_nemotron_h_weights_test.py b/src/mobius/integrations/transformers/_nemotron_h_weights_test.py index 7c9f12567..4f77a4597 100644 --- a/src/mobius/integrations/transformers/_nemotron_h_weights_test.py +++ b/src/mobius/integrations/transformers/_nemotron_h_weights_test.py @@ -310,3 +310,172 @@ def test_cache_scale_contract_fails_before_initializer_binding(tmp_path, change) safetensors.torch.save_file(weights, path) with pytest.raises(ValueError): build(str(tmp_path), keep_quantized=False, dtype="bf16") + + +def _write_mtp_checkpoint(tmp_path, *, mixed): + raw, expected = _write_mixed_checkpoint(tmp_path) + path = tmp_path / "model.safetensors" + weights = safetensors.torch.load_file(path) + if not mixed: + del raw["quantization_config"] + weights = {_source_name(name): tensor for name, tensor in expected.items()} + # Omission accounting includes all storage/scale kinds, not only weights. + mtp = { + "mtp.layers.0.hnorm.weight": torch.ones(32, dtype=torch.bfloat16), + "mtp.layers.0.mixer.weight": torch.zeros(4, 16, dtype=torch.uint8), + "mtp.layers.0.mixer.weight_scale": torch.ones(4, 2).to(torch.float8_e4m3fn), + "mtp.layers.0.mixer.weight_scale_2": torch.tensor(0.25), + } + weights.update(mtp) + raw["num_nextn_predict_layers"] = 1 + (tmp_path / "config.json").write_text(json.dumps(raw), encoding="utf-8") + safetensors.torch.save_file(weights, path) + return raw, expected, mtp + + +@pytest.mark.parametrize("mixed", [False, True]) +def test_explicit_target_variant_accounts_for_every_mtp_tensor(tmp_path, mixed, monkeypatch): + raw, expected, mtp = _write_mtp_checkpoint(tmp_path, mixed=mixed) + source = SimpleNamespace(**raw) + monkeypatch.setattr(_builder, "_load_transformers_config", lambda *a, **k: (source, True)) + monkeypatch.setattr( + _builder, "_download_weights", lambda *a, **k: pytest.fail("must use strict streaming") + ) + options = {"dtype": "bf16", "keep_quantized": not mixed} + with pytest.raises(NotImplementedError, match="NextN/MTP"): + build(str(tmp_path), **options) + package = build(str(tmp_path), target_decoder_only=True, **options) + assert source.num_nextn_predict_layers == 1 + assert source.__dict__ == raw + report = package.weight_loading_report + assert report["export_variant"] == "target-decoder-only" + assert report["source_num_nextn_predict_layers"] == 1 + assert report["mtp_preserved"] is False + assert report["omitted_mtp_tensor_count"] == len(mtp) + assert report["ignored_tensors"] == len(mtp) + assert set(report["omitted_mtp_tensors"]) == set(mtp) + for name, tensor in mtp.items(): + entry = report["omitted_mtp_tensors"][name] + assert entry["shape"] == list(tensor.shape) + assert "explicit target-decoder-only" in entry["reason"] + model = package["model"] + assert "/target-decoder-only/" in model.graph.name + assert model.metadata_props["mobius.source_num_nextn_predict_layers"] == "1" + assert model.metadata_props["mobius.mtp_inventory_status"] == "inspected" + assert not any(name.startswith("mtp.") for name in model.graph.initializers) + for name, tensor in expected.items(): + actual = model.graph.initializers[name].const_value.numpy().astype("float32") + torch.testing.assert_close(torch.from_numpy(actual), tensor.float(), rtol=0, atol=0) + output = tmp_path / "target-decoder-only" + package.save(output, external_data="onnx") + loaded = ir.load(output / "model.onnx") + assert loaded.metadata_props["mobius.export_variant"] == "target-decoder-only" + assert set(json.loads(loaded.metadata_props["mobius.omitted_mtp_tensors"])) == set(mtp) + saved_report = json.loads((output / "weight-loading-report.json").read_text()) + assert saved_report["omitted_mtp_tensors"] == report["omitted_mtp_tensors"] + + +@pytest.mark.parametrize("mixed", [False, True]) +def test_target_variant_cli_and_graph_only_provenance(tmp_path, mixed): + raw, _, mtp = _write_mtp_checkpoint(tmp_path, mixed=mixed) + graph_only = build( + str(tmp_path), + dtype="bf16", + keep_quantized=not mixed, + target_decoder_only=True, + load_weights=False, + )["model"] + assert graph_only.metadata_props["mobius.mtp_inventory_status"] == "not_inspected" + assert "mobius.omitted_mtp_tensors" not in graph_only.metadata_props + output = tmp_path / "cli-target" + main( + [ + "build", + "--config", + str(tmp_path), + "--output", + str(output), + "--ep", + "cpu", + "--dtype", + "bf16", + "--target-decoder-only", + "--external-data", + "onnx", + *(["--dequantize"] if mixed else []), + ] + ) + loaded = ir.load(output / "model.onnx") + assert loaded.metadata_props["mobius.export_variant"] == "target-decoder-only" + report = json.loads((output / "weight-loading-report.json").read_text()) + assert set(report["omitted_mtp_tensors"]) == set(mtp) + assert json.loads((tmp_path / "config.json").read_text()) == raw + + +@pytest.mark.parametrize( + "extra", ["unexpected.weight", "unexpected.weight_scale", "mtpx.weight"] +) +def test_target_variant_does_not_ignore_non_mtp_extras(tmp_path, extra): + _write_mtp_checkpoint(tmp_path, mixed=True) + path = tmp_path / "model.safetensors" + tensors = safetensors.torch.load_file(path) + tensors[extra] = torch.ones(1) + safetensors.torch.save_file(tensors, path) + with pytest.raises(ValueError, match=r"Unmapped|unclassified"): + build(str(tmp_path), dtype="bf16", keep_quantized=False, target_decoder_only=True) + + +def test_target_variant_still_blocks_native_modelopt(tmp_path): + _write_mtp_checkpoint(tmp_path, mixed=True) + with pytest.raises(NotImplementedError, match="Native"): + build(str(tmp_path), dtype="bf16", target_decoder_only=True) + + +def test_target_variant_only_accepts_nemotron(tmp_path): + (tmp_path / "config.json").write_text(json.dumps({"model_type": "qwen2"})) + with pytest.raises(ValueError, match="only supported for Nemotron"): + build(str(tmp_path), load_weights=False, target_decoder_only=True) + + +def test_declared_nextn_requires_source_inventory_when_loading(tmp_path): + raw, _ = _write_mixed_checkpoint(tmp_path) + raw["num_nextn_predict_layers"] = 1 + (tmp_path / "config.json").write_text(json.dumps(raw)) + with pytest.raises(ValueError, match="no mtp"): + build(str(tmp_path), dtype="bf16", keep_quantized=False, target_decoder_only=True) + + +def test_target_variant_forwards_cli_model_and_revision(tmp_path, monkeypatch): + from mobius import __main__ as cli + + received = {} + + def capture(source, **options): + received.update(source=source, **options) + return {} + + monkeypatch.setattr(cli, "build", capture) + monkeypatch.setattr(cli, "_save_package", lambda *args: None) + monkeypatch.setattr( + "mobius.integrations.diffusers._builder._load_diffusers_pipeline_index", + lambda *args, **kwargs: pytest.fail("target variant cannot bypass build policy"), + ) + main( + [ + "build", + "--model", + "nvidia/lightning", + "--revision", + "a" * 40, + "--output", + str(tmp_path / "target"), + "--dtype", + "bf16", + "--target-decoder-only", + "--dequantize", + ] + ) + assert received["source"] == "nvidia/lightning" + assert received["target_decoder_only"] is True + assert received["keep_quantized"] is False + assert received["revision"] == "a" * 40 diff --git a/src/mobius/models/_nemotron_h_test.py b/src/mobius/models/_nemotron_h_test.py index d6a1a7598..c2430ab35 100644 --- a/src/mobius/models/_nemotron_h_test.py +++ b/src/mobius/models/_nemotron_h_test.py @@ -235,6 +235,18 @@ def test_mtp_and_duplicate_source_weights_are_not_silently_dropped(): _rename_nemotron_h_weight("backbone.layers.3.mixer.q_proj.weight", ["mamba2"]) +def test_target_decoder_constructor_preserves_source_nextn_contract(): + config = dataclasses.replace(_tiny_config(), num_nextn_predict_layers=1) + with pytest.raises(NotImplementedError, match="NextN/MTP"): + NemotronHCausalLMModel(config) + module = NemotronHCausalLMModel(config, target_decoder_only=True) + assert module.config is config + assert config.num_nextn_predict_layers == 1 + # The public opt-in uses the accounted streaming path, never eager omission. + with pytest.raises(NotImplementedError, match="MTP tensors"): + module.preprocess_weights({"mtp.layers.0.hnorm.weight": torch.ones(32)}) + + def test_reject_unimplemented_grouped_expert_routing(): with pytest.raises(NotImplementedError, match="grouped expert"): NemotronHConfig.from_transformers(SimpleNamespace(**{**_LIGHTNING, "n_group": 2})) diff --git a/src/mobius/models/nemotron_h.py b/src/mobius/models/nemotron_h.py index 27e41d814..ea38f7009 100644 --- a/src/mobius/models/nemotron_h.py +++ b/src/mobius/models/nemotron_h.py @@ -524,13 +524,15 @@ class NemotronHCausalLMModel(nn.Module): category: str = "Hybrid SSM+Attention" config_class: type = NemotronHConfig - def __init__(self, config: NemotronHConfig): + def __init__(self, config: NemotronHConfig, *, target_decoder_only: bool = False): super().__init__() self.config = config - if config.num_nextn_predict_layers: + self.target_decoder_only = target_decoder_only + if config.num_nextn_predict_layers and not target_decoder_only: raise NotImplementedError( "Nemotron-H NextN/MTP export is not implemented; auxiliary tensors " - "must not be silently skipped by a target-only graph." + "must not be silently skipped. Explicit target_decoder_only=True " + "(--target-decoder-only) exports a separately labeled variant." ) self.model = _NemotronHTextModel(config) self.lm_head = Linear(config.hidden_size, config.vocab_size, bias=False) From d6f970893a11a4882f699e22f4f1d0216681e75a Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 7 Oct 2026 21:46:07 -0700 Subject: [PATCH 07/15] Protect Nemotron lazy checkpoint sources during package save Resolve shard paths once and register lazy shard/config/index files plus snapshot and resolved blob parent directories with ModelPackage's existing save guard. Share the directory-identity helper with GPT-OSS instead of duplicating it. Add source-directory, indexed/unindexed, floating/reconstructed and hard-link collision regressions that prove no source mutation before lazy reads, with successful fresh-output serialization. Correct reconstructed NVFP4 recipe output precision to bf16 without claiming execution. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> --- docs/nemotron-lightning-export.md | 4 ++ examples/nemotron_lightning/recipe.json | 2 +- src/mobius/integrations/_weight_loading.py | 10 +++ .../integrations/transformers/_builder.py | 18 ++++++ .../transformers/_gptoss_weights.py | 12 +--- .../transformers/_nemotron_h_weights_test.py | 64 +++++++++++++++++++ 6 files changed, 98 insertions(+), 12 deletions(-) diff --git a/docs/nemotron-lightning-export.md b/docs/nemotron-lightning-export.md index d5fa9d408..68489cb43 100644 --- a/docs/nemotron-lightning-export.md +++ b/docs/nemotron-lightning-export.md @@ -73,6 +73,10 @@ or `--target-decoder-only`. Neither opt-in implies the other. and missing scales. FP8 activation and KV-cache scales are explicitly validated/accounted for, **not simulated**. Export metadata and the loading report identify the reconstructed representation and lost quantization. + The package registers lazy shard/config/index sources and both snapshot and + resolved blob directories. Saving into source checkpoint directories or + onto source-file aliases is rejected before serialization; use a fresh, + separate output directory. - `models/nemotron_h.py` casts router weights/bias to FP32 and accumulates routed outputs in FP32 before the shared-expert path, retaining sigmoid selection, unbiased routing weights and squared ReLU. It also retains FP32 router diff --git a/examples/nemotron_lightning/recipe.json b/examples/nemotron_lightning/recipe.json index 8a7e41c22..d3d8eb4e7 100644 --- a/examples/nemotron_lightning/recipe.json +++ b/examples/nemotron_lightning/recipe.json @@ -111,7 +111,7 @@ "weight_formats": ["FP8", "NVFP4"], "nvfp4_group_size": 16, "compatibility": "NATIVE_UNSUPPORTED_EXPLICIT_RECONSTRUCTED_BF16_TARGET_ONLY", - "output_precision": null, + "output_precision": "bf16", "export_variant": "target-decoder-only", "mtp_preserved": false, "proposed_graph_export_argv": [ diff --git a/src/mobius/integrations/_weight_loading.py b/src/mobius/integrations/_weight_loading.py index acac583ee..f4fb99847 100644 --- a/src/mobius/integrations/_weight_loading.py +++ b/src/mobius/integrations/_weight_loading.py @@ -476,6 +476,16 @@ def _dequantize_fp8_weights(state_dict: dict[str, torch.Tensor]) -> dict[str, to return {k: v for k, v in result.items() if not any(k.endswith(s) for s in aux_suffixes)} +def _lazy_safetensors_source_parent_aliases(paths: list[str]) -> frozenset[pathlib.Path]: + """Return both directory identities from which lazy shards may be read.""" + # Keep HF snapshot directories as well as the blob directories reached by + # shard symlinks so neither can be reused as an external-data destination. + return frozenset( + {pathlib.Path(path).parent.resolve() for path in paths} + | {pathlib.Path(path).resolve().parent for path in paths} + ) + + def _resolve_shard_paths(model_id: str, revision: str | None = None) -> list[str]: """Resolve local safetensors shard paths for the streaming loader. diff --git a/src/mobius/integrations/transformers/_builder.py b/src/mobius/integrations/transformers/_builder.py index 821bbb7aa..afe49c663 100644 --- a/src/mobius/integrations/transformers/_builder.py +++ b/src/mobius/integrations/transformers/_builder.py @@ -26,6 +26,8 @@ from mobius._registry import registry from mobius.integrations._weight_loading import ( _download_weights, + _lazy_safetensors_source_parent_aliases, + _resolve_shard_paths, stream_preprocessed_safetensors_to_model, stream_qdq_safetensors_to_model, ) @@ -777,6 +779,7 @@ def build_transformers_model( if len(package) != 1: raise ValueError("Nemotron-H dense streaming requires one decoder component") + paths = _resolve_shard_paths(model_id, revision) report = stream_preprocessed_safetensors_to_model( next(iter(package.values())), model_id, @@ -788,6 +791,21 @@ def build_transformers_model( target_decoder_only=target_decoder_only, ), revision=revision, + _resolved_paths=paths, + ) + source_roots = {pathlib.Path(path).parent for path in paths} + if pathlib.Path(model_id).is_dir(): + source_roots.add(pathlib.Path(model_id)) + control_paths = [ + str(path) + for root in source_roots + for name in ("config.json", "model.safetensors.index.json") + if (path := root / name).is_file() + ] + all_source_paths = [*paths, *control_paths] + package.register_lazy_source_artifacts( + files=all_source_paths, + directories=_lazy_safetensors_source_parent_aliases(all_source_paths), ) package.weight_loading_report = report for model in package.values(): diff --git a/src/mobius/integrations/transformers/_gptoss_weights.py b/src/mobius/integrations/transformers/_gptoss_weights.py index 214c8a11e..62e24cda8 100644 --- a/src/mobius/integrations/transformers/_gptoss_weights.py +++ b/src/mobius/integrations/transformers/_gptoss_weights.py @@ -20,6 +20,7 @@ StreamingTransformedWeightSource, StreamingWeightPlan, StreamingWeightSource, + _lazy_safetensors_source_parent_aliases, _local_weight_paths, _resolve_shard_paths, stream_preprocessed_safetensors_to_model, @@ -36,17 +37,6 @@ _FLOAT_DTYPES = frozenset({"BF16", "F16", "F32"}) -def _lazy_safetensors_source_parent_aliases(paths: list[str]) -> frozenset[pathlib.Path]: - """Return both directory identities from which lazy shards may be read.""" - # HF snapshots can contain shard symlinks into the blob cache. Resolving - # the parent preserves the snapshot directory, while resolving the whole - # shard path follows the symlink and identifies the blob directory. - return frozenset( - {pathlib.Path(path).parent.resolve() for path in paths} - | {pathlib.Path(path).resolve().parent for path in paths} - ) - - def _repack_blocks(tensor: torch.Tensor, _source_name: str) -> torch.Tensor: return repack_gptoss_mxfp4_blocks(tensor) diff --git a/src/mobius/integrations/transformers/_nemotron_h_weights_test.py b/src/mobius/integrations/transformers/_nemotron_h_weights_test.py index 4f77a4597..13e8f64cd 100644 --- a/src/mobius/integrations/transformers/_nemotron_h_weights_test.py +++ b/src/mobius/integrations/transformers/_nemotron_h_weights_test.py @@ -479,3 +479,67 @@ def capture(source, **options): assert received["target_decoder_only"] is True assert received["keep_quantized"] is False assert received["revision"] == "a" * 40 + + +@pytest.mark.parametrize("mixed", [False, True]) +@pytest.mark.parametrize("indexed", [False, True]) +def test_lazy_source_guard_prevents_checkpoint_overwrite( + tmp_path, monkeypatch, mixed, indexed +): + _write_mtp_checkpoint(tmp_path, mixed=mixed) + if indexed: + shard = tmp_path / "model-00001-of-00001.safetensors" + (tmp_path / "model.safetensors").rename(shard) + names = safetensors.torch.load_file(shard) + (tmp_path / "model.safetensors.index.json").write_text( + json.dumps({"weight_map": dict.fromkeys(names, shard.name)}) + ) + before = {path: path.read_bytes() for path in tmp_path.iterdir() if path.is_file()} + resolve = _builder._resolve_shard_paths + calls = [] + + def resolve_once(*args): + calls.append(args) + return resolve(*args) + + monkeypatch.setattr(_builder, "_resolve_shard_paths", resolve_once) + package = build( + str(tmp_path), dtype="bf16", keep_quantized=not mixed, target_decoder_only=True + ) + assert len(calls) == 1 + assert set(package._native_streaming_source_files) == set(before) + assert tmp_path.resolve() in package._native_streaming_source_directories + with pytest.raises(ValueError, match="source checkpoint directory"): + package.save(tmp_path, external_data="safetensors") + assert {path: path.read_bytes() for path in before} == before + output = tmp_path / "fresh-output" + package.save(output, external_data="safetensors") + loaded = ir.load(output / "model.onnx") + assert loaded.metadata_props["mobius.export_variant"] == "target-decoder-only" + + +@pytest.mark.parametrize( + "source_file", ["model.safetensors", "config.json", "model.safetensors.index.json"] +) +def test_lazy_source_guard_rejects_hardlinked_output_before_materializing( + tmp_path, monkeypatch, source_file +): + _write_mixed_checkpoint(tmp_path) + if source_file.endswith(".index.json"): + weights = safetensors.torch.load_file(tmp_path / "model.safetensors") + (tmp_path / source_file).write_text( + json.dumps({"weight_map": dict.fromkeys(weights, "model.safetensors")}) + ) + package = build(str(tmp_path), dtype="bf16", keep_quantized=False) + source = tmp_path / source_file + before = source.read_bytes() + output = tmp_path / "alias-output" + output.mkdir() + (output / "model.onnx.data").hardlink_to(source) + monkeypatch.setattr( + "mobius.integrations._weight_loading._materialize_preprocessed_source", + lambda *args: pytest.fail("alias rejection must precede lazy reads"), + ) + with pytest.raises(ValueError, match="aliases lazy source file"): + package.save(output, external_data="onnx") + assert source.read_bytes() == before From 897a0e05bc60625088ab100b8e3da0c6550de032 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 8 Oct 2026 10:56:38 -0700 Subject: [PATCH 08/15] Add explicit Nemotron Lightning target and MTP graph exports Implement a separate nemotron-h-mtp task that preserves the pinned Lightning attention/MoE NextN weights, exports a post-final-norm target seed, and gives the draft its own KV cache. Keep default refusal, target-only omission reporting, and ModelOpt reconstruction opt-ins unchanged. Account for both graph source inventories and shared embedding/head copies with strict streaming and lazy-source protection. Validate tiny independent CPU equations across prefill, cached decode, padding, source precisions and the actual target-to-draft bridge. Record pinned header/config evidence and explicit full-checkpoint, BF16/CUDA and generation-runtime limitations; this is not speculative runtime integration or native NVFP4. Signed-off-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> --- docs/nemotron-lightning-export.md | 116 +++- examples/nemotron_lightning/recipe.json | 61 +- src/mobius/_configs/_base.py | 13 +- src/mobius/_model_package.py | 1 + .../integrations/ort_genai/auto_export.py | 20 +- .../integrations/transformers/_builder.py | 85 ++- .../transformers/_nemotron_h_weights.py | 112 +++- src/mobius/models/__init__.py | 3 + src/mobius/models/_nemotron_h_mtp_test.py | 549 ++++++++++++++++++ src/mobius/models/nemotron_h.py | 11 +- src/mobius/models/nemotron_h_mtp.py | 171 ++++++ src/mobius/tasks/__init__.py | 3 + src/mobius/tasks/_nemotron_h_mtp.py | 81 +++ 13 files changed, 1190 insertions(+), 36 deletions(-) create mode 100644 src/mobius/models/_nemotron_h_mtp_test.py create mode 100644 src/mobius/models/nemotron_h_mtp.py create mode 100644 src/mobius/tasks/_nemotron_h_mtp.py diff --git a/docs/nemotron-lightning-export.md b/docs/nemotron-lightning-export.md index 68489cb43..eff3bb4d4 100644 --- a/docs/nemotron-lightning-export.md +++ b/docs/nemotron-lightning-export.md @@ -6,7 +6,9 @@ records immutable inputs and compatibility limits. Native NVFP4 remains unsupported. Explicit dense BF16 reconstruction is implemented and tested with miniature mixed-format checkpoints, **not** the complete 30B models. An explicitly selected **target-decoder-only** variant can omit NextN/MTP -with complete named reporting. Faithful MTP export remains unsupported. +with complete named reporting. A separate explicit `nemotron-h-mtp` task now +exports target and draft graphs with strict MTP coverage and tiny independent +CPU parity. This is **graph support**, not a speculative-generation runtime. The original preflight host was Windows ARM64 with a Qualcomm Adreno X1-85, no NVIDIA/CUDA GPU or `nvidia-smi`, Python 3.12.10, 68,171,038,720 bytes RAM @@ -28,7 +30,7 @@ retrieved from the real checkpoints; no tensor payloads were downloaded. | Unique index shards | 14 | 52 | | Config context limit | 262,144 | 1,048,576 | | Source quantization | Unquantized BF16 | ModelOpt mixed FP8/NVFP4 | -| Current classification | Explicit target-decoder-only; full MTP unsupported | Explicit reconstructed-BF16 target-decoder-only; native/full MTP unsupported | +| Current classification | Explicit target-only or target+MTP graphs; full weighted MTP untested | Explicit reconstructed-BF16 target-only or target+MTP graphs; native NVFP4 unsupported | Both configs declare `NemotronHForCausalLM` / `nemotron_h`, 52 layers (23 Mamba, 23 MoE, 6 attention), 128 routed experts, top-6 routing, and one @@ -56,7 +58,9 @@ The current code adds an explicit Nemotron-H dense reconstruction route: `--dequantize --dtype bf16`. Both options are required; implicit dtype, FP16, native requests and unsupported layouts fail before weight loading. For the pinned Lightning sources, also select `target_decoder_only=True` -or `--target-decoder-only`. Neither opt-in implies the other. +/ `--target-decoder-only`, or the separate `task="nemotron-h-mtp"` / +`--task nemotron-h-mtp` route. Neither MTP selection nor target-only selection +implies reconstruction; both remain explicit. - `integrations/modelopt/_config.py` validates exact FP8/NVFP4 module groups, static block-16 format, redundant `quantized_layers` declarations and FP8 @@ -102,11 +106,11 @@ quantization and FP8 KV-cache quantization.** It is neither native NVFP4 preservation nor re-quantized ORT INT4, and does not promise equivalence to the original quantized execution. -### NextN/MTP: explicit target-decoder-only, default fail closed +### NextN/MTP: default fail closed, two explicit export contracts Model construction rejects nonzero `num_nextn_predict_layers`; eager and streaming weight paths also reject `mtp.*` tensors even if a config omits that -field, **unless the public builder's separate target-only option is explicit**. +field, **unless the public builder's separate target-only option or MTP task is explicit**. The approved `target_decoder_only=True` / `--target-decoder-only` variant preserves source `num_nextn_predict_layers`, uses strict streaming for both floating per-expert and reconstructed weights, and excludes only `mtp.*` keys. @@ -122,10 +126,104 @@ and logs warn about omissions. With `--no-weights`, inventory status is `not_inspected`: no omitted-tensor coverage is claimed from config alone. The [inspected HuggingFace reference](https://github.com/huggingface/transformers/blob/a96730c8c97b8efbf35bbaf7f5da33ec99231a49/src/transformers/models/nemotron_h/modeling_nemotron_h.py) ignores MTP; the [inspected vLLM MTP implementation](https://github.com/vllm-project/vllm/blob/ef63c23d35acccbba8e014fc88da19d541da38a9/vllm/model_executor/models/nemotron_h_mtp.py) -appears to construct shared experts absent from the pinned MTP indexes. -That discrepancy is not resolved by inventing missing weights. -The target-only variant is not a resolution of that faithful MTP discrepancy, -nor a claim to preserve all checkpoint behavior. +and its [parent layer definitions](https://github.com/vllm-project/vllm/blob/ef63c23d35acccbba8e014fc88da19d541da38a9/vllm/model_executor/models/nemotron_h.py) +provide the MTP reference. **The earlier shared-expert-absence statement was +incorrect:** both pinned indexes contain +`mtp.layers.1.mixer.shared_experts.up_proj.weight` and +`mtp.layers.1.mixer.shared_experts.down_proj.weight`. +Independent bounded reads of BF16 shard 14 and NVFP4 shard 52 headers confirm +their shapes are `[3712,2688]` and `[2688,3712]`. No missing weights or +architectural substitutions are needed. + +### Explicit target + Lightning MTP graph package + +`build(..., task="nemotron-h-mtp")` produces `decoder/model.onnx` and +`mtp/model.onnx`. The default constructor guard and existing target-only +contract are unchanged. In PowerShell, for the pinned BF16 checkpoint: + +```powershell +mobius build --model nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16 ` + --revision a9904d24bcc1d289a1950fa9d2b978c47cf903b9 ` + --task nemotron-h-mtp --dtype bf16 --ep cuda ` + --external-data onnx --output C:\private\lightning-target-and-mtp +``` + +For the NVFP4 model, use its pinned revision above and **also** add +`--dequantize --dtype bf16`; native requests still fail. These full-size +commands are **NOT_RUN here**. They are not instructions to overwrite the +checkpoint directory, nor evidence of H100 parity. + +The source architecture is one NextN prediction step containing two physical +blocks, `attention` then `moe`. The first block independently RMS-normalizes +the next-token embedding and target hidden state, concatenates them in that +order, projects `2H -> H`, then applies pre-norm GQA and residual addition. +The second applies pre-norm sigmoid/correction-bias top-6 ReLU2 MoE, including +its own always-active shared expert, residual addition, and final RMSNorm. +There is **no RoPE, position-ID input, Mamba block, or SSM state in MTP**. +Both graphs borrow the same source embedding and LM-head tensors; serialized +graphs contain separate copies rather than an undocumented runtime-sharing ABI. + +| Pinned MTP source | Shape / inventory | +| --- | --- | +| `mtp.layers.0.{enorm,hnorm,norm}.weight` | `[2688]` each | +| `mtp.layers.0.eh_proj.weight` | `[2688,5376]` | +| `mtp.layers.0.mixer.{q,k,v,o}_proj.weight` | `[4096,2688]`, `[256,2688]`, `[256,2688]`, `[2688,4096]` | +| `mtp.layers.1.{norm,final_layernorm}.weight` | `[2688]` each | +| `mtp.layers.1.mixer.gate.{weight,e_score_correction_bias}` | `[128,2688]`, `[128]` | +| `mtp.layers.1.mixer.experts.0..127.{up,down}_proj.weight` | 256 tensors: `[1856,2688]`, `[2688,1856]` | +| `mtp.layers.1.mixer.shared_experts.{up,down}_proj.weight` | `[3712,2688]`, `[2688,3712]` | + +Both independent headers have exactly **270 MTP tensors**, all BF16 except +the FP32 correction bias. The pinned NVFP4 config has **zero MTP quantized +module targets**, and its MTP index/header has **zero MTP scale tensors**. +The target's reconstructed LM head is reused for the MTP graph; this does not +simulate the original quantized activations or KV cache. + +**Decoder-to-MTP contract:** the explicit decoder emits `mtp_seed`, +`[batch,sequence,2688]`, **after the target's final `norm_f`**, exactly the +hidden tensor feeding its LM head. Pair `mtp_seed[:,i]` (target `h_i`) with +`input_ids[:,i+1]` (`t_(i+1)`) in the draft graph. Draft logits predict +`t_(i+2)`; the graph does not shift tokens itself. For prompt initialization, +feed known shifted pairs `ids[:,1:]` and `seed[:,:-1]`. For the next draft +step, pair the last target seed with the sampled next target token. Further +recursive drafting may pair `mtp_hidden` with a sampled draft token. + +The MTP graph consumes `input_ids` INT64 `[batch,sequence]`, `hidden_states` +in model dtype `[batch,sequence,2688]`, `attention_mask` INT64 +`[batch,past+sequence]`, and `past_key_values.0.{key,value}` in model dtype +`[batch,2,past,128]`. It emits `logits` `[batch,sequence,131072]`, +`mtp_hidden` `[batch,sequence,2688]`, and `present.0.{key,value}` with +`past+sequence` length. Initialize draft KV with zero sequence length; +**never borrow target-layer KV**. Target hybrid caches remain in their +existing sparse layer-index namespaces. + +Strict streaming validates the union of both graph source inventories before +returning the package. Every `mtp.*` source must bind to the draft graph; +missing, extra, duplicate-mapped and malformed tensors fail. The package +report records MTP names, shared table copies and per-component partition +counts, with **zero package-level MTP omissions**. Metadata with `--no-weights` +marks inventory `not_inspected`; graph construction alone is not tensor coverage. +Lazy source-file/directory protection applies to both graphs and package roots. + +**Validated scope:** tiny FP32 CPU independent-equation parity for prefill, +cached one-token and multi-token blocks, two rows with different padding, +FP32 and BF16 source weights, and the actual exported target-to-MTP bridge. +Tiny FP16 CPU draft prefill and cached decode also pass independent +FP32-reference comparison at `rtol=atol=1e-2`. +Tiny tests also cover mixed FP8/NVFP4 reconstruction with floating MTP and a +shared reconstructed head, exact graph/cache I/O, f32/f16/bf16 graph dtypes, +CLI serialization, strict source accounting and checkpoint-save protection. +BF16 execution/parity, full weighted 30B MTP, CUDA/H100, +performance, Olive requantization and L4/L5 real-weight goldens are **NOT_RUN**; +no full-model numerical or speedup claim is made. + +ORT-GenAI metadata can describe the external coordination contract and marks +it `runtime_unvalidated`. The exporter does **not** implement speculative +acceptance/verification. A generation runtime still needs shifted token/hidden +pairing, independent draft-KV management, target block verification/sampling, +accepted-prefix commit, and rejected-state restore/replay for **both** draft +KV and target KV/Mamba conv/SSM states. Existing generic acceptance/rollback +graphs do not establish Nemotron-H hybrid-state orchestration. ## Reproduce metadata/source verification diff --git a/examples/nemotron_lightning/recipe.json b/examples/nemotron_lightning/recipe.json index d3d8eb4e7..ec62b7553 100644 --- a/examples/nemotron_lightning/recipe.json +++ b/examples/nemotron_lightning/recipe.json @@ -1,10 +1,42 @@ { "schema_version": 1, "kind": "development_export_recipe", - "checked_on": "2026-10-07", - "status": "CPU_REGRESSION_TESTED_TARGET_DECODER_ONLY_FULL_SIZE_NOT_RUN", + "checked_on": "2026-10-08", + "status": "CPU_REGRESSION_TESTED_TARGET_AND_MTP_GRAPHS_FULL_WEIGHTED_MTP_NOT_RUN", "export_status": "NOT_RUN", "certified": false, + "mtp_export_contract": { + "task": "nemotron-h-mtp", + "explicit_opt_in": true, + "mutually_exclusive_with": "target_decoder_only", + "package_graphs": ["decoder/model.onnx", "mtp/model.onnx"], + "source_implementation": "https://github.com/vllm-project/vllm/blob/ef63c23d35acccbba8e014fc88da19d541da38a9/vllm/model_executor/models/nemotron_h_mtp.py", + "source_layer_implementation": "https://github.com/vllm-project/vllm/blob/ef63c23d35acccbba8e014fc88da19d541da38a9/vllm/model_executor/models/nemotron_h.py", + "physical_mtp_blocks": ["attention", "moe"], + "prediction_steps": 1, + "target_hidden_output": "mtp_seed", + "target_hidden_normalization": "post_final_norm", + "token_alignment": "target_h_i_with_token_t_i_plus_1", + "draft_prediction": "token_t_i_plus_2", + "target_shared_source_tables": ["backbone.embeddings.weight", "lm_head.weight"], + "graph_tables": "separate_copies_of_same_source", + "draft_cache": "independent_attention_kv_only", + "mtp_tensor_count": 270, + "mtp_source_dtype": "BF16_except_FP32_correction_bias", + "nvfp4_mtp_quantized_modules": 0, + "nvfp4_mtp_scale_tensors": 0, + "shared_expert_weights": "present_in_both_pinned_indexes_and_independent_headers", + "validation": "TINY_INDEPENDENT_CPU_REFERENCE_PREFILL_CACHED_DECODE_REAL_GRAPH_BRIDGE", + "full_weighted_mtp_export": "NOT_RUN", + "fp16_execution": "TINY_CPU_PREFILL_CACHED_DECODE_REFERENCE_TESTED", + "bf16_execution": "NOT_RUN", + "h100_cuda": "NOT_RUN", + "speculative_generation_runtime": "NOT_IMPLEMENTED", + "runtime_metadata_status": "runtime_unvalidated", + "native_nvfp4": "UNSUPPORTED", + "nvfp4_required_additional_options": ["--dequantize", "--dtype", "bf16"], + "default_mtp_guard": "UNCHANGED_FAIL_CLOSED" + }, "exporter": { "repository": "https://github.com/onnxruntime/mobius", "revision": "dc1798465f77cb2b7070e4d422479c2f890ee24b", @@ -62,9 +94,17 @@ "max_position_embeddings": 262144, "checkpoint_tensor_bytes": 65842365568, "checkpoint_shards": 14, - "compatibility": "EXPLICIT_TARGET_DECODER_ONLY_FULL_MTP_UNSUPPORTED", + "compatibility": "EXPLICIT_TARGET_ONLY_OR_TARGET_WITH_MTP_FULL_WEIGHTED_MTP_NOT_RUN", "output_precision": "bf16", "checkpoint_mtp_tensors": 270, + "mtp_header_probe": { + "shard": "model-00014-of-00014.safetensors", + "header_bytes": 33072, + "scope": "all_270_mtp_names_shapes_and_dtypes_no_tensor_payloads", + "status": "VERIFIED", + "shared_expert_projections": "PRESENT", + "mtp_scale_tensors": 0 + }, "checkpoint_non_mtp_tensors": 6243, "export_variant": "target-decoder-only", "mtp_preserved": false, @@ -80,7 +120,7 @@ "No NVIDIA CUDA GPU; executed validation uses an isolated x64 CPU environment.", "A suitable host, dependency lock and resource budget have not been validated.", "Full Lightning checkpoint weight coverage and source/export parity are untested.", - "Faithful NextN/MTP is unsupported; the explicit target-only variant names all omitted mtp.* tensors. No default fallback.", + "The separate explicit MTP task has tiny CPU graph parity, not full weighted MTP or speculative-generation validation. Target-only still names all omitted mtp.* tensors; no default fallback.", "Mamba conv_state/ssm_state packaging is not verified against ORT-GenAI." ] }, @@ -110,7 +150,7 @@ "quant_method": "modelopt", "weight_formats": ["FP8", "NVFP4"], "nvfp4_group_size": 16, - "compatibility": "NATIVE_UNSUPPORTED_EXPLICIT_RECONSTRUCTED_BF16_TARGET_ONLY", + "compatibility": "NATIVE_UNSUPPORTED_EXPLICIT_RECONSTRUCTED_BF16_TARGET_ONLY_OR_MTP_GRAPHS", "output_precision": "bf16", "export_variant": "target-decoder-only", "mtp_preserved": false, @@ -122,6 +162,15 @@ ], "command_status": "RECONSTRUCTED_BF16_TARGET_DECODER_ONLY_NOT_RUN", "checkpoint_mtp_tensors": 270, + "mtp_header_probe": { + "shard": "model-00052-of-00052.safetensors", + "header_bytes": 131824, + "scope": "all_270_mtp_names_shapes_and_dtypes_no_tensor_payloads", + "status": "VERIFIED_INDEPENDENTLY", + "shared_expert_projections": "PRESENT", + "mtp_scale_tensors": 0, + "mtp_quantized_config_modules": 0 + }, "source_format_counts": {"FP8": 46, "NVFP4": 5935}, "source_fp8_input_scales": 46, "source_fp8_kv_cache_scales": 12, @@ -143,7 +192,7 @@ "Mixed-format loading/reconstruction is CPU-tested, not native routed-expert NVFP4 emission.", "Portable ReLU2/sigmoid routing is CPU-tested; a generic QMoE is not equivalent.", "BF16 reconstruction or ORT INT4 quantization would be different artifacts, not preserved NVFP4.", - "The BF16 host, checkpoint parity, faithful MTP and runtime packaging blockers also apply." + "The BF16 host, full weighted MTP/checkpoint parity and runtime-orchestration blockers also apply." ] } ] diff --git a/src/mobius/_configs/_base.py b/src/mobius/_configs/_base.py index acebe31aa..bf6e2124b 100644 --- a/src/mobius/_configs/_base.py +++ b/src/mobius/_configs/_base.py @@ -4233,6 +4233,8 @@ class NemotronHConfig(ArchitectureConfig): mamba_time_step_min: float = 0.001 moe_latent_size: int | None = None mamba_ssm_dtype: ir.DataType | None = None + mtp_layers_block_type: list[str] | None = None + residual_in_fp32: bool = False @classmethod def from_transformers(cls, config, parent_config=None) -> NemotronHConfig: @@ -4307,7 +4309,7 @@ def from_transformers(cls, config, parent_config=None) -> NemotronHConfig: **base_fields, num_hidden_layers=n, layer_types=layers_block_type, - hidden_act="relu2", + hidden_act=getattr(config, "mlp_hidden_act", "relu2"), mamba_n_heads=mamba_n_heads, mamba_d_head=mamba_d_head, mamba_d_state=getattr(config, "ssm_state_size", 128), @@ -4320,6 +4322,15 @@ def from_transformers(cls, config, parent_config=None) -> NemotronHConfig: mamba_ssm_dtype=ssm_dtype, moe_latent_size=getattr(config, "moe_latent_size", None), shared_expert_intermediate_size=shared_expert_intermediate_size, + mtp_layers_block_type=( + [ + "attention" if kind == "full_attention" else kind + for kind in config.mtp_layers_block_type + ] + if getattr(config, "mtp_layers_block_type", None) is not None + else None + ), + residual_in_fp32=getattr(config, "residual_in_fp32", False), ) diff --git a/src/mobius/_model_package.py b/src/mobius/_model_package.py index e80dcdb00..0c31da750 100644 --- a/src/mobius/_model_package.py +++ b/src/mobius/_model_package.py @@ -218,6 +218,7 @@ def _validate_streaming_safetensors_destination( current_directory = Path(directory) current_layout = component_layout while True: + output_directories.add(current_directory.resolve()) component_directories = { _component_output_directory( current_directory, diff --git a/src/mobius/integrations/ort_genai/auto_export.py b/src/mobius/integrations/ort_genai/auto_export.py index 767936227..2d1908255 100644 --- a/src/mobius/integrations/ort_genai/auto_export.py +++ b/src/mobius/integrations/ort_genai/auto_export.py @@ -1943,18 +1943,26 @@ def _write_mtp_config(pkg: ModelPackage, directory: str) -> str | None: raise ValueError("MTP runtime metadata requires a target decoder graph") dedicated_embeddings = bool(getattr(proposer_config, "use_dedicated_embeddings", False)) dedicated_lm_head = bool(getattr(proposer_config, "use_dedicated_lm_head", False)) + conditioning = { + "target_hidden_output": "mtp_seed", + "target_hidden_input": "hidden_states", + "embedding": "dedicated" if dedicated_embeddings else "shared_target", + "lm_head": "dedicated" if dedicated_lm_head else "shared_target", + } + if mtp_model.metadata_props.get("mobius.mtp_contract") == "nemotron-h-lightning-nextn@1": + conditioning.update( + target_hidden_normalization="post_final_norm", + token_alignment="target_h_i_with_token_t_i_plus_1", + draft_prediction="token_t_i_plus_2", + tables="target_shared_sources_copied_into_draft_graph", + ) payload = { "schema_version": 1, "status": "runtime_unvalidated", "model": {"filename": model_filename}, "inputs": [value.name for value in mtp_model.graph.inputs if value.name is not None], "outputs": [value.name for value in mtp_model.graph.outputs if value.name is not None], - "conditioning": { - "target_hidden_output": "mtp_seed", - "target_hidden_input": "hidden_states", - "embedding": "dedicated" if dedicated_embeddings else "shared_target", - "lm_head": "dedicated" if dedicated_lm_head else "shared_target", - }, + "conditioning": conditioning, "cache_namespaces": { "target": { "namespace": "target", diff --git a/src/mobius/integrations/transformers/_builder.py b/src/mobius/integrations/transformers/_builder.py index afe49c663..78f94baec 100644 --- a/src/mobius/integrations/transformers/_builder.py +++ b/src/mobius/integrations/transformers/_builder.py @@ -403,6 +403,9 @@ def build_transformers_model( ``target_decoder_only`` explicitly exports Nemotron-H without NextN/MTP. Source config facts remain unchanged; the strict streaming report names every omitted ``mtp.*`` tensor. This is not faithful MTP or a default fallback. + ``task='nemotron-h-mtp'`` explicitly exports Lightning target and MTP + graphs with a post-final-norm target seed and independently cached draft + computation. It does not implement speculative generation orchestration. """ if input_sampling_rate is not None and bwe_sampling_rate is not None: raise ValueError("input_sampling_rate and bwe_sampling_rate are mutually exclusive") @@ -564,6 +567,26 @@ def build_transformers_model( ) hf_config, parent_config, model_type = _select_primary_config(hf_config) + from mobius.tasks import NemotronHMtpTask + + export_lightning_mtp = task == "nemotron-h-mtp" or isinstance(task, NemotronHMtpTask) + if export_lightning_mtp: + if model_type != "nemotron_h" or target_decoder_only: + raise ValueError( + "nemotron-h-mtp requires Nemotron-H and excludes target_decoder_only" + ) + if module_class is not None: + raise ValueError("nemotron-h-mtp does not support a module_class override") + if ( + text_only + or fp8_kv_cache + or kv_cache_scales is not None + or prune_prefill_prefix + or export_paged_attention + or glm_full_attention + or output_layer_indices is not None + ): + raise ValueError("nemotron-h-mtp does not support decoder ABI/feature overrides") if target_decoder_only and model_type != "nemotron_h": raise ValueError("target_decoder_only is only supported for Nemotron-H") hf_config, parent_config, modelopt_source = _prepare_nemotron_modelopt_source( @@ -704,7 +727,11 @@ def build_transformers_model( hf_config=parent_config, ) model_module: nn.Module - if target_decoder_only: + if export_lightning_mtp: + from mobius.models.nemotron_h_mtp import NemotronHSpeculativeModel + + model_module = NemotronHSpeculativeModel(config) + elif target_decoder_only: from mobius.models.nemotron_h import NemotronHCausalLMModel if module_class is not NemotronHCausalLMModel: @@ -735,6 +762,11 @@ def build_transformers_model( ) for name, model in package.items(): model.graph.name = f"{graph_source_name}/{name}" + if export_lightning_mtp: + model.metadata_props["mobius.source_revision"] = revision or "local-or-unpinned" + model.metadata_props["mobius.source_num_nextn_predict_layers"] = str( + config.num_nextn_predict_layers + ) if target_decoder_only: model.graph.name = f"{graph_source_name}/target-decoder-only/{name}" model.metadata_props["mobius.export_variant"] = "target-decoder-only" @@ -772,16 +804,17 @@ def build_transformers_model( if load_weights: _reject_unsupported_affine_qwen4(model_type, config) - if modelopt_source is not None or target_decoder_only: + if modelopt_source is not None or target_decoder_only or export_lightning_mtp: from mobius.integrations.transformers._nemotron_h_weights import ( + build_nemotron_h_mtp_streaming_plan, build_nemotron_h_streaming_plan, ) - if len(package) != 1: + if not export_lightning_mtp and len(package) != 1: raise ValueError("Nemotron-H dense streaming requires one decoder component") paths = _resolve_shard_paths(model_id, revision) report = stream_preprocessed_safetensors_to_model( - next(iter(package.values())), + package["decoder"] if export_lightning_mtp else next(iter(package.values())), model_id, lambda index, initializers: build_nemotron_h_streaming_plan( config, @@ -789,10 +822,51 @@ def build_transformers_model( initializers, modelopt=modelopt_source, target_decoder_only=target_decoder_only, + mtp_companion=export_lightning_mtp, ), revision=revision, _resolved_paths=paths, ) + if export_lightning_mtp: + mtp_report = stream_preprocessed_safetensors_to_model( + package["mtp"], + model_id, + lambda index, initializers: build_nemotron_h_mtp_streaming_plan( + config, + index, + initializers, + package["decoder"].graph.initializers, + modelopt=modelopt_source, + ), + revision=revision, + _resolved_paths=paths, + ) + # Partitioning into two graphs is not a target-only omission. + decoder_report = dict(report) + decoder_report["mtp_preserved"] = True + decoder_report["mtp_preservation_component"] = "mtp" + report["decoder_component"] = decoder_report + report["export_variant"] = "target-with-mtp" + report["mtp_preserved"] = True + report["mtp_tensor_count"] = mtp_report["mtp_tensor_count"] + report["mtp_tensors"] = mtp_report["mtp_tensors"] + report["omitted_mtp_tensor_count"] = 0 + report["omitted_mtp_tensors"] = {} + report["decoder_component_mtp_partition_count"] = mtp_report[ + "mtp_tensor_count" + ] + report["mtp_component"] = mtp_report + report["generation_runtime_integration"] = False + report["ignored_tensors"] = 0 + decoder_assigned = report["assigned_tensors"] + mtp_assigned = mtp_report["assigned_tensors"] + if not isinstance(decoder_assigned, int) or not isinstance(mtp_assigned, int): + raise TypeError("Malformed Nemotron-H component tensor accounting") + report["assigned_tensors"] = decoder_assigned + mtp_assigned + report["shared_table_copies"] = 2 + package["decoder"].metadata_props["mobius.weight_loading"] = json.dumps( + decoder_report, sort_keys=True + ) source_roots = {pathlib.Path(path).parent for path in paths} if pathlib.Path(model_id).is_dir(): source_roots.add(pathlib.Path(model_id)) @@ -810,6 +884,9 @@ def build_transformers_model( package.weight_loading_report = report for model in package.values(): model.metadata_props["mobius.weight_loading_status"] = "lazy-bound" + if export_lightning_mtp: + model.metadata_props["mobius.mtp_inventory_status"] = "inspected" + model.metadata_props["mobius.mtp_preserved"] = "true" if target_decoder_only: model.metadata_props["mobius.mtp_inventory_status"] = "inspected" model.metadata_props["mobius.omitted_mtp_tensors"] = json.dumps( diff --git a/src/mobius/integrations/transformers/_nemotron_h_weights.py b/src/mobius/integrations/transformers/_nemotron_h_weights.py index 2fd15ee97..34a864e9a 100644 --- a/src/mobius/integrations/transformers/_nemotron_h_weights.py +++ b/src/mobius/integrations/transformers/_nemotron_h_weights.py @@ -19,6 +19,7 @@ from mobius.integrations.modelopt._config import ModelOptConfig from mobius.integrations.modelopt._weights import validate_modelopt_scale from mobius.models.nemotron_h import _rename_nemotron_h_weight +from mobius.models.nemotron_h_mtp import mtp_source_name def build_nemotron_h_streaming_plan( @@ -28,19 +29,29 @@ def build_nemotron_h_streaming_plan( *, modelopt: ModelOptConfig | None = None, target_decoder_only: bool = False, + mtp_companion: bool = False, ) -> StreamingWeightPlan: """Map exact per-expert checkpoint keys and account for every source tensor.""" mtp_names = sorted(name for name in key_index if name.startswith("mtp.")) - if (config.num_nextn_predict_layers or mtp_names) and not target_decoder_only: + if (config.num_nextn_predict_layers or mtp_names) and not ( + target_decoder_only or mtp_companion + ): raise NotImplementedError( "Nemotron-H NextN/MTP tensors require a validated MTP export contract; " "they must not be silently discarded. Use explicit target_decoder_only=True " "for the separately labeled target decoder variant." ) - if target_decoder_only and config.num_nextn_predict_layers and not mtp_names: + if ( + (target_decoder_only or mtp_companion) + and config.num_nextn_predict_layers + and not mtp_names + ): raise ValueError("Source config declares NextN/MTP but no mtp.* tensors were found") omitted = dict.fromkeys( - mtp_names, "explicit target-decoder-only variant excludes NextN/MTP" + mtp_names, + "bound by companion MTP graph" + if mtp_companion + else "explicit target-decoder-only variant excludes NextN/MTP", ) targets: dict[str, StreamingWeightSource] = {} constants: dict[str, torch.Tensor] = {} @@ -112,11 +123,18 @@ def build_nemotron_h_streaming_plan( ignored=ignored, constants=constants, report={ - "export_variant": "target-decoder-only" if target_decoder_only else "standard", + "export_variant": ( + "target-with-mtp" + if mtp_companion + else "target-decoder-only" + if target_decoder_only + else "standard" + ), "source_num_nextn_predict_layers": config.num_nextn_predict_layers, + "accounted_source_tensor_count": len(key_index), "mtp_preserved": False, "mtp_inventory_status": "inspected", - "omitted_mtp_tensor_count": len(omitted), + "omitted_mtp_tensor_count": 0 if mtp_companion else len(omitted), "omitted_mtp_tensors": { name: { "reason": reason, @@ -124,7 +142,9 @@ def build_nemotron_h_streaming_plan( "dtype": key_index[name][2], } for name, reason in omitted.items() + if not mtp_companion }, + "companion_mtp_tensors": mtp_names if mtp_companion else [], "source_weight_format": modelopt.source_format if modelopt else "floating", "storage_policy": "explicit-dense-bf16-reconstruction" if modelopt else "dense", "native_nvfp4": False, @@ -134,3 +154,85 @@ def build_nemotron_h_streaming_plan( "reconstructed_projection_count": len(quantized_modules), }, ) + + +def build_nemotron_h_mtp_streaming_plan( + config: NemotronHConfig, + key_index: Mapping[str, tuple[str, list[int], str]], + initializers: Mapping[str, ir.Value], + decoder_initializers: Mapping[str, ir.Value], + *, + modelopt: ModelOptConfig | None = None, +) -> StreamingWeightPlan: + """Bind all MTP tensors and shared tables; classify only validated decoder sources. + + The decoder plan is independently applied to the companion graph before + this package is returned. Missing/extra MTP weights fail here; missing, + extra, duplicate or malformed target sources fail in that companion plan. + The pinned ModelOpt source leaves MTP weights floating, not NVFP4. + """ + decoder_plan = build_nemotron_h_streaming_plan( + config, + key_index, + decoder_initializers, + modelopt=modelopt, + mtp_companion=True, + ) + targets: dict[str, StreamingWeightSource] = {} + consumed: set[str] = set() + for name, initializer in initializers.items(): + if initializer.const_value is not None: + continue + source = mtp_source_name(name) + if source.startswith("mtp."): + metadata = key_index.get(source) + if metadata is None: + raise ValueError(f"Missing Lightning MTP tensor {source!r}") + if metadata[1] != list(initializer.shape or ()): + raise ValueError(f"Malformed Lightning MTP tensor shape {source!r}") + if metadata[2] not in {"BF16", "F16", "F32"}: + raise ValueError(f"Lightning MTP requires floating source tensor {source!r}") + binding = StreamingWeightSource(source) + else: + decoder_name = ( + "model.embed_tokens.weight" if name == "embed_tokens.weight" else name + ) + shared_binding = decoder_plan.targets[decoder_name] + if not isinstance(shared_binding, StreamingWeightSource): + raise TypeError(f"Unsupported shared Lightning table binding {source!r}") + binding = shared_binding + targets[name] = binding + consumed.add(binding.source_name) + if binding.scale_name is not None: + consumed.add(binding.scale_name) + if binding.global_scale_name is not None: + consumed.add(binding.global_scale_name) + mtp_inventory = {name for name in key_index if name.startswith("mtp.")} + if mtp_inventory != {name for name in consumed if name.startswith("mtp.")}: + raise ValueError( + f"Unexpected Lightning MTP tensor(s): {sorted(mtp_inventory - consumed)[:5]}" + ) + return StreamingWeightPlan( + targets, + constants=decoder_plan.constants, + ignored={ + name: "bound/validated by companion decoder graph" + for name in key_index + if name not in consumed and name not in decoder_plan.constants + }, + report={ + "export_variant": "target-with-mtp", + "mtp_preserved": True, + "mtp_tensor_count": len(mtp_inventory), + "mtp_tensors": sorted(mtp_inventory), + "shared_target_sources": ["backbone.embeddings.weight", "lm_head.weight"], + "source_num_nextn_predict_layers": config.num_nextn_predict_layers, + "source_weight_format": modelopt.source_format if modelopt else "floating", + "storage_policy": "explicit-dense-bf16-reconstruction" if modelopt else "dense", + "native_nvfp4": False, + "generation_runtime_integration": False, + "companion_decoder_sources": sorted( + name for name in key_index if name not in consumed + ), + }, + ) diff --git a/src/mobius/models/__init__.py b/src/mobius/models/__init__.py index 638ae99db..0ed884c51 100644 --- a/src/mobius/models/__init__.py +++ b/src/mobius/models/__init__.py @@ -138,6 +138,8 @@ "NanoChatCausalLMModel", "NemotronCausalLMModel", "NemotronHCausalLMModel", + "NemotronHMtpModel", + "NemotronHSpeculativeModel", "NeoBertGGUFModel", "NomicBertGGUFModel", "NomicBertMoEGGUFModel", @@ -374,6 +376,7 @@ from mobius.models.nemotron import NemotronCausalLMModel from mobius.models.nemotron3_diarization import Nemotron3DiarizationModel from mobius.models.nemotron_h import NemotronHCausalLMModel +from mobius.models.nemotron_h_mtp import NemotronHMtpModel, NemotronHSpeculativeModel from mobius.models.nemotron_parse import NemotronParseForConditionalGeneration from mobius.models.olmo import OLMo2CausalLMModel, OLMoCausalLMModel from mobius.models.opt import OPTCausalLMModel diff --git a/src/mobius/models/_nemotron_h_mtp_test.py b/src/mobius/models/_nemotron_h_mtp_test.py new file mode 100644 index 000000000..0ab50d8aa --- /dev/null +++ b/src/mobius/models/_nemotron_h_mtp_test.py @@ -0,0 +1,549 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Offline Lightning NextN inventory, independent math, and cache regressions.""" + +from __future__ import annotations + +import dataclasses +import json +from types import SimpleNamespace + +import numpy as np +import onnx_ir as ir +import onnxruntime as ort +import pytest +import safetensors.torch +import torch +import torch.nn.functional as functional + +from mobius import build +from mobius._builder import build_from_module +from mobius._configs import NemotronHConfig +from mobius._model_package import ModelPackage +from mobius.integrations.transformers._nemotron_h_weights_test import ( + _config, + _source_name, + _write_mixed_checkpoint, +) +from mobius.models.nemotron_h_mtp import ( + NemotronHMtpModel, + NemotronHSpeculativeModel, + mtp_source_name, +) + + +def _raw(): + return { + **_config(), + "num_nextn_predict_layers": 1, + "mtp_layers_block_type": ["attention", "moe"], + "n_shared_experts": 1, + "norm_topk_prob": True, + "mlp_hidden_act": "relu2", + } + + +def _graph_package(dtype="f32"): + config = NemotronHConfig.from_transformers(SimpleNamespace(**_raw())) + config.dtype = { + "f32": ir.DataType.FLOAT, + "f16": ir.DataType.FLOAT16, + "bf16": ir.DataType.BFLOAT16, + }[dtype] + return build_from_module(NemotronHSpeculativeModel(config), config, task="nemotron-h-mtp") + + +def _checkpoint(tmp_path, *, dtype=torch.float32): + raw = _raw() + (tmp_path / "config.json").write_text(json.dumps(raw), encoding="utf-8") + package = build(str(tmp_path), task="nemotron-h-mtp", dtype="f32", load_weights=False) + generator = torch.Generator().manual_seed(187) + tensors = {} + for component, model in package.items(): + for name, value in model.graph.initializers.items(): + if value.const_value is not None: + continue + source = _source_name(name) if component == "decoder" else mtp_source_name(name) + if source in tensors: + continue + shape = tuple(value.shape) + tensor = torch.randn(shape, generator=generator) * 0.04 + if len(shape) == 1 and source.endswith("weight"): + tensor += 1 + if source.endswith("A_log"): + tensor.zero_() + tensors[source] = tensor.to( + torch.float32 + if source.endswith(("e_score_correction_bias", "A_log")) + else dtype + ) + safetensors.torch.save_file(tensors, tmp_path / "model.safetensors") + return tensors + + +def _reference(ids, seed, mask, past_k, past_v, weights): + """Independent dense PyTorch equations, not the Mobius module/component tree.""" + + def norm(x, key): + return x * torch.rsqrt(x.float().square().mean(-1, keepdim=True) + 1e-5) * weights[key] + + def linear(x, key): + return functional.linear(x, weights[key]) + + prefix = "mtp.layers.0." + embed = functional.embedding(ids, weights["backbone.embeddings.weight"]) + fused = torch.cat( + [norm(embed, prefix + "enorm.weight"), norm(seed, prefix + "hnorm.weight")], -1 + ) + h = linear(fused, prefix + "eh_proj.weight") + n = norm(h, prefix + "norm.weight") + b, s, _ = n.shape + q = linear(n, prefix + "mixer.q_proj.weight").view(b, s, 4, 8).transpose(1, 2) + k = linear(n, prefix + "mixer.k_proj.weight").view(b, s, 2, 8).transpose(1, 2) + v = linear(n, prefix + "mixer.v_proj.weight").view(b, s, 2, 8).transpose(1, 2) + k, v = torch.cat([past_k, k], 2), torch.cat([past_v, v], 2) + scores = q @ k.repeat_interleave(2, 1).transpose(-2, -1) / 8**0.5 + causal = torch.arange(k.shape[2])[None, :] <= (torch.arange(s)[:, None] + past_k.shape[2]) + allowed = causal[None, None] & mask[:, None, None, :].bool() + probabilities = scores.masked_fill(~allowed, -torch.inf).softmax(-1) + attended = (probabilities @ v.repeat_interleave(2, 1)).transpose(1, 2).reshape(b, s, 32) + h = h + linear(attended, prefix + "mixer.o_proj.weight") + prefix = "mtp.layers.1." + n = norm(h, prefix + "norm.weight") + probabilities = linear(n.float(), prefix + "mixer.gate.weight").sigmoid() + choice = probabilities + weights[prefix + "mixer.gate.e_score_correction_bias"] + indices = choice.topk(2, -1).indices + routing = probabilities.gather(-1, indices) + routing = routing / (routing.sum(-1, keepdim=True) + 1e-20) * 2.5 + routed = torch.zeros_like(n) + for i in range(4): + expert = prefix + f"mixer.experts.{i}." + output = linear( + functional.relu(linear(n, expert + "up_proj.weight")).square(), + expert + "down_proj.weight", + ) + weight = (routing * (indices == i)).sum(-1, keepdim=True) + routed += output * weight + shared = prefix + "mixer.shared_experts." + shared_output = linear( + functional.relu(linear(n, shared + "up_proj.weight")).square(), + shared + "down_proj.weight", + ) + result = norm(h + routed + shared_output, prefix + "final_layernorm.weight") + return linear(result, "lm_head.weight"), result, k, v + + +def _session(model, path): + ir.save(model, path) + options = ort.SessionOptions() + options.intra_op_num_threads = 1 + return ort.InferenceSession(str(path), options, providers=["CPUExecutionProvider"]) + + +def test_explicit_graph_contract_and_shared_sources(): + package = _graph_package() + assert set(package) == {"decoder", "mtp"} + assert "mtp_seed" in {v.name for v in package["decoder"].graph.outputs} + mtp = package["mtp"] + assert [v.name for v in mtp.graph.inputs] == [ + "input_ids", + "hidden_states", + "attention_mask", + "past_key_values.0.key", + "past_key_values.0.value", + ] + assert [v.name for v in mtp.graph.outputs] == [ + "logits", + "mtp_hidden", + "present.0.key", + "present.0.value", + ] + assert mtp.graph.inputs[1].shape[-1] == 32 + assert list(mtp.graph.inputs[3].shape)[1::2] == [2, 8] + for value in mtp.graph.inputs: + assert value.dtype == ( + ir.DataType.INT64 + if value.name in {"input_ids", "attention_mask"} + else ir.DataType.FLOAT + ) + expected_shapes = { + "logits": ["batch", "sequence_len", "48"], + "mtp_hidden": ["batch", "sequence_len", "32"], + "present.0.key": ["batch", "2", "past_sequence_len + sequence_len", "8"], + "present.0.value": ["batch", "2", "past_sequence_len + sequence_len", "8"], + } + for value in mtp.graph.outputs: + assert value.dtype == ir.DataType.FLOAT + assert [str(dimension) for dimension in value.shape] == expected_shapes[value.name] + assert "position_ids" not in {v.name for v in mtp.graph.inputs} + assert not any(n.op_type in {"RotaryEmbedding", "Sin", "Cos", "Scan"} for n in mtp.graph) + + +@pytest.mark.parametrize("dtype", ["f32", "f16", "bf16"]) +def test_mtp_graph_dtypes(dtype): + package = _graph_package(dtype) + expected = { + "f32": ir.DataType.FLOAT, + "f16": ir.DataType.FLOAT16, + "bf16": ir.DataType.BFLOAT16, + }[dtype] + for component in package.values(): + assert next(v for v in component.graph.outputs if v.name == "logits").dtype == expected + router = package["mtp"].graph.initializers["layers.1.moe.gate.weight"] + assert router.dtype == ir.DataType.FLOAT + + +@pytest.mark.parametrize("source_dtype", [torch.float32, torch.bfloat16]) +def test_public_package_independent_cpu_prefill_cached_decode_and_seed(tmp_path, source_dtype): + weights = {n: t.float() for n, t in _checkpoint(tmp_path, dtype=source_dtype).items()} + package = build(str(tmp_path), task="nemotron-h-mtp", dtype="f32") + session = _session(package["mtp"], tmp_path / "mtp.onnx") + generator = torch.Generator().manual_seed(32) + past_k = past_v = torch.zeros(2, 2, 0, 8) + for length in (3, 1, 2): + ids = torch.randint(0, 48, (2, length), generator=generator) + seed = torch.randn(2, length, 32, generator=generator) + mask = torch.ones(2, past_k.shape[2] + length, dtype=torch.int64) + mask[1, 1] = 0 + expected = _reference(ids, seed, mask, past_k, past_v, weights) + actual = session.run( + None, + { + "input_ids": ids.numpy(), + "hidden_states": seed.numpy(), + "attention_mask": mask.numpy(), + "past_key_values.0.key": past_k.numpy(), + "past_key_values.0.value": past_v.numpy(), + }, + ) + for result, reference in zip(actual, expected): + np.testing.assert_allclose(result, reference.numpy(), rtol=1e-4, atol=1e-5) + past_k, past_v = expected[2:] + decoder = _session(package["decoder"], tmp_path / "decoder.onnx") + inputs = { + "input_ids": np.array([[4, 5, 6], [7, 8, 9]], dtype=np.int64), + "attention_mask": np.ones((2, 3), dtype=np.int64), + "past_key_values.0.conv_state": np.zeros((2, 64, 3), dtype=np.float32), + "past_key_values.0.ssm_state": np.zeros((2, 4, 8, 8), dtype=np.float32), + "past_key_values.2.key": np.zeros((2, 2, 0, 8), dtype=np.float32), + "past_key_values.2.value": np.zeros((2, 2, 0, 8), dtype=np.float32), + } + logits, seed = decoder.run(["logits", "mtp_seed"], inputs) + np.testing.assert_allclose( + logits, seed @ weights["lm_head.weight"].numpy().T, rtol=1e-5, atol=1e-6 + ) + # Exercise the real decoder→MTP bridge with shifted target tokens. + aligned_ids = inputs["input_ids"][:, 1:] + aligned_seed = seed[:, :-1] + actual = session.run( + None, + { + "input_ids": aligned_ids, + "hidden_states": aligned_seed, + "attention_mask": np.ones((2, 2), dtype=np.int64), + "past_key_values.0.key": np.zeros((2, 2, 0, 8), dtype=np.float32), + "past_key_values.0.value": np.zeros((2, 2, 0, 8), dtype=np.float32), + }, + ) + expected = _reference( + torch.from_numpy(aligned_ids), + torch.from_numpy(aligned_seed), + torch.ones(2, 2, dtype=torch.int64), + torch.zeros(2, 2, 0, 8), + torch.zeros(2, 2, 0, 8), + weights, + ) + np.testing.assert_allclose(actual[0], expected[0].numpy(), rtol=1e-4, atol=1e-5) + + +@pytest.mark.parametrize("mutation", ["missing", "extra", "shape", "duplicate"]) +def test_strict_mtp_inventory(tmp_path, mutation): + tensors = _checkpoint(tmp_path) + if mutation == "missing": + del tensors["mtp.layers.1.mixer.shared_experts.up_proj.weight"] + elif mutation == "extra": + tensors["mtp.unexpected.weight"] = torch.zeros(1) + elif mutation == "shape": + tensors["mtp.layers.0.eh_proj.weight"] = torch.zeros(32, 32) + else: + tensors["model.embeddings.weight"] = tensors["backbone.embeddings.weight"].clone() + safetensors.torch.save_file(tensors, tmp_path / "model.safetensors") + with pytest.raises(ValueError): + build(str(tmp_path), task="nemotron-h-mtp", dtype="f32") + + +def test_package_accounting_save_and_default_guard(tmp_path): + tensors = _checkpoint(tmp_path, dtype=torch.bfloat16) + with pytest.raises(NotImplementedError, match="NextN/MTP"): + build(str(tmp_path), dtype="bf16", load_weights=False) + package = build(str(tmp_path), task="nemotron-h-mtp", dtype="bf16") + report = package.weight_loading_report + assert report["mtp_preserved"] is True + assert report["omitted_mtp_tensor_count"] == 0 + assert set(report["mtp_tensors"]) == {n for n in tensors if n.startswith("mtp.")} + assert report["mtp_tensor_count"] == 22 + assert report["ignored_tensors"] == 0 + assert report["accounted_source_tensor_count"] == len(tensors) + assert report["assigned_tensors"] == len(tensors) + 2 # Shared table copies. + assert report["decoder_component"]["mtp_preservation_component"] == "mtp" + assert package.config.num_nextn_predict_layers == 1 + before = (tmp_path / "model.safetensors").read_bytes() + with pytest.raises(ValueError): + package.save(tmp_path, external_data="safetensors") + assert (tmp_path / "model.safetensors").read_bytes() == before + output = tmp_path / "export" + package.save(output, external_data="onnx") + assert (output / "decoder" / "model.onnx").is_file() + assert (output / "mtp" / "model.onnx").is_file() + assert json.loads((output / "weight-loading-report.json").read_text())["mtp_preserved"] + restored = ModelPackage.load(output) + assert set(restored) == {"decoder", "mtp"} + for model in restored.values(): + assert model.metadata_props["mobius.source_num_nextn_predict_layers"] == "1" + assert model.metadata_props["mobius.mtp_contract"] == "nemotron-h-lightning-nextn@1" + assert restored.weight_loading_report["mtp_preserved"] is True + + +@pytest.mark.parametrize( + "field,value", + [ + ("num_nextn_predict_layers", 2), + ("mtp_layers_block_type", ["attention"]), + ("n_shared_experts", 0), + ("moe_latent_size", 16), + ("residual_in_fp32", True), + ], +) +def test_architecture_variants_rejected(field, value): + config = NemotronHConfig.from_transformers(SimpleNamespace(**_raw())) + with pytest.raises(NotImplementedError): + NemotronHSpeculativeModel(dataclasses.replace(config, **{field: value})) + + +def test_mixed_reconstruction_with_floating_mtp_and_shared_quantized_head(tmp_path): + floating = tmp_path / "floating" + floating.mkdir() + mtp_tensors = { + n: t + for n, t in _checkpoint(floating, dtype=torch.bfloat16).items() + if n.startswith("mtp.") + } + raw, expected = _write_mixed_checkpoint(tmp_path) + raw.update( + num_nextn_predict_layers=1, + mtp_layers_block_type=["attention", "moe"], + n_shared_experts=1, + ) + (tmp_path / "config.json").write_text(json.dumps(raw), encoding="utf-8") + tensors = safetensors.torch.load_file(tmp_path / "model.safetensors") + tensors.update(mtp_tensors) + safetensors.torch.save_file(tensors, tmp_path / "model.safetensors") + with pytest.raises(NotImplementedError, match="ModelOpt"): + build(str(tmp_path), task="nemotron-h-mtp", dtype="bf16") + package = build(str(tmp_path), task="nemotron-h-mtp", keep_quantized=False, dtype="bf16") + for name, value in package["mtp"].graph.initializers.items(): + if name == "lm_head.weight": + tensor = expected[name] + elif name == "embed_tokens.weight": + tensor = expected["model.embed_tokens.weight"] + elif name.startswith("layers."): + tensor = mtp_tensors[mtp_source_name(name)] + else: + continue + np.testing.assert_array_equal( + value.const_value.numpy().astype("float32"), tensor.float().numpy() + ) + assert package.weight_loading_report["mtp_tensor_count"] == 22 + assert package.weight_loading_report["source_weight_format"] == "modelopt-mixed-fp8-nvfp4" + assert package.weight_loading_report["generation_runtime_integration"] is False + package.save(tmp_path / "export", external_data="onnx") + + +def test_cli_explicit_mtp_route(tmp_path): + from mobius.__main__ import main + + _checkpoint(tmp_path, dtype=torch.bfloat16) + output = tmp_path / "cli-export" + main( + [ + "build", + "--config", + str(tmp_path), + "--output", + str(output), + "--task", + "nemotron-h-mtp", + "--dtype", + "bf16", + "--ep", + "cpu", + "--external-data", + "onnx", + ] + ) + assert (output / "decoder" / "model.onnx").is_file() + assert (output / "mtp" / "model.onnx").is_file() + assert json.loads((output / "weight-loading-report.json").read_text())["mtp_preserved"] + + +def test_intrinsic_runtime_metadata_has_precise_external_contract(tmp_path): + from mobius.integrations.ort_genai.auto_export import write_ort_genai_config + + package = _graph_package() + write_ort_genai_config(package, str(tmp_path)) + payload = json.loads((tmp_path / "mtp_config.json").read_text()) + assert payload["status"] == "runtime_unvalidated" + assert payload["conditioning"]["target_hidden_output"] == "mtp_seed" + assert payload["conditioning"]["target_hidden_normalization"] == "post_final_norm" + assert payload["conditioning"]["token_alignment"] == "target_h_i_with_token_t_i_plus_1" + assert payload["conditioning"]["draft_prediction"] == "token_t_i_plus_2" + assert payload["conditioning"]["tables"] == "target_shared_sources_copied_into_draft_graph" + assert payload["cache_namespaces"]["target"]["ports"] + assert payload["cache_namespaces"]["mtp"]["ports"] + assert payload["num_nextn_predict_layers"] == 1 + + +def test_pinned_production_mtp_parameter_shapes(): + # Derived from independent range-read headers at the recipe's immutable + # BF16/NVFP4 revisions, not inferred from this implementation's graph. + raw = { + **_raw(), + "hidden_size": 2688, + "vocab_size": 131072, + "num_attention_heads": 32, + "num_key_value_heads": 2, + "head_dim": 128, + "n_routed_experts": 128, + "num_experts_per_tok": 6, + "moe_intermediate_size": 1856, + "moe_shared_expert_intermediate_size": 3712, + } + config = NemotronHConfig.from_transformers(SimpleNamespace(**raw)) + model = NemotronHMtpModel(config) + expected = { + "mtp.layers.0.enorm.weight": [2688], + "mtp.layers.0.hnorm.weight": [2688], + "mtp.layers.0.eh_proj.weight": [2688, 5376], + "mtp.layers.0.norm.weight": [2688], + "mtp.layers.0.mixer.q_proj.weight": [4096, 2688], + "mtp.layers.0.mixer.k_proj.weight": [256, 2688], + "mtp.layers.0.mixer.v_proj.weight": [256, 2688], + "mtp.layers.0.mixer.o_proj.weight": [2688, 4096], + "mtp.layers.1.norm.weight": [2688], + "mtp.layers.1.final_layernorm.weight": [2688], + "mtp.layers.1.mixer.gate.weight": [128, 2688], + "mtp.layers.1.mixer.gate.e_score_correction_bias": [128], + "mtp.layers.1.mixer.shared_experts.up_proj.weight": [3712, 2688], + "mtp.layers.1.mixer.shared_experts.down_proj.weight": [2688, 3712], + } + for index in range(128): + expected[f"mtp.layers.1.mixer.experts.{index}.up_proj.weight"] = [1856, 2688] + expected[f"mtp.layers.1.mixer.experts.{index}.down_proj.weight"] = [2688, 1856] + actual = { + mtp_source_name(name): list(value.shape) + for name, value in model.named_parameters() + if name.startswith("layers.") + } + assert actual == expected + assert len(actual) == 270 + + +def test_transformers_mtp_vocabulary_preserves_source_facts(): + from transformers import AutoConfig + + raw = _raw() + source = AutoConfig.for_model(**raw) + before = source.to_dict() + config = NemotronHConfig.from_transformers(source) + assert config.mtp_layers_block_type == ["attention", "moe"] + assert config.num_nextn_predict_layers == 1 + assert source.to_dict() == before + assert raw["mtp_layers_block_type"] == ["attention", "moe"] + + +def test_pinned_revision_reaches_config_shards_and_component_metadata(tmp_path, monkeypatch): + from mobius.integrations.transformers import _builder + + _checkpoint(tmp_path) + source = SimpleNamespace(**_raw()) + before = dict(source.__dict__) + calls = [] + revision = "a9904d24bcc1d289a1950fa9d2b978c47cf903b9" + + def load_config(model_id, **kwargs): + calls.append(("config", model_id, kwargs["revision"])) + return source, True + + def resolve_shards(model_id, selected_revision): + calls.append(("shards", model_id, selected_revision)) + return [str(tmp_path / "model.safetensors")] + + monkeypatch.setattr(_builder, "_load_transformers_config", load_config) + monkeypatch.setattr(_builder, "_resolve_shard_paths", resolve_shards) + package = build("nvidia/lightning", task="nemotron-h-mtp", revision=revision, dtype="f32") + assert calls == [ + ("config", "nvidia/lightning", revision), + ("shards", "nvidia/lightning", revision), + ] + assert source.__dict__ == before + assert package.weight_loading_report["revision"] == revision + assert package.weight_loading_report["mtp_component"]["revision"] == revision + for model in package.values(): + assert model.metadata_props["mobius.source_revision"] == revision + + +def test_fp16_cpu_draft_prefill_and_cached_decode(tmp_path): + weights = _checkpoint(tmp_path) + package = build(str(tmp_path), task="nemotron-h-mtp", dtype="f16") + session = _session(package["mtp"], tmp_path / "mtp-f16.onnx") + reference_weights = { + name: tensor.float() if ".gate." in name else tensor.half().float() + for name, tensor in weights.items() + } + generator = torch.Generator().manual_seed(819) + past_k = past_v = np.zeros((2, 2, 0, 8), dtype=np.float16) + for length in (3, 1): + ids = torch.randint(0, 48, (2, length), generator=generator) + seed = torch.randn(2, length, 32, generator=generator).half().numpy() + mask = torch.ones(2, past_k.shape[2] + length, dtype=torch.int64) + expected = _reference( + ids, + torch.from_numpy(seed).float(), + mask, + torch.from_numpy(past_k).float(), + torch.from_numpy(past_v).float(), + reference_weights, + ) + actual = session.run( + None, + { + "input_ids": ids.numpy(), + "hidden_states": seed, + "attention_mask": mask.numpy(), + "past_key_values.0.key": past_k, + "past_key_values.0.value": past_v, + }, + ) + for result, reference in zip(actual, expected): + assert result.dtype == np.float16 + np.testing.assert_allclose( + result.astype("float32"), reference.numpy(), rtol=1e-2, atol=1e-2 + ) + past_k, past_v = actual[2:] + + +@pytest.mark.parametrize( + "options", + [ + {"target_decoder_only": True}, + {"fp8_kv_cache": True}, + {"prune_prefill_prefix": True}, + {"export_paged_attention": True}, + {"output_layer_indices": [0]}, + ], +) +def test_explicit_mtp_rejects_unsupported_feature_combinations(tmp_path, options): + (tmp_path / "config.json").write_text(json.dumps(_raw()), encoding="utf-8") + with pytest.raises(ValueError, match="nemotron-h-mtp"): + build(str(tmp_path), task="nemotron-h-mtp", load_weights=False, **options) diff --git a/src/mobius/models/nemotron_h.py b/src/mobius/models/nemotron_h.py index ea38f7009..c6207e2b4 100644 --- a/src/mobius/models/nemotron_h.py +++ b/src/mobius/models/nemotron_h.py @@ -134,7 +134,7 @@ def forward( op: OpBuilder, hidden_states: ir.Value, attention_bias: ir.Value, - position_embeddings: tuple, + position_embeddings: tuple | None, past_key_value: tuple | None, ): """Forward pass. Returns (hidden_states, (key, value)).""" @@ -413,8 +413,8 @@ def forward( self, op: OpBuilder, hidden_states: ir.Value, - attention_bias: ir.Value, - position_embeddings: tuple, + attention_bias: ir.Value | None, + position_embeddings: tuple | None, past_key_value: tuple | None, ): """Forward pass. Returns (hidden_states, (None, None)). @@ -530,8 +530,9 @@ def __init__(self, config: NemotronHConfig, *, target_decoder_only: bool = False self.target_decoder_only = target_decoder_only if config.num_nextn_predict_layers and not target_decoder_only: raise NotImplementedError( - "Nemotron-H NextN/MTP export is not implemented; auxiliary tensors " - "must not be silently skipped. Explicit target_decoder_only=True " + "Nemotron-H NextN/MTP requires explicit task='nemotron-h-mtp'; " + "default export must not silently skip auxiliary tensors. " + "Explicit target_decoder_only=True " "(--target-decoder-only) exports a separately labeled variant." ) self.model = _NemotronHTextModel(config) diff --git a/src/mobius/models/nemotron_h_mtp.py b/src/mobius/models/nemotron_h_mtp.py new file mode 100644 index 000000000..12e9bb131 --- /dev/null +++ b/src/mobius/models/nemotron_h_mtp.py @@ -0,0 +1,171 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Lightning's single NextN step: fusion -> attention -> MoE -> RMSNorm. + +Replicates vLLM's NemotronHMultiTokenPredictor at +ef63c23d35acccbba8e014fc88da19d541da38a9. The pinned Lightning checkpoints +store two physical MTP blocks, including dedicated shared-expert weights. +Embedding and vocabulary projection weights are shared with the target. +No positional encoding or Mamba state is used by the MTP blocks. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from onnxscript import OpBuilder, nn + +from mobius._configs import NemotronHConfig +from mobius.components import Embedding, Linear, RMSNorm, create_padding_mask +from mobius.models.nemotron_h import ( + NemotronHAttentionLayer, + NemotronHCausalLMModel, + NemotronHMoELayer, +) + +if TYPE_CHECKING: + import onnx_ir as ir + + +def validate_lightning_mtp_config(config: NemotronHConfig) -> None: + """Reject unsupported NextN variants rather than guessing their architecture.""" + if config.num_nextn_predict_layers != 1: + raise NotImplementedError("Lightning MTP requires exactly one NextN prediction step") + if config.mtp_layers_block_type != ["attention", "moe"]: + raise NotImplementedError( + "Lightning MTP requires mtp_layers_block_type=['attention','moe']" + ) + if config.n_shared_experts != 1 or config.moe_latent_size is not None: + raise NotImplementedError("Lightning MTP requires one shared expert and no latent MoE") + if ( + config.residual_in_fp32 + or config.attn_qkv_bias + or config.attn_o_bias + or config.mlp_bias + ): + raise NotImplementedError("Lightning MTP supports bias-free, model-dtype residuals") + if config.hidden_act != "relu2" or not config.shared_expert_intermediate_size: + raise NotImplementedError("Lightning MTP requires ReLU2 and an explicit shared width") + + +class _FusionAttentionLayer(NemotronHAttentionLayer): + def __init__(self, config: NemotronHConfig): + super().__init__(config) + self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.eh_proj = Linear(2 * config.hidden_size, config.hidden_size, bias=False) + + def forward( + self, + op: OpBuilder, + inputs_embeds: ir.Value, + hidden_states: ir.Value, + attention_bias: ir.Value, + past_key_value: tuple, + ): + # Condition on t_(i+1)'s embedding and the target's post-final-norm h_i. + fused = op.Concat( + self.enorm(op, inputs_embeds), self.hnorm(op, hidden_states), axis=-1 + ) # (batch, sequence, 2H) + hidden_states = self.eh_proj(op, fused) # (batch, sequence, H) + return super().forward(op, hidden_states, attention_bias, None, past_key_value) + + +class _FinalMoELayer(NemotronHMoELayer): + def __init__(self, config: NemotronHConfig): + super().__init__(config) + self.final_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward(self, op: OpBuilder, hidden_states: ir.Value): + hidden_states, _ = super().forward(op, hidden_states, None, None, None) + return self.final_layernorm(op, hidden_states) + + +class NemotronHMtpModel(nn.Module): + """Standalone draft computation with target-shared embedding and LM head. + + Inputs: token IDs for t_(i+1), post-final-norm target h_i, a padding mask, + and independent MTP attention KV. Outputs: draft logits for t_(i+2), + normalized MTP hidden states for recursive drafting, and updated MTP KV. + Token shifting, speculative acceptance and state rollback belong to the caller. + """ + + config_class: type = NemotronHConfig + default_task: str = "nemotron-h-mtp" + category: str = "Mixture of Experts" + + def __init__(self, config: NemotronHConfig): + super().__init__() + validate_lightning_mtp_config(config) + self.config = config + self.embed_tokens = Embedding( + config.vocab_size, config.hidden_size, config.pad_token_id + ) + self.layers = nn.ModuleList([_FusionAttentionLayer(config), _FinalMoELayer(config)]) + self.lm_head = Linear(config.hidden_size, config.vocab_size, bias=False) + + def forward( + self, + op: OpBuilder, + input_ids: ir.Value, + hidden_states: ir.Value, + attention_mask: ir.Value, + past_key_value: tuple, + ): + inputs_embeds = self.embed_tokens(op, input_ids) + attention_bias = create_padding_mask( + op, input_ids=input_ids, attention_mask=attention_mask + ) + hidden_states, present = self.layers[0]( + op, inputs_embeds, hidden_states, attention_bias, past_key_value + ) + hidden_states = self.layers[1](op, hidden_states) + return self.lm_head(op, hidden_states), hidden_states, present + + +class _MtpSeedTarget(NemotronHCausalLMModel): + def forward( + self, + op: OpBuilder, + input_ids: ir.Value, + attention_mask: ir.Value, + position_ids: ir.Value, + past_key_values: list | None = None, + ): + hidden_states, present = self.model( + op, input_ids, attention_mask, position_ids, past_key_values + ) + return self.lm_head(op, hidden_states), present, None, hidden_states + + +class NemotronHSpeculativeModel(nn.Module): + """Explicit target + Lightning MTP graph package, not a generation runtime.""" + + config_class: type = NemotronHConfig + default_task: str = "nemotron-h-mtp" + category: str = "Mixture of Experts" + + def __init__(self, config: NemotronHConfig): + super().__init__() + validate_lightning_mtp_config(config) + self.config = config + # The explicit package preserves auxiliary tensors in a separate graph. + self.decoder = _MtpSeedTarget(config, target_decoder_only=True) + self.mtp = NemotronHMtpModel(config) + # Each component is its own graph root, not a nested parameter namespace. + self.decoder._set_name("") + self.mtp._set_name("") + + +def mtp_source_name(target_name: str) -> str: + """Exact checkpoint mapping; tables share target sources, not random fallbacks.""" + if target_name == "embed_tokens.weight": + return "backbone.embeddings.weight" + if target_name == "lm_head.weight": + return target_name + if target_name.startswith("layers."): + return "mtp." + target_name.replace(".self_attn.", ".mixer.").replace( + ".moe.", ".mixer." + ) + raise ValueError(f"Unexpected Lightning MTP graph parameter {target_name!r}") diff --git a/src/mobius/tasks/__init__.py b/src/mobius/tasks/__init__.py index 72ca3bfd6..3b4929715 100644 --- a/src/mobius/tasks/__init__.py +++ b/src/mobius/tasks/__init__.py @@ -35,6 +35,7 @@ "DraftTargetCausalLMTask", "Eagle3DraftTask", "Qwen35MtpTask", + "NemotronHMtpTask", "Qwen4ExpCausalLMTask", "Qwen4ExpVisionLanguageTask", "DenoisingTask", @@ -188,6 +189,7 @@ from mobius.tasks._mistral4_gguf import Mistral4GGUFCausalLMTask from mobius.tasks._moshi import MoshiDepformerTask, MoshiTemporalTask from mobius.tasks._multimodal import MultiModalTask +from mobius.tasks._nemotron_h_mtp import NemotronHMtpTask from mobius.tasks._object_detection import ObjectDetectionTask from mobius.tasks._phi4mm_multimodal import Phi4MMMultiModalTask from mobius.tasks._plamo import PlamoCausalLMTask @@ -278,6 +280,7 @@ "dflash-draft": DFlashDraftTask, "eagle3-draft": Eagle3DraftTask, "qwen35-mtp": Qwen35MtpTask, + "nemotron-h-mtp": NemotronHMtpTask, "qwen4-exp-text-generation": Qwen4ExpCausalLMTask, "qwen4-exp-vision-language": Qwen4ExpVisionLanguageTask, "vae": VAETask, diff --git a/src/mobius/tasks/_nemotron_h_mtp.py b/src/mobius/tasks/_nemotron_h_mtp.py new file mode 100644 index 000000000..bf15b8ce3 --- /dev/null +++ b/src/mobius/tasks/_nemotron_h_mtp.py @@ -0,0 +1,81 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Explicit target and draft graph contracts for Lightning NextN/MTP.""" + +from __future__ import annotations + +import onnx_ir as ir +from onnxscript import nn + +from mobius._configs import NemotronHConfig +from mobius._model_package import ModelPackage +from mobius.models.nemotron_h_mtp import NemotronHSpeculativeModel +from mobius.tasks._base import ComponentSpec, ModelTask, _make_graph, _make_model +from mobius.tasks._cache_utils import _make_kv_cache_inputs, _register_kv_cache_outputs +from mobius.tasks._causal_lm import HybridCausalLMTask + + +class NemotronHMtpTask(ModelTask): + """Build target ``decoder`` with ``mtp_seed`` and independent ``mtp`` graph. + + The caller pairs target h_i with t_(i+1), initializes/advances draft KV + independently, verifies draft tokens against the target, and restores + rejected target SSM/conv/KV and draft KV state. No runtime ABI is implied. + """ + + components = ComponentSpec(decoder="decoder", mtp="mtp") + + def build(self, module: nn.Module, config: NemotronHConfig) -> ModelPackage: + self._validate_components(module) + if not isinstance(module, NemotronHSpeculativeModel): + raise TypeError("NemotronHMtpTask requires NemotronHSpeculativeModel") + target = HybridCausalLMTask().build(module.decoder, config)["model"] + graph, builder = _make_graph("lightning_mtp") + batch = ir.SymbolicDim("batch") + input_ids = builder.input( + "input_ids", dtype=ir.DataType.INT64, shape=["batch", "sequence_len"] + ) + hidden = builder.input( + "hidden_states", + dtype=config.dtype, + shape=["batch", "sequence_len", config.hidden_size], + ) + mask = builder.input( + "attention_mask", + dtype=ir.DataType.INT64, + shape=["batch", "past_sequence_len + sequence_len"], + ) + past = _make_kv_cache_inputs( + builder, + 1, + config.num_key_value_heads, + config.head_dim, + config.dtype, + batch, + ir.SymbolicDim("past_sequence_len"), + ) + logits, mtp_hidden, present = module.mtp(builder.op, input_ids, hidden, mask, past[0]) + builder.add_output(logits, "logits") + builder.add_output(mtp_hidden, "mtp_hidden") + _register_kv_cache_outputs( + builder, + [present], + batch=batch, + num_kv_heads=config.num_key_value_heads, + key_head_dim=config.head_dim, + value_head_dim=config.head_dim, + total_seq_len="past_sequence_len + sequence_len", + dtype=config.dtype, + ) + mtp = _make_model(graph) + for model in (target, mtp): + model.metadata_props["mobius.export_variant"] = "target-with-mtp" + model.metadata_props["mobius.mtp_contract"] = "nemotron-h-lightning-nextn@1" + model.metadata_props["mobius.mtp_inventory_status"] = "not_inspected" + model.metadata_props["mobius.runtime_support"] = ( + "Standalone graphs only: requires caller-managed token/hidden alignment, " + "draft KV, target verification and hybrid-state rollback; not an " + "ORT GenAI speculative-generation ABI." + ) + return ModelPackage({"decoder": target, "mtp": mtp}, config=config) From c5d47af4e6c8c26cdca15b4897b14d7b03130f4f Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 8 Oct 2026 13:57:19 -0700 Subject: [PATCH 09/15] Use source-faithful grouped Mamba normalization for Nemotron Keep gating and per-group normalization in FP32, then cast the normalized activation before applying native-dtype gamma, matching the pinned HF Zamba2RMSNormGated implementation. Wire only Nemotron to the explicit component so it avoids the demonstrated CUDA RMSNormalization two-dimensional-scale kernel failure without altering other Mamba defaults or the existing opt-in workaround. Add CPU FP32/FP16 numerical coverage, BF16 graph precision and CUDA-gated probes, production group dimensions, optimized-graph checks, and unchanged weight/cache contract regressions. Full H100 decoder parity and generation remain blocked pending first-divergence diagnostics. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> --- docs/nemotron-lightning-export.md | 40 ++++++- src/mobius/components/__init__.py | 2 + src/mobius/components/_rms_norm.py | 31 +++++ src/mobius/components/_rms_norm_test.py | 149 +++++++++++++++++++++++- src/mobius/models/_nemotron_h_test.py | 27 +++++ src/mobius/models/nemotron_h.py | 6 + 6 files changed, 253 insertions(+), 2 deletions(-) diff --git a/docs/nemotron-lightning-export.md b/docs/nemotron-lightning-export.md index eff3bb4d4..b53bf31a7 100644 --- a/docs/nemotron-lightning-export.md +++ b/docs/nemotron-lightning-export.md @@ -1,6 +1,6 @@ # Nemotron 3.5 Lightning: reconstruction and compatibility gates -**CPU-tested exporter improvements; both full Lightning exports `NOT_RUN`.** +**CPU-tested exporter improvements; full-model generation remains blocked.** The [machine-readable recipe](../examples/nemotron_lightning/recipe.json) records immutable inputs and compatibility limits. Native NVFP4 remains unsupported. Explicit dense BF16 reconstruction is implemented and tested @@ -10,6 +10,44 @@ with complete named reporting. A separate explicit `nemotron-h-mtp` task now exports target and draft graphs with strict MTP coverage and tiny independent CPU parity. This is **graph support**, not a speculative-generation runtime. +### Grouped Mamba normalization correction and H100 gate + +User-supplied H100 evidence for export commit `897a0e05` reports exact source +weight fidelity passing, but full decoder numerical parity and eight-token +generation failing. One isolated failure is ORT +`1.28.0.dev20260722004` CUDA `RMSNormalization` with a two-dimensional grouped +scale. Scratch decomposition improves the decoder substantially but does not +clear its numerical or generation gates. Conditional MTP parity under those +decoder seeds is not end-to-end parity. + +Nemotron's Mamba mixer now uses `GroupedGatedRMSNorm`, with explicit FP32 +gating, per-group variance and normalization, followed by a cast to the input +dtype **before** native-dtype gamma multiplication. This matches the pinned +HF [`Zamba2RMSNormGated`](https://github.com/huggingface/transformers/blob/a96730c8c97b8efbf35bbaf7f5da33ec99231a49/src/transformers/models/zamba2/modeling_zamba2.py) +used by `NemotronHMamba2Mixer`. It avoids the faulty two-dimensional-scale +kernel without changing shared Mamba defaults, weight names, cache contracts +or task selection. The existing `MOBIUS_ORT_CUDA_GROUPED_RMSNORM_WORKAROUND` +flag is unchanged for other users; its FP32 gamma-before-cast ordering is +not the same as this source-faithful variant. + +Focused regressions execute FP32/FP16 grouped math on CPU, including +production normalization dimensions (4096 channels, groups of 512), and +assert FP32/FP16/BF16 graph cast ordering. The component test also includes a +BF16 CUDA case, skipped when CUDA is unavailable. Run on the H100: + +```bash +python -m pytest src/mobius/components/_rms_norm_test.py \ + -k grouped_gated_norm_source_precision_parity -q +``` + +The H100 gate is to re-export to a **fresh** directory with this correction, +verify the optimized graph does not recreate a two-dimensional-scale norm, +rerun the isolated CUDA norm check and full decoder prefill/cached-state +comparisons, then require complete generation-token agreement. Original +exports and checkpoint files must remain unchanged. The remaining decoder +discrepancy needs a first-divergent-stage comparison and an inspected source +reference; this correction alone is **not** a claim of resolved full parity. + The original preflight host was Windows ARM64 with a Qualcomm Adreno X1-85, no NVIDIA/CUDA GPU or `nvidia-smi`, Python 3.12.10, 68,171,038,720 bytes RAM and approximately 818 GB free disk. Mobius, PyTorch, Transformers, ORT and diff --git a/src/mobius/components/__init__.py b/src/mobius/components/__init__.py index d9cee1d06..464f8e4c9 100644 --- a/src/mobius/components/__init__.py +++ b/src/mobius/components/__init__.py @@ -45,6 +45,7 @@ "GatedDeltaNet", "GatedMLP", "GatedRMSNorm", + "GroupedGatedRMSNorm", "KimiDeltaAttention", "KimiMLAAttention", "Gemma3nAudioEncoder", @@ -448,6 +449,7 @@ from mobius.components._radio_vision import RadioVisionModel from mobius.components._rms_norm import ( GatedRMSNorm, + GroupedGatedRMSNorm, OffsetRMSNorm, PostGatedRMSNorm, RMSNorm, diff --git a/src/mobius/components/_rms_norm.py b/src/mobius/components/_rms_norm.py index 1888880d9..928f3ae13 100644 --- a/src/mobius/components/_rms_norm.py +++ b/src/mobius/components/_rms_norm.py @@ -154,6 +154,37 @@ def forward(self, op: OpBuilder, hidden_states: ir.Value, gate: ir.Value): return op.CastLike(normed, hidden_states) +class GroupedGatedRMSNorm(GatedRMSNorm): + """Grouped pre-gated RMSNorm with FP32 normalization and native-dtype scale. + + Matches HuggingFace's ``Zamba2RMSNormGated``: gating and variance stay + in FP32, then the normalized activation is cast before multiplying + the learned scale. Explicit reduction avoids CUDA RMSNormalization + kernels that misindex a two-dimensional grouped scale. + """ + + def __init__(self, hidden_size: int, group_size: int, eps: float = 1e-6): + if hidden_size <= 0 or group_size <= 0 or hidden_size % group_size: + raise ValueError("group_size must be positive and divide positive hidden_size") + super().__init__(hidden_size, eps=eps, group_size=group_size) + self.num_groups = hidden_size // group_size + + def forward(self, op: OpBuilder, hidden_states: ir.Value, gate: ir.Value) -> ir.Value: + hidden_f32 = op.Cast(hidden_states, to=ir.DataType.FLOAT) + gate_f32 = op.Cast(gate, to=ir.DataType.FLOAT) + gated = op.Mul(hidden_f32, op.Swish(gate_f32)) + original_shape = op.Shape(hidden_states) + grouped = op.Reshape( + gated, [-1, self.num_groups, self.group_size] + ) # (flattened leading dimensions, groups, group_size) + variance = op.ReduceMean(op.Mul(grouped, grouped), axes=[-1], keepdims=True) + reciprocal_rms = op.Reciprocal(op.Sqrt(op.Add(variance, self.variance_epsilon))) + normalized = op.Reshape(op.Mul(grouped, reciprocal_rms), original_shape) + # HF casts before gamma, not before variance or after gamma multiplication. + normalized = op.CastLike(normalized, hidden_states) + return op.Mul(normalized, op.CastLike(self.weight, hidden_states)) + + class ScaleFreeRMSNorm(nn.Module): """RMSNorm with a constant all-ones scale (no learnable parameter). diff --git a/src/mobius/components/_rms_norm_test.py b/src/mobius/components/_rms_norm_test.py index 6c1ce9551..c9d432098 100644 --- a/src/mobius/components/_rms_norm_test.py +++ b/src/mobius/components/_rms_norm_test.py @@ -5,10 +5,24 @@ from __future__ import annotations +import ml_dtypes +import numpy as np +import onnx_ir as ir +import onnxruntime as ort +import pytest +import torch from onnxscript import nn +from mobius._builder import _cast_module_dtype +from mobius._flags import override_flags from mobius._testing import count_op_type, create_test_builder, create_test_input -from mobius.components._rms_norm import RMSNorm, apply_rms_norm +from mobius._testing.ort_inference import OnnxModelSession +from mobius.components._rms_norm import ( + GatedRMSNorm, + GroupedGatedRMSNorm, + RMSNorm, + apply_rms_norm, +) class TestRMSNorm: @@ -70,3 +84,136 @@ def test_multiple_rms_norms_in_same_graph(self): assert r1 is not None assert r2 is not None assert count_op_type(graph, "RMSNormalization") == 2 + + +def _grouped_gated_norm_model(shape, group_size, dtype): + norm = GroupedGatedRMSNorm(shape[-1], group_size, eps=1e-5) + _cast_module_dtype(norm, dtype) + builder, op, graph = create_test_builder() + hidden = create_test_input(builder, "hidden", shape, dtype) + gate = create_test_input(builder, "gate", shape, dtype) + builder.add_output(norm(op, hidden, gate), "out") + return ir.Model(graph, ir_version=11) + + +@pytest.mark.parametrize( + "dtype", [ir.DataType.FLOAT, ir.DataType.FLOAT16, ir.DataType.BFLOAT16] +) +@pytest.mark.parametrize("workaround", [False, True]) +def test_grouped_gated_norm_precision_graph(dtype, workaround): + with override_flags(ort_cuda_grouped_rmsnorm_workaround=workaround): + model = _grouped_gated_norm_model([2, 3, 24], 8, dtype) + graph = model.graph + assert [ + name for name, value in graph.initializers.items() if value.const_value is None + ] == ["weight"] + assert graph.initializers["weight"].dtype == dtype + assert not any(node.op_type == "RMSNormalization" for node in graph) + variance = next(node for node in graph if node.op_type == "ReduceMean") + square = variance.inputs[0].producer() + assert square.op_type == "Mul" + grouped = square.inputs[0].producer() + assert grouped.op_type == "Reshape" + gated = grouped.inputs[0].producer() + assert gated.op_type == "Mul" + assert gated.inputs[0].producer().attributes["to"].as_int() == ir.DataType.FLOAT + swish = gated.inputs[1].producer() + assert swish.op_type == "Swish" + assert swish.inputs[0].producer().attributes["to"].as_int() == ir.DataType.FLOAT + gamma_mul = graph.outputs[0].producer() + assert gamma_mul.op_type == "Mul" + # The rounding boundary is after normalization but before the native gamma multiply. + for value in gamma_mul.inputs: + cast = value.producer() + assert cast.op_type == "CastLike" + assert cast.inputs[1] is graph.inputs[0] + assert gamma_mul.inputs[0].producer().inputs[0].producer().op_type == "Reshape" + + +@pytest.mark.parametrize( + "shape,group_size", + [ + ([2, 24], 8), + ([2, 3, 24], 8), + ([2, 3, 24], 24), + ([2, 3, 24], 1), + ([2, 3, 4096], 512), + ], +) +@pytest.mark.parametrize( + "dtype", [ir.DataType.FLOAT, ir.DataType.FLOAT16, ir.DataType.BFLOAT16] +) +def test_grouped_gated_norm_source_precision_parity(tmp_path, shape, group_size, dtype): + if ( + dtype == ir.DataType.BFLOAT16 + and "CUDAExecutionProvider" not in ort.get_available_providers() + ): + pytest.skip("BF16 grouped norm arithmetic requires CUDA execution for this probe") + model = _grouped_gated_norm_model(shape, group_size, dtype) + np_dtype, torch_dtype = { + ir.DataType.FLOAT: (np.float32, torch.float32), + ir.DataType.FLOAT16: (np.float16, torch.float16), + ir.DataType.BFLOAT16: (ml_dtypes.bfloat16, torch.bfloat16), + }[dtype] + rng = np.random.default_rng(761) + # Distinct magnitudes and nonuniform gamma expose cross-group scale indexing errors. + magnitudes = np.geomspace(0.01, 20, shape[-1] // group_size).repeat(group_size) + hidden_np = (rng.normal(size=shape) * magnitudes).astype(np_dtype) + gate_np = rng.normal(size=shape).astype(np_dtype) + weight_np = rng.uniform(0.1, 2, shape[-1]).astype(np_dtype) + model.graph.initializers["weight"].const_value = ir.tensor(weight_np) + hidden = torch.from_numpy(hidden_np.astype(np.float32)).to(torch_dtype) + gate = torch.from_numpy(gate_np.astype(np.float32)).to(torch_dtype) + weight = torch.from_numpy(weight_np.astype(np.float32)).to(torch_dtype) + # Pinned HF Zamba2RMSNormGated: FP32 SiLU/variance, cast normalized x, then gamma. + gated = hidden.float() * torch.nn.functional.silu(gate.float()) + grouped = gated.reshape(*shape[:-1], -1, group_size) + normalized = grouped * torch.rsqrt(grouped.square().mean(-1, keepdim=True) + 1e-5) + expected = weight * normalized.reshape(shape).to(torch_dtype) + device = "cuda" if dtype == ir.DataType.BFLOAT16 else "cpu" + optimized_path = tmp_path / "optimized.onnx" + session = OnnxModelSession( + model, device=device, optimized_model_filepath=str(optimized_path) + ) + try: + actual = session.run({"hidden": hidden_np, "gate": gate_np})["out"] + optimized = ir.load(optimized_path) + assert all( + node.inputs[1].shape.rank() == 1 + for node in optimized.graph + if node.op_type == "RMSNormalization" + ) + finally: + session.close() + tolerance = { + ir.DataType.FLOAT: 2e-6, + ir.DataType.FLOAT16: 1e-3, + ir.DataType.BFLOAT16: 1e-2, + }[dtype] + np.testing.assert_allclose( + actual.astype(np.float32), expected.float().numpy(), rtol=tolerance, atol=tolerance + ) + if dtype != ir.DataType.FLOAT: + wrong_early_cast = gated.to(torch_dtype).float().reshape_as(grouped) + wrong_early_cast *= torch.rsqrt( + wrong_early_cast.square().mean(-1, keepdim=True) + 1e-5 + ) + wrong_early_cast = weight * wrong_early_cast.reshape(shape).to(torch_dtype) + wrong_late_cast = (normalized.reshape(shape) * weight.float()).to(torch_dtype) + assert torch.any(expected != wrong_early_cast) + assert torch.any(expected != wrong_late_cast) + + +@pytest.mark.parametrize("hidden,group", [(0, 4), (8, 0), (8, -1), (8, 3)]) +def test_grouped_gated_norm_rejects_invalid_groups(hidden, group): + with pytest.raises(ValueError, match="group_size"): + GroupedGatedRMSNorm(hidden, group) + + +def test_default_gated_norm_retains_existing_grouped_operator(): + with override_flags(ort_cuda_grouped_rmsnorm_workaround=False): + builder, op, graph = create_test_builder() + hidden = create_test_input(builder, "hidden", [2, 3, 24]) + gate = create_test_input(builder, "gate", [2, 3, 24]) + GatedRMSNorm(24, group_size=8)(op, hidden, gate) + assert count_op_type(graph, "RMSNormalization") == 1 diff --git a/src/mobius/models/_nemotron_h_test.py b/src/mobius/models/_nemotron_h_test.py index c2430ab35..1fd48aba8 100644 --- a/src/mobius/models/_nemotron_h_test.py +++ b/src/mobius/models/_nemotron_h_test.py @@ -18,6 +18,7 @@ from mobius._builder import _cast_module_dtype from mobius._configs import NemotronHConfig from mobius._testing import create_test_builder, create_test_input +from mobius.components import GroupedGatedRMSNorm from mobius.integrations._weight_loading import apply_weights from mobius.models.nemotron_h import ( NemotronHCausalLMModel, @@ -177,6 +178,32 @@ def test_mamba_cache_abi_preserves_configured_float32_ssm(dtype): assert len(dtypes) <= 1, (node.name, dtypes) +@pytest.mark.parametrize( + "dtype", [ir.DataType.FLOAT, ir.DataType.FLOAT16, ir.DataType.BFLOAT16] +) +def test_nemotron_mamba_uses_source_precision_grouped_norm(dtype): + config = _tiny_config(dtype=dtype) + module = NemotronHCausalLMModel(config) + norm = module.model.layers[0].mamba.norm + assert isinstance(norm, GroupedGatedRMSNorm) + package = build_from_module(module, config, "hybrid-text-generation") + graph = package["model"].graph + assert "model.layers.0.mamba.norm.weight" in graph.initializers + # Other decoder norms may fuse, but none uses a two-dimensional RMS scale. + norms = [node for node in graph if node.op_type == "RMSNormalization"] + assert norms + assert all(node.inputs[1].shape.rank() == 1 for node in norms) + gamma_mul = next( + node + for node in graph + if node.op_type == "Mul" + and node.inputs[1] is graph.initializers["model.layers.0.mamba.norm.weight"] + ) + assert gamma_mul.inputs[0].dtype == dtype + if dtype != ir.DataType.FLOAT: + assert gamma_mul.inputs[0].producer().op_type in {"Cast", "CastLike"} + + def test_expert_routing_relu2_and_sigmoid_matches_independent_reference(tmp_path): config = _tiny_config() block = NemotronHMoEBlock(config) diff --git a/src/mobius/models/nemotron_h.py b/src/mobius/models/nemotron_h.py index c6207e2b4..2e1e76b28 100644 --- a/src/mobius/models/nemotron_h.py +++ b/src/mobius/models/nemotron_h.py @@ -41,6 +41,7 @@ FCMLP, Attention, Embedding, + GroupedGatedRMSNorm, Linear, Mamba2Block, RMSNorm, @@ -85,6 +86,11 @@ def __init__(self, config: NemotronHConfig): norm_group_size=d_inner // config.mamba_n_groups, time_step_min=config.mamba_time_step_min, ) + # HF Zamba2RMSNormGated keeps gating/variance in FP32 and casts before gamma. + # Preserve the existing norm.weight namespace without the CUDA 2D-scale op. + self.mamba.norm = GroupedGatedRMSNorm( + d_inner, d_inner // config.mamba_n_groups, eps=config.rms_norm_eps + ) self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) def forward( From 4ad49ac438af8024f0a501a3dd3f17b53762c5b9 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 8 Oct 2026 14:12:55 -0700 Subject: [PATCH 10/15] Refresh Nemotron graph census after source-faithful grouped norm Measure raw and optimized graphs using the pinned GGUF config and existing CPU export route. Preserve exact census assertions, including the two additional decomposition constants, and regenerate the linked catalog/evidence. Runtime correctness remains blocked by the reported full-model decoder mismatch. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> --- docs/gguf-capability-catalog.md | 2 +- src/mobius/integrations/gguf/_quant_capabilities_test.py | 6 +++--- src/mobius/integrations/gguf/_runtime_blocker_evidence.py | 8 ++++---- src/mobius/integrations/gguf/_runtime_evidence_test.py | 4 ++-- testdata/evidence/gguf_quantization_capabilities.json | 8 ++++---- 5 files changed, 14 insertions(+), 14 deletions(-) diff --git a/docs/gguf-capability-catalog.md b/docs/gguf-capability-catalog.md index 58978c56b..6c3275601 100644 --- a/docs/gguf-capability-catalog.md +++ b/docs/gguf-capability-catalog.md @@ -55,7 +55,7 @@ Runtime support above is independent from tokenizer materialization support belo | Evidence ID | Pinned candidate | Bounded result | Withheld runtime claims | |---|---|---|---| -| `nemotron-h-moe-30b-iq2-xxs-runtime-blocker` | `bartowski/nvidia_Nemotron-3-Nano-30B-A3B-GGUF@1fc64d5b160654ec892df2708aa893b0e96e6491`
`nvidia_Nemotron-3-Nano-30B-A3B-IQ2_XXS.gguf`
18,010,755,296 B
`f3da710c046ce7cc6ff28a9b5f1a9153ac72e3f60603e51c7bb679d80716b58a` | config/tokenizer `nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16@bf77c3174f68ad409e1c2aa60daeb46e32d1c606`; GGUF tokenizer metadata `6089bcaf08b3fe0d49379ca7e85bd3c93e8705bac6130636425c159212971225`; result=blocked; 401 tensors / 31,577,940,288 parameters; ORT 1.29.0 / CPUExecutionProvider; ort-genai 0.15.2; The smallest model GGUF in the pinned 29-file repository revision is 18,010,755,296 bytes, above the 16 GiB bounded-artifact policy; explicit float16/float32 weights require 63,155,880,576/126,311,761,152 bytes.
The normal optimized CPU export has 40,109 nodes (43,180 before Mobius optimization) and 6,028 MatMul nodes because its truthful ReLU2 routed experts remain an ONNX loop with FP32 routing/accumulation. This graph-only census was refreshed locally with ORT 1.30.0 installed; the pinned ORT 1.29 runtime inspection below is historical, not a new runtime result. ORT 1.29 MoE/QMoE exposes ReLU but not ReLU2. QMoE has separate router_probs/router_weights for correction-biased selection with unbiased sigmoid mixing, while shared experts and optional latent projections can surround it; those are not fused-op blockers. The pinned 30B has no latent projection. QMoE does not support GGUF IQ2_XXS storage, and Mobius has no proven IQ2_XXS packer/kernel path.
ORT GenAI 0.15.2 discovers sparse/nonconsecutive KV and conv/recurrent slots, but derives recurrent_state names while this export uses ssm_state, does not beam-reorder recurrent state, and rejects nonzero recurrent-state rewind.
The GGUF tokenizer declares pre=pixtral, whose compiled llama.cpp behavior is not serialized; exact ORT tokenizer materialization is unavailable. | cached decode and deterministic generation, full-logit parity, package and report roundtrip, state replay, rollback, and reorder, tensor value closure | +| `nemotron-h-moe-30b-iq2-xxs-runtime-blocker` | `bartowski/nvidia_Nemotron-3-Nano-30B-A3B-GGUF@1fc64d5b160654ec892df2708aa893b0e96e6491`
`nvidia_Nemotron-3-Nano-30B-A3B-IQ2_XXS.gguf`
18,010,755,296 B
`f3da710c046ce7cc6ff28a9b5f1a9153ac72e3f60603e51c7bb679d80716b58a` | config/tokenizer `nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16@bf77c3174f68ad409e1c2aa60daeb46e32d1c606`; GGUF tokenizer metadata `6089bcaf08b3fe0d49379ca7e85bd3c93e8705bac6130636425c159212971225`; result=blocked; 401 tensors / 31,577,940,288 parameters; ORT 1.29.0 / CPUExecutionProvider; ort-genai 0.15.2; The smallest model GGUF in the pinned 29-file repository revision is 18,010,755,296 bytes, above the 16 GiB bounded-artifact policy; explicit float16/float32 weights require 63,155,880,576/126,311,761,152 bytes.
The normal optimized CPU export has 40,222 nodes (43,272 before Mobius optimization) and 6,028 MatMul nodes because its truthful ReLU2 routed experts remain an ONNX loop with FP32 routing/accumulation. This graph-only census was refreshed locally with ORT 1.30.0 installed; the pinned ORT 1.29 runtime inspection below is historical, not a new runtime result. ORT 1.29 MoE/QMoE exposes ReLU but not ReLU2. QMoE has separate router_probs/router_weights for correction-biased selection with unbiased sigmoid mixing, while shared experts and optional latent projections can surround it; those are not fused-op blockers. The pinned 30B has no latent projection. QMoE does not support GGUF IQ2_XXS storage, and Mobius has no proven IQ2_XXS packer/kernel path.
ORT GenAI 0.15.2 discovers sparse/nonconsecutive KV and conv/recurrent slots, but derives recurrent_state names while this export uses ssm_state, does not beam-reorder recurrent state, and rejects nonzero recurrent-state rewind.
The GGUF tokenizer declares pre=pixtral, whose compiled llama.cpp behavior is not serialized; exact ORT tokenizer materialization is unavailable. | cached decode and deterministic generation, full-logit parity, package and report roundtrip, state replay, rollback, and reorder, tensor value closure | | `glm-dsa-iq1-s-artifact-budget-blocker` | `unsloth/GLM-5.2-GGUF@abc55e72527792c6e77069c99b4cb7de16fa9f23`
`UD-IQ1_S/GLM-5.2-UD-IQ1_S-00001-of-00006.gguf` 9,423,744 B `46b6148389219ae45167cb8124fbb18ef7d432daf619b4faf9e06ea80d3f4777`
`UD-IQ1_S/GLM-5.2-UD-IQ1_S-00002-of-00006.gguf` 49,208,128,256 B `f2180207285e04fcaa5b8c53ba6e77ad5cc58666b6e7c6b04a5eded3fe8bef09`
`UD-IQ1_S/GLM-5.2-UD-IQ1_S-00003-of-00006.gguf` 49,684,417,024 B `b1c0c5a302cc8d5d9ea0bcd4467c01db72c26839f820f7e882079582ea0a8d2b`
`UD-IQ1_S/GLM-5.2-UD-IQ1_S-00004-of-00006.gguf` 49,396,052,864 B `a6a42da6975e29f89866dcde2956e9e50e6ea26635fb5063b74f3973f4f863b6`
`UD-IQ1_S/GLM-5.2-UD-IQ1_S-00005-of-00006.gguf` 49,246,275,936 B `a4a9851a50db533f21ef824e5d8038f04e6782e7d602d18e5fdd6643f68ccccb`
`UD-IQ1_S/GLM-5.2-UD-IQ1_S-00006-of-00006.gguf` 19,171,063,136 B `3b767f55df64e0432d52fcf1a14eb47a1ef3bbc91339e2ae220f38602237d7d7`
total 216,715,360,960 B | blocked; The smallest mainline GLM-5.2 GGUF is a 216,715,360,960-byte split set whose trunk also declares an unsupported routed DSA/MLA MTP block. | real-weight full-logit parity, runtime packaging, deterministic generation | | `minimax-m2-iq1-s-artifact-budget-blocker` | `mradermacher/MiniMax-M2-i1-GGUF@2d4f9b1a86d32ce4dfc47db312c8d6fcae8d7b37`
`MiniMax-M2.i1-IQ1_S.gguf` 46,514,882,176 B `7bae986e3cd380c28c6177d612fce1d52373241a0dfa13a6fca25de79abf15fb`
total 46,514,882,176 B | blocked; The smallest immutable public MiniMax-M2 GGUF is 46,514,882,176 bytes; real-weight parity cannot enter the 16 GiB bounded evidence set. | real-weight full-logit parity, runtime packaging, deterministic generation | | `mistral4-iq1-m-artifact-budget-blocker` | `unsloth/Mistral-Small-4-119B-2603-GGUF@bd93c721735aa32c035c0f19e738cb3371fd56ff`
`Mistral-Small-4-119B-2603-UD-IQ1_M.gguf` 32,306,941,632 B `40fcdee4869110938638c6b8bac253f442b196518d4623f4afdf4b885cd961c7`
total 32,306,941,632 B | blocked; The smallest immutable public Mistral4 GGUF is 32,306,941,632 bytes; real-weight parity cannot enter the 16 GiB bounded evidence set. | real-weight full-logit parity, runtime packaging, deterministic generation | diff --git a/src/mobius/integrations/gguf/_quant_capabilities_test.py b/src/mobius/integrations/gguf/_quant_capabilities_test.py index e60fd5e8c..51f60c131 100644 --- a/src/mobius/integrations/gguf/_quant_capabilities_test.py +++ b/src/mobius/integrations/gguf/_quant_capabilities_test.py @@ -42,9 +42,9 @@ def test_runtime_blocker_candidate_is_metadata_only_and_not_budgeted_as_support( assert record["size"] > policy["max_selected_artifact_bytes"] assert record["size"] not in {artifact["size"] for artifact in selected_artifacts} assert record["tokenizer"]["revision"] == "bf77c3174f68ad409e1c2aa60daeb46e32d1c606" - assert record["graph"]["pre_optimization_node_count"] == 43_180 - assert record["graph"]["node_count"] == 40_109 - assert record["graph"]["initializer_count"] == 6_255 + assert record["graph"]["pre_optimization_node_count"] == 43_272 + assert record["graph"]["node_count"] == 40_222 + assert record["graph"]["initializer_count"] == 6_257 assert record["graph"]["matmul_count"] == 6_028 assert record["graph"]["state_slots"]["mamba2.ssm_state"] == 23 assert record["runtime_schema_issue"].endswith("/issues/605") diff --git a/src/mobius/integrations/gguf/_runtime_blocker_evidence.py b/src/mobius/integrations/gguf/_runtime_blocker_evidence.py index 5d1f294a2..f624b1026 100644 --- a/src/mobius/integrations/gguf/_runtime_blocker_evidence.py +++ b/src/mobius/integrations/gguf/_runtime_blocker_evidence.py @@ -200,9 +200,9 @@ def __post_init__(self) -> None: expert_count=128, experts_per_token=6, layer_counts=(("full_attention", 6), ("mamba2", 23), ("moe", 23)), - pre_optimization_graph_node_count=43_180, - graph_node_count=40_109, - graph_initializer_count=6_255, + pre_optimization_graph_node_count=43_272, + graph_node_count=40_222, + graph_initializer_count=6_257, graph_matmul_count=6_028, state_slots=( ("attention.key", 6), @@ -222,7 +222,7 @@ def __post_init__(self) -> None: "float16/float32 weights require 63,155,880,576/126,311,761,152 bytes." ), ( - "The normal optimized CPU export has 40,109 nodes (43,180 before Mobius " + "The normal optimized CPU export has 40,222 nodes (43,272 before Mobius " "optimization) and 6,028 MatMul nodes because its truthful ReLU2 routed " "experts remain an ONNX loop with FP32 routing/accumulation. " "This graph-only census was refreshed locally with ORT 1.30.0 installed; " diff --git a/src/mobius/integrations/gguf/_runtime_evidence_test.py b/src/mobius/integrations/gguf/_runtime_evidence_test.py index 160c51481..d767284d1 100644 --- a/src/mobius/integrations/gguf/_runtime_evidence_test.py +++ b/src/mobius/integrations/gguf/_runtime_evidence_test.py @@ -215,8 +215,8 @@ def test_nemotron_h_runtime_blocker_is_pinned_without_support_claim() -> None: assert record.runtime_schema_issue.endswith("/issues/605") assert _runtime_evidence.runtime_evidence(record.evidence_id) is None assert "full-logit parity" in record.withheld_checks - assert record.graph_node_count == 40_109 - assert record.pre_optimization_graph_node_count == 43_180 + assert record.graph_node_count == 40_222 + assert record.pre_optimization_graph_node_count == 43_272 assert "separate router_probs/router_weights" in record.blockers[1] assert "not fused-op blockers" in record.blockers[1] assert "has no latent projection" in record.blockers[1] diff --git a/testdata/evidence/gguf_quantization_capabilities.json b/testdata/evidence/gguf_quantization_capabilities.json index bad8ebf7a..06ae91ebc 100644 --- a/testdata/evidence/gguf_quantization_capabilities.json +++ b/testdata/evidence/gguf_quantization_capabilities.json @@ -2842,7 +2842,7 @@ "architecture": "nemotron_h_moe", "blockers": [ "The smallest model GGUF in the pinned 29-file repository revision is 18,010,755,296 bytes, above the 16 GiB bounded-artifact policy; explicit float16/float32 weights require 63,155,880,576/126,311,761,152 bytes.", - "The normal optimized CPU export has 40,109 nodes (43,180 before Mobius optimization) and 6,028 MatMul nodes because its truthful ReLU2 routed experts remain an ONNX loop with FP32 routing/accumulation. This graph-only census was refreshed locally with ORT 1.30.0 installed; the pinned ORT 1.29 runtime inspection below is historical, not a new runtime result. ORT 1.29 MoE/QMoE exposes ReLU but not ReLU2. QMoE has separate router_probs/router_weights for correction-biased selection with unbiased sigmoid mixing, while shared experts and optional latent projections can surround it; those are not fused-op blockers. The pinned 30B has no latent projection. QMoE does not support GGUF IQ2_XXS storage, and Mobius has no proven IQ2_XXS packer/kernel path.", + "The normal optimized CPU export has 40,222 nodes (43,272 before Mobius optimization) and 6,028 MatMul nodes because its truthful ReLU2 routed experts remain an ONNX loop with FP32 routing/accumulation. This graph-only census was refreshed locally with ORT 1.30.0 installed; the pinned ORT 1.29 runtime inspection below is historical, not a new runtime result. ORT 1.29 MoE/QMoE exposes ReLU but not ReLU2. QMoE has separate router_probs/router_weights for correction-biased selection with unbiased sigmoid mixing, while shared experts and optional latent projections can surround it; those are not fused-op blockers. The pinned 30B has no latent projection. QMoE does not support GGUF IQ2_XXS storage, and Mobius has no proven IQ2_XXS packer/kernel path.", "ORT GenAI 0.15.2 discovers sparse/nonconsecutive KV and conv/recurrent slots, but derives recurrent_state names while this export uses ssm_state, does not beam-reorder recurrent state, and rejects nonzero recurrent-state rewind.", "The GGUF tokenizer declares pre=pixtral, whose compiled llama.cpp behavior is not serialized; exact ORT tokenizer materialization is unavailable." ], @@ -2861,10 +2861,10 @@ "explicit_float32_bytes": 126311761152, "filename": "nvidia_Nemotron-3-Nano-30B-A3B-IQ2_XXS.gguf", "graph": { - "initializer_count": 6255, + "initializer_count": 6257, "matmul_count": 6028, - "node_count": 40109, - "pre_optimization_node_count": 43180, + "node_count": 40222, + "pre_optimization_node_count": 43272, "state_slots": { "attention.key": 6, "attention.value": 6, From 4ba865c0dec8ae1e2640f6691ca248a1fa51cb23 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 8 Oct 2026 14:15:36 -0700 Subject: [PATCH 11/15] Preserve the Lightning final MoE forward override contract Accept and forward the base layer's optional attention, position and cache parameters while retaining the draft's existing two-argument call. Add a positional forwarding regression; independent MTP parity and cached decode tests remain passing. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> --- src/mobius/models/_nemotron_h_mtp_test.py | 25 +++++++++++++++++++++++ src/mobius/models/nemotron_h_mtp.py | 13 ++++++++++-- 2 files changed, 36 insertions(+), 2 deletions(-) diff --git a/src/mobius/models/_nemotron_h_mtp_test.py b/src/mobius/models/_nemotron_h_mtp_test.py index 0ab50d8aa..0a1b75577 100644 --- a/src/mobius/models/_nemotron_h_mtp_test.py +++ b/src/mobius/models/_nemotron_h_mtp_test.py @@ -21,14 +21,17 @@ from mobius._builder import build_from_module from mobius._configs import NemotronHConfig from mobius._model_package import ModelPackage +from mobius._testing import create_test_builder, create_test_input from mobius.integrations.transformers._nemotron_h_weights_test import ( _config, _source_name, _write_mixed_checkpoint, ) +from mobius.models.nemotron_h import NemotronHMoELayer from mobius.models.nemotron_h_mtp import ( NemotronHMtpModel, NemotronHSpeculativeModel, + _FinalMoELayer, mtp_source_name, ) @@ -54,6 +57,28 @@ def _graph_package(dtype="f32"): return build_from_module(NemotronHSpeculativeModel(config), config, task="nemotron-h-mtp") +def test_final_moe_layer_accepts_and_forwards_base_positional_contract(monkeypatch): + config = NemotronHConfig.from_transformers(SimpleNamespace(**_raw())) + layer = _FinalMoELayer(config) + builder, op, _graph = create_test_builder() + hidden = create_test_input(builder, "hidden", [1, 2, config.hidden_size]) + received = [] + + def base_forward( + self, op, hidden_states, attention_bias, position_embeddings, past_key_value + ): + received.append((attention_bias, position_embeddings, past_key_value)) + return hidden_states, (None, None) + + monkeypatch.setattr(NemotronHMoELayer, "forward", base_forward) + layer(op, hidden) + bias = create_test_input(builder, "bias", [1, 1, 2, 2]) + positions = (hidden, hidden) + cache = (hidden, hidden) + layer(op, hidden, bias, positions, cache) + assert received == [(None, None, None), (bias, positions, cache)] + + def _checkpoint(tmp_path, *, dtype=torch.float32): raw = _raw() (tmp_path / "config.json").write_text(json.dumps(raw), encoding="utf-8") diff --git a/src/mobius/models/nemotron_h_mtp.py b/src/mobius/models/nemotron_h_mtp.py index 12e9bb131..0b3bb8650 100644 --- a/src/mobius/models/nemotron_h_mtp.py +++ b/src/mobius/models/nemotron_h_mtp.py @@ -77,8 +77,17 @@ def __init__(self, config: NemotronHConfig): super().__init__(config) self.final_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - def forward(self, op: OpBuilder, hidden_states: ir.Value): - hidden_states, _ = super().forward(op, hidden_states, None, None, None) + def forward( + self, + op: OpBuilder, + hidden_states: ir.Value, + attention_bias: ir.Value | None = None, + position_embeddings: tuple | None = None, + past_key_value: tuple | None = None, + ): + hidden_states, _ = super().forward( + op, hidden_states, attention_bias, position_embeddings, past_key_value + ) return self.final_layernorm(op, hidden_states) From a72b658f0874c8d8761bc64f5222a27d0e3dc1c7 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 8 Oct 2026 14:17:34 -0700 Subject: [PATCH 12/15] Keep local checkpoint paths out of persisted streaming reports Use a portable local-safetensors-checkpoint source label in dense streaming reports and ONNX metadata. Transient lazy bindings and overlap protections retain actual paths. Cover reconstructed target exports and both MTP component reports through save/reload, preserving remote IDs. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> --- docs/nemotron-lightning-export.md | 3 +++ src/mobius/integrations/_weight_loading.py | 4 +++- .../integrations/transformers/_nemotron_h_weights_test.py | 4 ++++ src/mobius/models/_nemotron_h_mtp_test.py | 7 +++++++ 4 files changed, 17 insertions(+), 1 deletion(-) diff --git a/docs/nemotron-lightning-export.md b/docs/nemotron-lightning-export.md index b53bf31a7..3745b0075 100644 --- a/docs/nemotron-lightning-export.md +++ b/docs/nemotron-lightning-export.md @@ -119,6 +119,9 @@ implies reconstruction; both remain explicit. resolved blob directories. Saving into source checkpoint directories or onto source-file aliases is rejected before serialization; use a fresh, separate output directory. + Local checkpoint paths remain only in transient lazy-source bindings and + overlap guards; persisted loading reports and ONNX metadata use + `local-safetensors-checkpoint`, including both target/MTP component reports. - `models/nemotron_h.py` casts router weights/bias to FP32 and accumulates routed outputs in FP32 before the shared-expert path, retaining sigmoid selection, unbiased routing weights and squared ReLU. It also retains FP32 router diff --git a/src/mobius/integrations/_weight_loading.py b/src/mobius/integrations/_weight_loading.py index f4fb99847..145493b7b 100644 --- a/src/mobius/integrations/_weight_loading.py +++ b/src/mobius/integrations/_weight_loading.py @@ -1243,7 +1243,9 @@ def validate_source( report = { "format": "mobius.weight-loading-report.v1", - "source": model_id, + "source": ( + "local-safetensors-checkpoint" if pathlib.Path(model_id).is_dir() else model_id + ), "revision": revision, "output_weight_format": "dense", "native_fp8": False, diff --git a/src/mobius/integrations/transformers/_nemotron_h_weights_test.py b/src/mobius/integrations/transformers/_nemotron_h_weights_test.py index 13e8f64cd..78b468b45 100644 --- a/src/mobius/integrations/transformers/_nemotron_h_weights_test.py +++ b/src/mobius/integrations/transformers/_nemotron_h_weights_test.py @@ -153,6 +153,7 @@ def test_public_local_build_reconstructs_all_projection_families(tmp_path, monke ) package = build(str(tmp_path), keep_quantized=False, dtype="bf16") report = package.weight_loading_report + assert report["source"] == "local-safetensors-checkpoint" assert report["storage_policy"] == "explicit-dense-bf16-reconstruction" assert report["native_nvfp4"] is False assert report["ignored_tensors"] == 0 @@ -169,6 +170,9 @@ def test_public_local_build_reconstructs_all_projection_families(tmp_path, monke package.save(output, external_data="onnx") assert (output / "model.onnx").is_file() loaded = ir.load(output / "model.onnx") + assert json.loads(loaded.metadata_props["mobius.weight_loading"])["source"] == ( + "local-safetensors-checkpoint" + ) assert loaded.metadata_props["mobius.source_weight_format"] == "modelopt-mixed-fp8-nvfp4" diff --git a/src/mobius/models/_nemotron_h_mtp_test.py b/src/mobius/models/_nemotron_h_mtp_test.py index 0a1b75577..8ce1253a1 100644 --- a/src/mobius/models/_nemotron_h_mtp_test.py +++ b/src/mobius/models/_nemotron_h_mtp_test.py @@ -304,6 +304,9 @@ def test_package_accounting_save_and_default_guard(tmp_path): build(str(tmp_path), dtype="bf16", load_weights=False) package = build(str(tmp_path), task="nemotron-h-mtp", dtype="bf16") report = package.weight_loading_report + assert report["source"] == "local-safetensors-checkpoint" + assert report["decoder_component"]["source"] == "local-safetensors-checkpoint" + assert report["mtp_component"]["source"] == "local-safetensors-checkpoint" assert report["mtp_preserved"] is True assert report["omitted_mtp_tensor_count"] == 0 assert set(report["mtp_tensors"]) == {n for n in tensors if n.startswith("mtp.")} @@ -325,9 +328,13 @@ def test_package_accounting_save_and_default_guard(tmp_path): restored = ModelPackage.load(output) assert set(restored) == {"decoder", "mtp"} for model in restored.values(): + loading = json.loads(model.metadata_props["mobius.weight_loading"]) + assert loading["source"] == "local-safetensors-checkpoint" + assert str(tmp_path) not in json.dumps(loading) assert model.metadata_props["mobius.source_num_nextn_predict_layers"] == "1" assert model.metadata_props["mobius.mtp_contract"] == "nemotron-h-lightning-nextn@1" assert restored.weight_loading_report["mtp_preserved"] is True + assert str(tmp_path) not in json.dumps(restored.weight_loading_report) @pytest.mark.parametrize( From 089c161b8a68ec9c6807af9309293c46e256e83b Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Sun, 11 Oct 2026 09:37:40 -0700 Subject: [PATCH 13/15] Clarify Lightning draft ABI and updated H100 validation limits Declare physically dedicated draft tables despite shared checkpoint sources and keep the draft implementation private behind the public paired wrapper. Add metadata/public-task regressions. Document the user-supplied norm and generation passes while retaining failed full decoder parity and reference-contract qualifications; no recurrence or threshold changes. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> --- docs/nemotron-lightning-export.md | 52 +++++++++++++++++-- .../integrations/ort_genai/auto_export.py | 8 ++- src/mobius/models/__init__.py | 3 +- src/mobius/models/_nemotron_h_mtp_test.py | 25 ++++++++- src/mobius/models/nemotron_h_mtp.py | 5 +- 5 files changed, 81 insertions(+), 12 deletions(-) diff --git a/docs/nemotron-lightning-export.md b/docs/nemotron-lightning-export.md index 3745b0075..c6b1fbf98 100644 --- a/docs/nemotron-lightning-export.md +++ b/docs/nemotron-lightning-export.md @@ -1,6 +1,6 @@ # Nemotron 3.5 Lightning: reconstruction and compatibility gates -**CPU-tested exporter improvements; full-model generation remains blocked.** +**CPU-tested exporter improvements; full-model numerical readiness remains blocked.** The [machine-readable recipe](../examples/nemotron_lightning/recipe.json) records immutable inputs and compatibility limits. Native NVFP4 remains unsupported. Explicit dense BF16 reconstruction is implemented and tested @@ -48,6 +48,43 @@ exports and checkpoint files must remain unchanged. The remaining decoder discrepancy needs a first-divergent-stage comparison and an inspected source reference; this correction alone is **not** a claim of resolved full parity. +### Updated user-supplied H100 evidence (2026-10-11) + +The supplied report at checkout `a72b658f0874c8d8761bc64f5222a27d0e3dc1c7` +and the original pinned BF16 revision reports **15 normalization cases passing, +including all five BF16 CUDA cases**, source fidelity passing for 6243 target +and 270 MTP tensors (6515 comparisons including shared copies), and exact +eight-token generation passing. This replaces the earlier generation failure +**for this run only**. Conditional MTP passes 12 comparisons, with maximum +NRMSE 0.0047783538309905 and minimum cosine 0.9999894304916273. + +Full decoder parity still fails **141 of 240 comparisons**, with maximum NRMSE +0.10049961190760874 and minimum cosine 0.9952473196233278, against unchanged +NRMSE <= 0.03 / cosine >= 0.999 / finite gates. Short token agreement and +conditional MTP do not clear this gate; speculative integration is NOT_RUN. +These are user-provided remote observations, not locally rerun H100 tests. + +The report separates precision amplification from identical-input operators: +prefill first fails the layer-4 convolution-state gate, while the first hidden +failure is layer 34. Identical-input audits of Mamba layers 0/4/16 show FP32 +recurrence readout NRMSE <= 6.1e-7 with the same timestep, near-exact grouped +norm, and matching native FP32 router expert sets on the same BF16 inputs. +Different FP32 upstream inputs change some layer-34 expert selections. +Layer-4 pre-Mamba norm agrees on the same exposed rounded BF16 residual, +qualified because fused SkipNorm may normalize before residual-output rounding. +None of these observations proves full-model parity or justifies a blanket +recurrence or normalization change. + +The native FP32 CPU-fallback reference audit agrees on prefill (112 comparisons, +maximum NRMSE 1.55e-6), but differs on cached steps. A **reference-only** +diagnostic omitting the timestep minimum clamp on cached updates matches native +fallback behavior; the baseline exporter and validation reference were unchanged. +This establishes a reference-contract question, not fused native BF16 CUDA proof +or permission to replace the acceptance reference. The report is structured +evidence, not the runnable `run_validation.py`, reference implementation, +input/cache fixtures, provider profile, or optimized graph; those remote +`/home/mobius/...` paths are not accessible in the local session. + The original preflight host was Windows ARM64 with a Qualcomm Adreno X1-85, no NVIDIA/CUDA GPU or `nvidia-smi`, Python 3.12.10, 68,171,038,720 bytes RAM and approximately 818 GB free disk. Mobius, PyTorch, Transformers, ORT and @@ -254,9 +291,16 @@ FP32-reference comparison at `rtol=atol=1e-2`. Tiny tests also cover mixed FP8/NVFP4 reconstruction with floating MTP and a shared reconstructed head, exact graph/cache I/O, f32/f16/bf16 graph dtypes, CLI serialization, strict source accounting and checkpoint-save protection. -BF16 execution/parity, full weighted 30B MTP, CUDA/H100, -performance, Olive requantization and L4/L5 real-weight goldens are **NOT_RUN**; -no full-model numerical or speedup claim is made. +Local full-weight BF16/CUDA execution, performance, Olive requantization and +L4/L5 real-weight goldens remain **NOT_RUN**. The separately supplied H100 +conditional MTP/generation results above are not full decoder parity or a +speedup claim. + +`NemotronHSpeculativeModel` is the public paired builder module; the draft +implementation is private and does not advertise a standalone public task. +External coordination metadata declares dedicated draft embedding/head tables +because the graph consumes token IDs and emits logits. Their checkpoint sources +are shared with the target, but their runtime storage/ABI is not `shared_target`. ORT-GenAI metadata can describe the external coordination contract and marks it `runtime_unvalidated`. The exporter does **not** implement speculative diff --git a/src/mobius/integrations/ort_genai/auto_export.py b/src/mobius/integrations/ort_genai/auto_export.py index 2d1908255..ed5a4fd4b 100644 --- a/src/mobius/integrations/ort_genai/auto_export.py +++ b/src/mobius/integrations/ort_genai/auto_export.py @@ -1943,13 +1943,19 @@ def _write_mtp_config(pkg: ModelPackage, directory: str) -> str | None: raise ValueError("MTP runtime metadata requires a target decoder graph") dedicated_embeddings = bool(getattr(proposer_config, "use_dedicated_embeddings", False)) dedicated_lm_head = bool(getattr(proposer_config, "use_dedicated_lm_head", False)) + lightning_contract = ( + mtp_model.metadata_props.get("mobius.mtp_contract") == "nemotron-h-lightning-nextn@1" + ) + if lightning_contract: + # Shared checkpoint sources are copied into physically dedicated draft tables. + dedicated_embeddings = dedicated_lm_head = True conditioning = { "target_hidden_output": "mtp_seed", "target_hidden_input": "hidden_states", "embedding": "dedicated" if dedicated_embeddings else "shared_target", "lm_head": "dedicated" if dedicated_lm_head else "shared_target", } - if mtp_model.metadata_props.get("mobius.mtp_contract") == "nemotron-h-lightning-nextn@1": + if lightning_contract: conditioning.update( target_hidden_normalization="post_final_norm", token_alignment="target_h_i_with_token_t_i_plus_1", diff --git a/src/mobius/models/__init__.py b/src/mobius/models/__init__.py index 0ed884c51..5c0710d66 100644 --- a/src/mobius/models/__init__.py +++ b/src/mobius/models/__init__.py @@ -138,7 +138,6 @@ "NanoChatCausalLMModel", "NemotronCausalLMModel", "NemotronHCausalLMModel", - "NemotronHMtpModel", "NemotronHSpeculativeModel", "NeoBertGGUFModel", "NomicBertGGUFModel", @@ -376,7 +375,7 @@ from mobius.models.nemotron import NemotronCausalLMModel from mobius.models.nemotron3_diarization import Nemotron3DiarizationModel from mobius.models.nemotron_h import NemotronHCausalLMModel -from mobius.models.nemotron_h_mtp import NemotronHMtpModel, NemotronHSpeculativeModel +from mobius.models.nemotron_h_mtp import NemotronHSpeculativeModel from mobius.models.nemotron_parse import NemotronParseForConditionalGeneration from mobius.models.olmo import OLMo2CausalLMModel, OLMoCausalLMModel from mobius.models.opt import OPTCausalLMModel diff --git a/src/mobius/models/_nemotron_h_mtp_test.py b/src/mobius/models/_nemotron_h_mtp_test.py index 8ce1253a1..219ca7de7 100644 --- a/src/mobius/models/_nemotron_h_mtp_test.py +++ b/src/mobius/models/_nemotron_h_mtp_test.py @@ -29,9 +29,9 @@ ) from mobius.models.nemotron_h import NemotronHMoELayer from mobius.models.nemotron_h_mtp import ( - NemotronHMtpModel, NemotronHSpeculativeModel, _FinalMoELayer, + _NemotronHMtpModel, mtp_source_name, ) @@ -57,6 +57,21 @@ def _graph_package(dtype="f32"): return build_from_module(NemotronHSpeculativeModel(config), config, task="nemotron-h-mtp") +def test_public_speculative_wrapper_builds_with_its_default_task(): + import mobius.models as models + + config = NemotronHConfig.from_transformers(SimpleNamespace(**_raw())) + module = models.NemotronHSpeculativeModel(config) + assert module.default_task == "nemotron-h-mtp" + assert set(build_from_module(module, config, task=module.default_task)) == { + "decoder", + "mtp", + } + assert "NemotronHMtpModel" not in models.__all__ + assert not hasattr(models, "NemotronHMtpModel") + assert not hasattr(module.mtp, "default_task") + + def test_final_moe_layer_accepts_and_forwards_base_positional_contract(monkeypatch): config = NemotronHConfig.from_transformers(SimpleNamespace(**_raw())) layer = _FinalMoELayer(config) @@ -431,6 +446,12 @@ def test_intrinsic_runtime_metadata_has_precise_external_contract(tmp_path): assert payload["conditioning"]["token_alignment"] == "target_h_i_with_token_t_i_plus_1" assert payload["conditioning"]["draft_prediction"] == "token_t_i_plus_2" assert payload["conditioning"]["tables"] == "target_shared_sources_copied_into_draft_graph" + assert payload["conditioning"]["embedding"] == "dedicated" + assert payload["conditioning"]["lm_head"] == "dedicated" + assert payload["shared_embedding"] is None + assert payload["shared_lm_head"] is None + assert "input_ids" in payload["inputs"] + assert "logits" in payload["outputs"] assert payload["cache_namespaces"]["target"]["ports"] assert payload["cache_namespaces"]["mtp"]["ports"] assert payload["num_nextn_predict_layers"] == 1 @@ -452,7 +473,7 @@ def test_pinned_production_mtp_parameter_shapes(): "moe_shared_expert_intermediate_size": 3712, } config = NemotronHConfig.from_transformers(SimpleNamespace(**raw)) - model = NemotronHMtpModel(config) + model = _NemotronHMtpModel(config) expected = { "mtp.layers.0.enorm.weight": [2688], "mtp.layers.0.hnorm.weight": [2688], diff --git a/src/mobius/models/nemotron_h_mtp.py b/src/mobius/models/nemotron_h_mtp.py index 0b3bb8650..aa8987bd9 100644 --- a/src/mobius/models/nemotron_h_mtp.py +++ b/src/mobius/models/nemotron_h_mtp.py @@ -91,7 +91,7 @@ def forward( return self.final_layernorm(op, hidden_states) -class NemotronHMtpModel(nn.Module): +class _NemotronHMtpModel(nn.Module): """Standalone draft computation with target-shared embedding and LM head. Inputs: token IDs for t_(i+1), post-final-norm target h_i, a padding mask, @@ -101,7 +101,6 @@ class NemotronHMtpModel(nn.Module): """ config_class: type = NemotronHConfig - default_task: str = "nemotron-h-mtp" category: str = "Mixture of Experts" def __init__(self, config: NemotronHConfig): @@ -161,7 +160,7 @@ def __init__(self, config: NemotronHConfig): self.config = config # The explicit package preserves auxiliary tensors in a separate graph. self.decoder = _MtpSeedTarget(config, target_decoder_only=True) - self.mtp = NemotronHMtpModel(config) + self.mtp = _NemotronHMtpModel(config) # Each component is its own graph root, not a nested parameter namespace. self.decoder._set_name("") self.mtp._set_name("") From 49d5f2c89e8794e5f7fc1d904ee962dc088caeb6 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Sun, 11 Oct 2026 09:38:03 -0700 Subject: [PATCH 14/15] test: isolate Nemotron Mamba cached reference precision contracts Exercise portable Mamba2 prefill and three reused-state steps with an independent FP32 recurrence, grouped B/C expansion, source-precision grouped normalization, and timesteps below the Lightning minimum. Preserve the distinction between bounded export semantics and the native fallback's unbounded cached timestep path instead of silently replacing the reference. Protect FP16/BF16 export cast boundaries for timestep addition, FP32 SSM state and readout/D-skip rounding. These regressions do not change production recurrence or establish full-model H100 numerical parity. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> --- src/mobius/components/_mamba_block_test.py | 180 +++++++++++++++++++++ 1 file changed, 180 insertions(+) diff --git a/src/mobius/components/_mamba_block_test.py b/src/mobius/components/_mamba_block_test.py index 3c5d948f8..14902009b 100644 --- a/src/mobius/components/_mamba_block_test.py +++ b/src/mobius/components/_mamba_block_test.py @@ -16,9 +16,189 @@ create_test_builder, create_test_input, ) +from mobius.components import GroupedGatedRMSNorm, Mamba2Block from mobius.components._mamba_block import MambaBlock, SequenceMambaBlock +@pytest.mark.parametrize("dtype", [ir.DataType.FLOAT16, ir.DataType.BFLOAT16]) +def test_mamba2_recurrence_precision_boundaries(dtype): + """Protect export precision; native BF16 fallback has different cast boundaries.""" + from mobius._builder import _cast_module_dtype + + block = Mamba2Block(4, 4, 2, 2, 3, time_step_min=0.001) + _cast_module_dtype(block, dtype) + builder, op, graph = create_test_builder() + hidden = builder.input("hidden", dtype=dtype, shape=[2, "seq", 4]) + conv = builder.input("conv", dtype=dtype, shape=[2, 10, 3]) + state = builder.input("state", dtype=ir.DataType.FLOAT, shape=[2, 2, 3, 2]) + block(op, hidden, conv, state) + + softplus = next(node for node in graph if node.op_type == "Softplus") + dt_add = softplus.inputs[0].producer() + assert dt_add.op_type == "Add" + # In contrast to the native BF16 fallback, both operands are widened BEFORE + # bias addition and softplus; this is not a claim of native backend parity. + for value in dt_add.inputs: + cast = value.producer() + assert cast.op_type == "Cast" + assert cast.attributes["to"].as_int() == int(ir.DataType.FLOAT) + clip = next(node for node in graph if node.op_type == "Clip") + assert clip.inputs[0] is softplus.outputs[0] + attention = next(node for node in graph if node.op_type == "LinearAttention") + assert attention.inputs[3].producer().op_type == "Cast" + assert attention.inputs[3].producer().attributes["to"].as_int() == int(ir.DataType.FLOAT) + readout_cast = next( + node + for node in graph + if node.op_type == "CastLike" and node.inputs[0] is attention.outputs[0] + ) + skip_add = next( + node + for node in graph + if node.op_type == "Add" and node.inputs[0] is readout_cast.outputs[0] + ) + assert readout_cast.inputs[1] is hidden + skip_cast = skip_add.inputs[1].producer() + assert skip_cast.op_type == "CastLike" and skip_cast.inputs[1] is hidden + assert skip_cast.inputs[0].producer().op_type == "Reshape" + d_product = skip_cast.inputs[0].producer().inputs[0].producer() + assert d_product.op_type == "Mul" + assert d_product.inputs[0].producer().op_type == "Reshape" + assert d_product.inputs[1].producer().inputs[0].producer().op_type == "Cast" + + +@pytest.mark.parametrize("n_groups", [1, 2]) +@pytest.mark.parametrize("time_step_min", [0.0, 0.001]) +def test_mamba2_prefill_and_cached_timestep_contract(tmp_path, n_groups, time_step_min): + """Exercise the bounded portable recurrence, not native cached-backend parity. + + Transformers Nemotron-H a96730c8c97b8efbf35bbaf7f5da33ec99231a49 + clamps dt in mamba2_chunk_scan, but not mamba2_selective_state_update. + The cached contrast below preserves that distinction without changing the + exporter or silently substituting a different reference. + """ + import onnxruntime as ort + import torch + import torch.nn.functional as torch_f + from onnx_ir.passes.common import InlinePass + + from mobius.functions import causal_conv_nd_with_state, linear_attention + from mobius.integrations._weight_loading import apply_weights + + d_model, heads, d_head, d_state, kernel = 4, 2, 2, 3, 3 + d_inner = heads * d_head + conv_dim = d_inner + 2 * n_groups * d_state + block = Mamba2Block( + d_model, + d_inner, + heads, + d_head, + d_state, + n_groups=n_groups, + conv_kernel=kernel, + time_step_min=time_step_min, + ) + block.norm = GroupedGatedRMSNorm(d_inner, group_size=d_head, eps=1e-5) + graph = ir.Graph( + [], [], nodes=[], name="mamba2", opset_imports={"": OPSET_VERSION, "com.microsoft": 1} + ) + builder = GraphBuilder(graph) + hidden = builder.input("hidden", dtype=ir.DataType.FLOAT, shape=[2, "seq", d_model]) + conv = builder.input("conv", dtype=ir.DataType.FLOAT, shape=[2, conv_dim, kernel - 1]) + state = builder.input("state", dtype=ir.DataType.FLOAT, shape=[2, heads, d_state, d_head]) + out, new_conv, new_state = block(builder.op, hidden, conv, state) + for value, name in [(out, "out"), (new_conv, "new_conv"), (new_state, "new_state")]: + builder.add_output(value, name) + timestep = next(node for node in graph if node.op_type == "Softplus").outputs[0] + if time_step_min: + timestep = next(node for node in graph if node.op_type == "Clip").outputs[0] + builder.add_output(timestep, "dt") + model = ir.Model(graph, ir_version=11) + for function in [ + causal_conv_nd_with_state( + kernel_size=kernel, channels=conv_dim, ndim=1, activation="silu" + ), + linear_attention(q_num_heads=heads, kv_num_heads=heads, update_rule="gated"), + ]: + model.functions[function.identifier()] = function + + rng = np.random.default_rng(43) + weights = { + name: torch.from_numpy(rng.normal(0, 0.3, list(p.shape)).astype(np.float32)) + for name, p in graph.initializers.items() + if p.const_value is None + } + # First head is below the Lightning minimum; second is safely above it. + weights["dt_bias"] = torch.tensor([-11.0, -3.0]) + weights["A_log"] = torch.tensor([3.0, 0.0]) + weights["D"] = torch.tensor([0.4, -0.3]) + weights["norm.weight"] = torch.tensor([0.7, 1.2, 0.9, 1.1]) + apply_weights(model, weights) + InlinePass()(model) + path = tmp_path / "mamba2.onnx" + ir.save(model, path) + session = ort.InferenceSession(str(path), providers=["CPUExecutionProvider"]) + + def reference(u, conv_state, ssm_state, minimum): + gate, xbc, raw_dt = torch.split( + u @ weights["in_proj.weight"].T, [d_inner, conv_dim, heads], dim=-1 + ) + history = torch.cat([conv_state, xbc.transpose(1, 2)], dim=-1) + activated = torch_f.silu( + torch_f.conv1d( + history, weights["conv1d.weight"], weights["conv1d.bias"], groups=conv_dim + ) + ).transpose(1, 2) + x, b, c = torch.split( + activated, [d_inner, n_groups * d_state, n_groups * d_state], dim=-1 + ) + dt = torch_f.softplus(raw_dt + weights["dt_bias"]).clamp(min=minimum) + b = b.reshape(2, -1, n_groups, d_state).repeat_interleave(heads // n_groups, 2) + c = c.reshape(2, -1, n_groups, d_state).repeat_interleave(heads // n_groups, 2) + x = x.reshape(2, -1, heads, d_head) + values = [] + for t in range(u.shape[1]): + delta = dt[:, t, :, None, None] + decay = torch.exp(-weights["A_log"].exp()[None, :, None, None] * delta) + ssm_state = ssm_state * decay + delta * b[:, t, :, :, None] * x[:, t, :, None, :] + readout = (ssm_state * c[:, t, :, :, None]).sum(-2) + values.append(readout + weights["D"][None, :, None] * x[:, t]) + y = torch.stack(values, 1).reshape(2, -1, d_inner) * torch_f.silu(gate) + grouped = y.reshape(2, -1, heads, d_head) + y = grouped * torch.rsqrt(grouped.square().mean(-1, keepdim=True) + 1e-5) + y = y.reshape(2, -1, d_inner) * weights["norm.weight"] + return y @ weights["out_proj.weight"].T, history[:, :, -kernel + 1 :], ssm_state, dt + + actual_conv = np.zeros((2, conv_dim, kernel - 1), np.float32) + actual_state = np.zeros((2, heads, d_state, d_head), np.float32) + ref_conv, ref_state = torch.from_numpy(actual_conv.copy()), torch.from_numpy(actual_state) + contrast_conv, contrast_state = ref_conv.clone(), ref_state.clone() + for step, length in enumerate([3, 1, 1, 1]): + u = torch.from_numpy(rng.normal(size=(2, length, d_model)).astype(np.float32)) + expected, ref_conv, ref_state, dt = reference(u, ref_conv, ref_state, time_step_min) + contrast, contrast_conv, contrast_state, _ = reference( + u, contrast_conv, contrast_state, time_step_min if step == 0 else 0.0 + ) + actual, actual_conv, actual_state, actual_dt = session.run( + None, {"hidden": u.numpy(), "conv": actual_conv, "state": actual_state} + ) + for got, want in [ + (actual, expected), + (actual_conv, ref_conv), + (actual_state, ref_state), + (actual_dt, dt), + ]: + np.testing.assert_allclose(got, want.numpy(), atol=2e-6, rtol=2e-5) + if time_step_min: + assert torch.all(dt[:, :, 0] == time_step_min) + assert torch.all(dt[:, :, 1] > time_step_min) + if step and time_step_min: + assert not torch.allclose(ref_state, contrast_state, atol=1e-7, rtol=1e-5) + assert not torch.allclose(expected, contrast, atol=1e-7, rtol=1e-5) + else: + torch.testing.assert_close(expected, contrast, atol=0, rtol=0) + + class TestMambaBlock: """Tests for MambaBlock graph construction.""" From e3d01a7712f38deb962e7107752da7fb668a7054 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Sun, 11 Oct 2026 09:40:00 -0700 Subject: [PATCH 15/15] Record bounded Mamba source diagnostic evidence Distinguish locally reproduced native fallback timestep/readout precision from the unresolved full-model H100 gate. Document the new grouped/cache regression coverage without changing recurrence or acceptance thresholds. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> --- docs/nemotron-lightning-export.md | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/docs/nemotron-lightning-export.md b/docs/nemotron-lightning-export.md index c6b1fbf98..8891d68c6 100644 --- a/docs/nemotron-lightning-export.md +++ b/docs/nemotron-lightning-export.md @@ -85,6 +85,19 @@ evidence, not the runnable `run_validation.py`, reference implementation, input/cache fixtures, provider profile, or optimized graph; those remote `/home/mobius/...` paths are not accessible in the local session. +Local diagnosis reproduces those **source distinctions**, not the full-model +failure. The pinned native fallback source bytes match the reported native +module SHA256 `5417d6f8dcbee945d4f1c4ae5327b0d1f443b29cc99098c99ec7befb9c9afcd8`. +Invoking its original CPU fallback bodies with dispatch decorators removed gives +synthetic prefill output/state NRMSE 7.13e-8 / 4.61e-8. Native unclamped cached +updates differ from bounded updates (NRMSE up to 0.00714 in this fixture, below +the full-model gate). Native BF16 bias/softplus and state-cast-before-readout +boundaries are independently reproduced, but their synthetic differences do +not establish the remaining H100 cause or authorize changing export semantics. +Regression tests cover two-row prefill plus three reused-state calls, grouped +B/C, active/inactive timestep minimum, output/conv/FP32 SSM and cast boundaries. +Acceptance thresholds and the production recurrence remain unchanged. + The original preflight host was Windows ARM64 with a Qualcomm Adreno X1-85, no NVIDIA/CUDA GPU or `nvidia-smi`, Python 3.12.10, 68,171,038,720 bytes RAM and approximately 818 GB free disk. Mobius, PyTorch, Transformers, ORT and