Skip to content

feat(nemotron_parse): port Nemotron-Parse on the seq2seq worker - #2034

Merged
inureyes merged 2 commits into
mainfrom
feature/issue-1369-nemotron-parse
Sep 29, 2026
Merged

inureyes merged 2 commits into
mainfrom
feature/issue-1369-nemotron-parse

Conversation

@inureyes

Copy link
Copy Markdown
Member

Summary

Ports Nemotron-Parse (nemotron_parse): C-RADIOv2-H ViT-H/16 tower plus compression neck (3329 encoder states for a 2048x1664 page) and a 10-layer pre-norm mBART decoder seeded with the task prompt. It reuses the Florence-2 attention sublayers, dual KV cache, and causal mask, and is wired into detection, loaders, mlxcel generate (early exit like Florence-2), mlxcel run refusal, the mlxcel arch catalog, a dedicated single-stream server worker (admission, warmup skip, built-in text-only template), and docs/supported-models.md.

Behavior worth reviewing

  • Tokenizer loading now drops padding / truncation serialized in tokenizer.json, matching transformers. This checkpoint serializes padding: Fixed(9000), which turned every prompt into 9000 ids. This touches every tokenizer.json model; the tokenizer, loading, and chat-template test modules pass.
  • The decoder always runs activations in f32. With bf16 activations the 8-bit export flips two near-ties (reference top-2 margins 0.04 to 0.07); f32 reproduces the reference.
  • For this family the built-in template replaces the Hub checkpoint's chat_template.jinja, which prints message['content'] raw and would render the typed image list as the decoder seed.
  • Server bounds on the task prompt: 1024 bytes, no control characters, at most 256 seed tokens; exactly one image; audio/video refused.

Deviations from the issue text

  • Seeding: the reference tokenizes with add_special_tokens=False and generate prepends decoder_start when missing, so that is implemented instead of the encode-then-strip rule (same result for the default prompt).
  • Repetition penalty covers the whole decoder sequence including the seed, as transformers does, not only generated ids.
  • Line numbers in the issue have drifted (for example the Florence-2 dispatch is at generate.rs near 2968, not 2431); the code was followed.
  • cargo test --workspace / clippy --workspace --all-targets were replaced by the narrow gate below (run rules for this machine).

Test plan

  • Real checkpoints (tests/nemotron_parse_real_model.rs, ignored-gated): Hub f32 and 8-bit match a CPU fp32 run of the checkpoint's own transformers code for all 79 model-chosen greedy tokens, at penalty 1.0 and 1.1 (8-bit reference ran on its dequantized weights; the reference's 80th token is a forced EOS). 4-bit reads the page.
  • mlxcel-server with the Hub checkpoint: /v1/chat/completions returns the reference text; empty prompt uses the default; two images, zero images, and control characters are refused; streaming works.
  • Unit tests: models::nemotron_parse (22), server::nemotron_parse_worker (7), detection, chat template, tokenizer padding; plus tokenizer::tests, loading::tests, server::chat_template, models::registry, model_metadata, execution::memory_estimate.
  • cargo fmt --all -- --check, cargo clippy --release -p mlxcel --lib --tests -- -D warnings, cli_help_consistency, dead_doc_pointers, llama_compat_manifest.

Pages larger than the target box are resized with image bilinear filtering, which is close to but not bit-identical with PIL; the parity page avoids resampling. The v1.x untied-head path is implemented but not run against a real v1.x checkpoint.

Closes #1369

@inureyes inureyes added status:review Under review type:enhancement New features, capabilities, or significant additions priority:medium Medium priority area:models Model architectures, weights, loading, metadata area:inference Generation, sampling, decoding (incl. speculative, DRY) area:cli Command-line interface / CLI flags area:docs User and developer documentation labels Sep 29, 2026
inureyes added a commit that referenced this pull request Sep 29, 2026
@inureyes inureyes added status:done Completed and removed status:review Under review labels Sep 29, 2026
Nemotron-Parse (model_type nemotron_parse) is NVIDIA's document OCR / layout model: a C-RADIOv2-H ViT-H/16 tower over a white-padded 2048x1664 page, a neck that compresses the patch grid 4x horizontally and appends one summary row (3329 encoder states), and a 10-layer pre-norm mBART decoder seeded with the tokenized task prompt. It is encoder-decoder, so it reuses the Florence-2 attention sublayers, dual KV cache, and causal-mask helper, and runs on its own CLI driver and a dedicated single-stream server worker.

- models/nemotron_parse: config parsing, Hub and MLX key canonicalization (with shape-gated conv rewrites), the C-RADIO tower (fused SDPA over 13320 tokens, CPE positional crop/resample), the neck, the pre-norm decoder, the greedy loop with transformers-style seeding and repetition penalty, the page processor, and the runtime unit.
- Detection, registries, loader, LoadedModel, CLI generate early exit, run/chat refusal, server worker selection, single-stream admission, warmup skip, and a built-in text-only chat template that replaces the Hub checkpoint's raw-content template.
- Tokenizer loading now drops the padding and truncation a tokenizer.json serializes, as transformers does; this checkpoint serializes padding Fixed(9000), which turned every prompt into 9000 ids.
- The decoder runs its activations in f32: bf16 activations flip two near-ties (reference margins 0.04 to 0.07) on the 8-bit export.

Validated token for token (79 greedy tokens, penalty 1.0 and 1.1) against a CPU fp32 run of the checkpoint's own transformers code on the Hub checkpoint and the 8-bit export (reference on its dequantized weights); the 4-bit export reads the page. Pinned in tests/nemotron_parse_real_model.rs.

Closes #1369
@inureyes
inureyes force-pushed the feature/issue-1369-nemotron-parse branch from f40c465 to fdc90c0 Compare September 29, 2026 09:28
@inureyes
inureyes merged commit 830f6b6 into main Sep 29, 2026
@inureyes
inureyes deleted the feature/issue-1369-nemotron-parse branch September 29, 2026 09:28
inureyes added a commit that referenced this pull request Sep 30, 2026
Three gate failures landed on main with #2034 and #2037. There is no issue for them; each is explained below.

## Changes

- **`make verify-binary-assets`**: `tests/fixtures/nemotron_parse_page.png` (#2034) was 41.1 KB, over the 32 KB ceiling on `tests/fixtures/*.png`. The page is black text on white, so every pixel is gray. Re-encoding it losslessly as 8-bit grayscale with Pillow `optimize=True` brings it to 20.6 KB (an RGB `optimize=True` re-encode only reached 39.6 KB, and no oxipng, optipng, pngcrush or zopflipng is installed here). Decoding the new file to RGB gives byte-identical pixels to the old one, checked with Pillow `convert("RGB")` and with the `image` 0.25.10 crate's `to_rgb8()`, which is what `NemotronParseImageProcessor::preprocess_to_vec` calls. No per-file rule was needed. `generate_nemotron_parse_page.py` now saves the same encoding.
- **`tests::family_order_is_exhaustive`**: #2037 added the `Speech` family (Nemotron VoiceChat, speech-to-speech) without listing it in `FAMILY_ORDER`. It now sits after `Text-to-speech`, the same order as the ASR, TTS and speech-to-speech sections of `docs/supported-models.md`, so no docs change.
- **`gemma3_backbone_tests::gelu_approx_matches_mlx_nn_bit_for_bit`**: the test hardcoded values measured on Metal, and ROCm's f32 result at x = 4.1 is 4.0999565 against the hardcoded 4.099957 (one ulp). The test now builds the reference from the pinned `python/mlx/nn/layers/activations.py` line (Python operator order, `mx.power` for `x**3`, weak scalars built as `array(float(v), x.dtype)` as the bindings do), runs it on the same device, and compares bit for bit with `gelu_approx` over 4097 points in [-8, 8] for f32, bf16 and f16. Backend-independent f64 sanity values with a tolerance keep it from passing on a formula that is wrong on both sides.

The reference is a transcription of the upstream expression, so it has the same ops as the port by design. What the test checks is that the port stays op-for-op equal to upstream; the f64 sanity values check the formula itself. Temporarily breaking `gelu_approx` fails the test in each case: `x * x * x` for `power` (2 of 4097 f32 elements differ), `0.0447` for `0.044715` (2011 f32 elements), computing in f32 and rounding once (1270 bf16 elements), and `0.0447` changed in both the port and the reference (caught by the f32 sanity check).

## Verification (gfx1151, `--features rocm`)

- `cargo test -p mlxcel --lib models::gemma3_backbone::gemma3_backbone_tests`: 5 passed.
- `cargo test -p mlxcel --bin mlxcel tests::family_order`: 2 passed.
- `cargo test --test nemotron_parse_real_model -- --ignored`: builds and runs; all three skip, since no Nemotron-Parse checkpoint is on this host. `cargo test --test cli_help_consistency`: 27 passed.
- `make verify-versions verify-kernel-dtype-keys verify-kernel-port-dispatch verify-llama-compat verify-rocm-overlay verify-fmt verify-binary-assets`: pass.
- `cargo clippy -p mlxcel --lib --bin mlxcel --tests -- -D warnings`: clean.

Not verified: Metal and CUDA (not available on this host). The gelu test should still pass on Metal: the port and the reference issue the same MLX ops in the same order on the same stream, so each backend's kernels round them identically, and the sanity tolerances (1e-6 in f32, 4e-3 plus 1% in bf16) cover the values #2037 pinned on Metal. The same argument holds for CUDA. The other two changes are not backend-specific.
inureyes pushed a commit that referenced this pull request Oct 21, 2026
…2525)

The scope page listed the supported non-chat vision families but not
nemotron_parse, the document OCR and layout model that already loads through
detection.rs and runs on the seq2seq worker (#1369, #2034).
Add its row beside florence2 and refresh the commit the page was checked
against.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:cli Command-line interface / CLI flags area:docs User and developer documentation area:inference Generation, sampling, decoding (incl. speculative, DRY) area:models Model architectures, weights, loading, metadata priority:medium Medium priority status:done Completed type:enhancement New features, capabilities, or significant additions

Projects

None yet

Development

Successfully merging this pull request may close these issues.

feat(nemotron_parse): port Nemotron-Parse (C-RADIO ViT-H encoder + pre-norm mBART decoder) on the seq2seq worker

1 participant