Add NVIDIA Nemotron 3.5 Lightning export support - #487
Open
justinchuby wants to merge 12 commits into
Open
Conversation
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
Performance Comparison
|
![github-code-quality[bot]](https://avatars.githubusercontent.com/u/9919?s=60&v=4){kind=small}
6 tasks
Contributor
There was a problem hiding this comment.
Pull request overview
Adds export + validation support for the official NVIDIA Nemotron 3.5 Lightning (nemotron_h) BF16 checkpoint by pinning the upstream revision, normalizing config/weight handling for the current Transformers vocabulary, and introducing reduced-real-weight L4/L5 evidence plus runtime/ORT GenAI guardrails.
Changes:
- Normalize NemotronH config + cache dtype behavior (layer-type vocabularies, SSM cache dtype, pinned revisions/dtypes in validation paths).
- Update NemotronH weight alignment and MoE routing numerics (preserve fp32 router bias / recurrent state precision; ignore only
mtp.*auxiliary heads). - Add reduced-real-weight fixtures + tests, plus an Olive/ORT-direct example pipeline; reject unsupported ORT GenAI mixed-cache export earlier (CLI + config writer).
Reviewed changes
Copilot reviewed 34 out of 34 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
| tests/weight_alignment_test.py | Adds NemotronH-specific assertions for mtp.* filtering and per-expert weight name mapping. |
| tests/synthetic_parity_test.py | Updates NemotronH HF layer-type vocabulary mapping and removes the prior NemotronH xfail. |
| tests/nemotron_h_real_weight_test.py | New reduced-real-weight L4/L5 integration tests comparing ONNX vs HF plus golden checks. |
| tests/model_coverage_test.py | Removes NemotronH “no test model id” skip now that a checkpoint is pinned. |
| tests/cli_test.py | Ensures --runtime ort-genai rejects NemotronH before build/weight download. |
| tests/build_graph_test.py | Adds tests for fp32 preservation (router bias + SSM cache) and BF16 rejection. |
| tests/arch_validation_test.py | Forwards YAML-pinned revisions into config load and applies YAML dtype for NemotronH. |
| tests/_test_configs.py | Updates NemotronH test configs (adds explicit rms_norm_eps). |
| testdata/golden/causal-lm/nemotron-3_5-lightning-30b-reduced.json | New reduced-real L4 golden reference (top-k IDs/logits + summary). |
| testdata/golden/causal-lm/nemotron-3_5-lightning-30b-reduced_generation.json | New reduced-real L5 greedy generation golden reference. |
| testdata/cases/causal-lm/nemotron-3_5-lightning-30b.yaml | New pinned YAML case documenting revision/dtype and CI skip rationale. |
| src/mobius/tasks/_cache_utils.py | Allows NemotronH SSM cache inputs to use explicit cache dtype (fp32 by default). |
| src/mobius/models/nemotron_h.py | Implements mtp.* filtering, bf16 rejection, and fp32-sensitive MoE routing details. |
| src/mobius/integrations/ort_genai/auto_export.py | Adds early model-type rejection for known unsupported ORT GenAI contracts (NemotronH). |
| src/mobius/integrations/ort_genai/auto_export_test.py | Verifies NemotronH is rejected before writing ORT GenAI artifacts. |
| src/mobius/_registry.py | Updates NemotronH test model id to the official 3.5 Lightning BF16 checkpoint. |
| src/mobius/_configs/_base.py | Adds NemotronH layer-type normalization, stricter validation, and SSM cache dtype config. |
| src/mobius/_configs_test.py | Adds coverage for NemotronH layer-type vocabulary normalization and unknown-type rejection. |
| src/mobius/main.py | Adds CLI preflight rejection for ORT GenAI runtime before build/weight download. |
| examples/olive/nemotron-3_5-lightning-30b/validate_reduced_checkpoint.py | New script to build/validate reduced-real-weight packages (incl. range-fetch). |
| examples/olive/nemotron-3_5-lightning-30b/requirements.txt | Example-only deps for Olive workflow + range-reading. |
| examples/olive/nemotron-3_5-lightning-30b/README.md | Documents Option A flow, constraints, reduced-real evidence, and usage. |
| examples/olive/nemotron-3_5-lightning-30b/optimize.py | Exports pinned FP16 ONNX and runs CPU-isolated Olive INT4 quantization + smoke test. |
| examples/olive/nemotron-3_5-lightning-30b/olive_q4.json | Example Olive workflow JSON for Q4 K-quant. |
| examples/olive/nemotron-3_5-lightning-30b/inference.py | Direct ORT token-by-token generation loop for hybrid caches. |
| examples/olive/nemotron-3_5-lightning-30b/.gitignore | Ignores example cache/output directories. |
| .github/workflows/gpu_l5_generation_e2e.yml | Includes new NemotronH reduced-real-weight generation test in GPU L5 job. |
| .github/workflows/gpu_l4_golden_parity.yml | Includes new NemotronH reduced-real-weight golden test in GPU L4 job. |
| .agents/skills/writing-tests/SKILL.md | Codifies reduced-real-weight fixture guidance and warns about stale layer-type mappings. |
| .agents/skills/ort-genai-config/SKILL.md | Documents NemotronH mixed-cache ORT GenAI incompatibility and required early rejection. |
| .agents/skills/onnx-export-quantization/SKILL.md | Adds guidance on isolating Olive from unrelated EP DLL registration. |
| .agents/skills/moe-models/SKILL.md | Adds NemotronH config vocabulary + reduced-precision routing and mtp.* filtering guidance. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
🏗️ Architecture Diff
No architecture changes detected. ✅ Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
justinchuby
commented
Aug 14, 2026
justinchuby
commented
Aug 14, 2026
justinchuby
commented
Aug 14, 2026
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 34 out of 34 changed files in this pull request and generated no new comments.
Suppressed comments (3)
examples/olive/nemotron-3_5-lightning-30b/inference.py:207
OnnxKQuantQuantizationcan rename the graph output tologits_Q4(the repository already handles this inexamples/qwen3_6_35b_a3b_ort_genai.py:574-581andexamples/gemma4_unified_ort_genai.py:222-239). Becausequantize_package()passes that Olive output directly to this function, indexing only"logits"can raiseValueError, preventing the advertised post-quantization smoke test and reduced Q4 generation from running. Accept both output names or normalize the graph output after quantization.
logits = _as_numpy(outputs[output_names.index("logits")])[0, -1].astype(np.float32)
src/mobius/main.py:345
- This direct
AutoConfigcall bypasses the supported raw-config.jsonfallback inbuild()(src/mobius/_builder.py:542-552). Registered models whose config type is not recognized by the installed Transformers version will now fail only when--runtime ort-genaiis requested, before either the compatibility validator or the normal builder runs. Resolve the model type through the same fallback path (ideally a shared helper) so the new preflight does not regress those models.
runtime_config = transformers.AutoConfig.from_pretrained(
model_id_or_path,
trust_remote_code=trust_remote_code,
)
src/mobius/integrations/ort_genai/auto_export.py:1187
- This actionable error points Nemotron 3.5 users to the older Nano-specific script, whose default checkpoint and compatibility patches differ from the package introduced by this PR. Point to the new direct-runtime example so users follow the pinned 3.5 export and inference path described here.
"Runtime (see examples/nemotron_3_nano_text_generation.py)."
justinchuby
force-pushed
the
justinchuby-add-nemotron-35-lightning
branch
3 times, most recently
from
August 15, 2026 02:29
eee55c8 to
fbffb25
Compare
This was referenced Aug 15, 2026
Normalize current and legacy NemotronH configs, preserve routing and recurrent-state precision, and filter only auxiliary MTP weights. Add pinned reduced-real L4/L5 evidence, CUDA/ORT GenAI guards, and an executable Olive Q4 direct-runtime recipe. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Apply the Linux formatter/import order to the Olive scripts and include the pending ModelPackage formatting fixes required by the repository-wide lint job. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Separate unsupported BF16 evidence from the supported validation matrix, add automated FP16 CUDA cached-logit parity, and verify fresh Olive Q4 quantization through final-package CUDA reload and generation. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Fail reduced-real validation whenever the maximum absolute prefill or cached-step error exceeds the advertised 1e-2 gate, independent of relative tolerance. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Cache test-case discovery, use generic reduced-real GPU workflow discovery, and emit an honest NemotronH GenAI config without a runtime-version-specific hard rejection. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Teach affected-model detection about reduced-real tests and example assets, and stop direct generation on configured EOS without an unused final decoder invocation. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Profile only the cached-generation session used for placement evidence so prefill session teardown does not leave untracked ORT profile files. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Reject structurally unrepresentable SSM cache graphs before artifact creation and persist the pinned reduced-real fixture with validated range retries, atomic schema-versioned writes, and shared L4/L5 CI caching. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Set the runner-temp fixture path on the L4/L5 execution steps, where the runner context is available, while retaining the shared cache action path and key. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Remove downstream ORT GenAI capability gates and derive decoder cache/output metadata directly from exported graphs. Use portable standard-ONNX cache operators for validated Nemotron CUDA and Olive generation, including quantized logits handling. Preserve GGUF special-token metadata and update model-development guidance to keep runtime acceptance downstream. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Drop serialized null-token sentinels and out-of-vocabulary special-token IDs while preserving valid BOS and padding metadata. Combine embedded EOS, EOT, and EOM IDs into the generation stop-token contract. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
justinchuby
force-pushed
the
justinchuby-add-nemotron-35-lightning
branch
from
August 15, 2026 05:38
cfbbc7a to
ae4fd70
Compare
Remove the seven cross-cutting skill updates from the Nemotron implementation branch now that they are published independently in PR #500. Retain Nemotron-specific MoE guidance and all implementation, test, example, and workflow changes. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
justinchuby
added a commit
that referenced
this pull request
Aug 15, 2026
## Summary - extract the model-conversion lessons from #487 into a standalone skills-only change - require faithful graph-derived metadata without gating Mobius artifacts on the current ORT GenAI registry, version, topology, or cache executor - clarify cached-decode diagnosis, standard-ONNX controls before blaming Olive, and provider isolation for weight-only quantization - preserve the quantized-by-default and Direct GGUF guidance merged in #496 ## Scope Only the seven intended `SKILL.md` files are changed. Model code, tests, workflows, examples, and `.agents/skills/moe-models/SKILL.md` are intentionally excluded. ## Validation - `python -m pytest tests/yaml_schema_test.py -q --tb=short` (246 passed) - all 19 skill frontmatter blocks parsed and matched their directory names - `lintrunner -a --output oneline` - `git diff --check` - independent code review completed; its optional-runtime waiver contradiction was resolved --------- Signed-off-by: Justin Chu <justinchu@microsoft.com> Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com> Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16atd468880b6ad3c6e0d21377ce7242adaea4cc884dnemotron_harchitecture (no alias), normalize legacy/current layer vocabularies, and preserve FP32 router correction bias, routed-expert accumulation, and Mamba recurrent statemtp.*tensors and prove complete base-decoder alignment--dequantizeopt-out without duplicating its commits or guidancemoe-modelsguidance hereFinal evidence
1299f396cd69ff662548b770d5a3e7dbc32f9b7e; heada8cd77570bca980861eeadab9e6fa077464f0138origin/mainorigin/main0.00781250.00292969,0.0078125,0.00390625,0.00390625[12, 13, 12, 12]Attention,Scan, andConv; noGroupQueryAttention,LinearAttention, orCausalConvWithStateMatMulNBitsnodes, smaller than source, CUDA reload, finite logits, exact generation1e-2numerical gate (~0.855469max abs); supported export converts source BF16 weights to FP16lintrunner -a -r origin/mainandgit diff --check: cleanorigin/main;moe-modelsretains its 19-line Nemotron guidanceThe four inherited canonical failures are the GGUF Q4_0 zero-point test, two Attention softcap decomposition tests, and Qwen Image FP16 CUDA low-precision parity. Each was reproduced on detached
origin/mainat1299f39; this PR does not alter those implementations or hide their failures.ORT GenAI metadata policy
Mobius emits the best faithful package from graph metadata: semantic inputs/outputs, representable key/value and convolution cache templates, actual
logits/logits_Q4names, and the graph-derived global cache-slot count. It does not reject export based on the current ORT GenAI registry, version, topology support, or cache executor capability. NemotronHssm_statehas no current config template, so that unrepresentable field is omitted without blocking artifact creation; downstream load/generation results are informational.The same policy applies to GGUF, NeMo, Parakeet, Mage-VL, Moonshine, generic multimodal packages, and mixed LinearAttention/standard-Attention graphs. Intrinsic Mobius input/config errors remain errors. GGUF config extraction validates special-token IDs against vocabulary bounds, drops serialized null-token sentinels, and emits a deduplicated EOS/EOT/EOM stop-token union.
Resource/runtime waivers
Skill retrospective
Especially effective: architecture detection prevented a false alias; MoE/alignment guidance exposed non-fused sigmoid routing, FP32 correction bias, and auxiliary MTP keys; testing guidance drove independent reduced-real parity.
Nemotron-specific lessons remain in
.agents/skills/moe-models/SKILL.md. The seven broader lessons—runtime/export ownership, optional downstream probes, cached-decode validation, provider-isolated Olive packaging, and graph-derived metadata—are published separately in #500 so this implementation PR stays model-focused. #496's Direct GGUF/default-quantized guidance remains inherited once frommain.Relates to #483.
Cross-cutting skill guidance: #500.