Add NVIDIA Nemotron Parse 2.0 support - #473
Open
justinchuby wants to merge 3 commits into
Open
Conversation
Implement the C-RADIO vision encoder, compressed visual neck, and cross-attentive mBART decoder as a standardized vision encoder-decoder package. Add config and registry integration, real-image CUDA goldens, synthetic and real-weight parity, processor/generation support, and explicit ORT GenAI runtime guarding. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
Performance Comparison
|
🏗️ Architecture Diff
No architecture changes detected. ✅ Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
Contributor
There was a problem hiding this comment.
Pull request overview
Adds first-class support for NVIDIA Nemotron Parse 2.0 as a two-model vision-encoder + cross-attentive decoder package, including C‑RADIOv2‑H vision backbone, weight preprocessing, and extensive synthetic/real validation plus golden generation support.
Changes:
- Introduce
nemotron_parsemodel + config extraction (NemotronParseConfig) and register it in the model registry. - Add a new
vision-encoder-decodertask that exportsvision_encoderanddecodercomponents (self-attention KV cache only). - Expand test + golden infrastructure to cover image-to-text prefill/generation and ORT GenAI export guarding.
Reviewed changes
Copilot reviewed 25 out of 26 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/weight_alignment_test.py | Adds Nemotron Parse weight-alignment identity roundtrip coverage. |
| tests/synthetic_parity_test.py | Adds L3 synthetic parity for tiny C‑RADIO + mBART-style decoder path. |
| tests/integration_test.py | Adds CUDA BF16 real-weight parity test against Hugging Face Nemotron Parse. |
| tests/e2e_golden_test.py | Adds “image-to-text” golden prefill + greedy generation execution path. |
| tests/build_graph_test.py | Updates build assertions to accept 2-model vision-encoder-decoder packages. |
| tests/_test_configs.py | Adds tiny nemotron_parse config to VL test matrix. |
| testdata/golden/vision-language/nemotron-parse-2.json | Adds L4 golden logits reference for Nemotron Parse. |
| testdata/golden/vision-language/nemotron-parse-2_generation.json | Adds L5 golden generation tokens/text for Nemotron Parse. |
| testdata/cases/vision-language/nemotron-parse-2.yaml | Adds real-image Nemotron Parse test case metadata and params. |
| src/mobius/tasks/_vision_encoder_decoder.py | Implements a split vision-encoder/decoder task for image-to-text. |
| src/mobius/tasks/init.py | Exposes and registers the new task name vision-encoder-decoder. |
| src/mobius/models/nemotron_parse.py | Adds Nemotron Parse model implementation + weight preprocessing. |
| src/mobius/models/init.py | Exports the Nemotron Parse model class. |
| src/mobius/integrations/ort_genai/auto_export.py | Adds explicit guard rejecting unsupported generic vision encoder-decoder packages. |
| src/mobius/integrations/ort_genai/auto_export_test.py | Adds unit test verifying the ORT GenAI guard prevents artifact emission. |
| src/mobius/components/_radio_vision.py | Adds C‑RADIO ViT components (patch generator, blocks, summaries). |
| src/mobius/components/_conv.py | Extends Conv2d/Conv2dNoBias to support asymmetric (h, w) kernel/stride/padding. |
| src/mobius/components/_conv_test.py | Adds tests for asymmetric Conv2d/Conv2dNoBias parameterization. |
| src/mobius/components/init.py | Exports RadioVisionModel from components public API. |
| src/mobius/_registry.py | Registers nemotron_parse model_type and default model id. |
| src/mobius/_configs/_base.py | Adds NemotronParseConfig extraction from HF config (decoder aliases, image sizing). |
| src/mobius/_configs/_base_test.py | Adds config extraction regression test for MBART decoder head aliases. |
| src/mobius/_configs/init.py | Exports NemotronParseConfig. |
| scripts/generate_golden.py | Adds golden-generation routine for image-to-text (Nemotron Parse-style). |
| README.md | Lists Nemotron Parse under supported multimodal models. |
Suppressed comments (2)
tests/synthetic_parity_test.py:989
- Same as above: feed
onnx_encoderinto the ONNX decoder so the unpadded/padded invariance checks validate the realvision_encoderoutputs, not the torch reference tensor.
"encoder_hidden_states": torch_encoder.numpy(),
tests/synthetic_parity_test.py:998
- Same as above:
padded_logitsshould useonnx_encoderto ensure the padding-mask invariance test covers the full ONNX pipeline.
"encoder_hidden_states": torch_encoder.numpy(),
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Expose the official processor's float32 pixel tensors at the vision graph boundary and cast internally for reduced-precision models. Update graph and GPU parity tests to exercise the deployment contract directly. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
justinchuby
changed the title
Feed the ONNX vision encoder output into the ONNX decoder so L3 covers the exported two-model data path while preserving independent stage parity checks. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.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
Add end-to-end support for
nvidia/NVIDIA-Nemotron-Parse-2.0(model_type: nemotron_parse) as a two-model image-to-text package.Architecture
vision_encoderanddecoder.Validation
GPU: NVIDIA RTX A1000, 8188 MiB, driver 573.44. CUDA runs explicitly requested
CUDAExecutionProvider; sessions reported['CUDAExecutionProvider', 'CPUExecutionProvider'](CPU is retained for ORT shape operations).4 passed, 1350 deselected in 8.74s1 passed, 226 deselected in 8.80s237 passed in 1.94s1 passed, 203 deselected in 7.43satrtol=atol=1e-3(ONNX vision output fed to ONNX decoder)1 passed, 122 deselected in 71.49s(74.683s wall); encoder cosine >0.99, logits cosine >0.995, identical argmax2 passed, 377 deselected in 103.13s(106.441s wall); real nonzero image, exact L4 top-1 and exact 24-token L5 golden3638 passed, 57 skipped in 47.45slintrunner;lintrunner -a:ok No lint issuesCUDA CLI and dtype matrix
All three exports used
mobius build --model nvidia/NVIDIA-Nemotron-Parse-2.0 --trust-remote-code --dtype <dtype> --ep cudaand producedvision_encoder/model.onnxplusdecoder/model.onnx.50251matched.50251matched.50251matched.Olive
Olive 0.13.0
OnnxBnb4QuantizationNF4 quantized the final f16 decoder successfully in 29.320s (237,677,671-byte model). CUDA generation produced 24 non-degenerate OCR tokens in 7.112s:# MOBIUS OCR VALIDATION ... Nemotron Parse 2.Checklist
lintrunner -acleanNarrow waivers / limitations
genai_config.jsonwould create an unloadable artifact, so export explicitly raisesNotImplementedError; the guard test passed (1 passed, 76 deselected in 7.08s). Direct ONNX CUDA runtime is fully validated above.EncoderDecoderAttentioncontract currently recomputes encoder K/V projections per decode token. Adding a lazy cross-cache requires a broader split-prefill/decode redesign; this is a performance limitation, not a correctness gap, and the exact 24-token CUDA L5 run passes.Baseline note
Before validation,
.pytest_cache/v/cache/lastfailedcontained onlytests/synthetic_parity_test.py::test_synthetic_parity[granitemoehybrid]. That cached unrelated baseline was not treated as a Nemotron failure; targeted Nemotron parity and the full non-integration regression suite pass as reported above.Post-push CI baseline status
The repository-wide GitHub matrix and L3 jobs currently fail only the untouched
bamba_0/bamba_1synthetic cases on every Python/OS matrix entry (max_abs_diff=0.001375/0.001662, argmax matches, cosine ≥0.999996). Nemotron's affected-model L1 job passes, and its targeted local L1-L5 results are above. The benchmark-base job separately failed while finalizing an artifact with GitHub403 Forbidden; benchmark head passed. These unrelated failures are not masked or changed in this PR.