feat(nemotron_parse): port Nemotron-Parse (C-RADIO ViT-H encoder + pre-norm mBART decoder) on the seq2seq worker
- Dominant language
- Rust
- Stars
- 467
- Forks
- 54
- Avg merge
- 4h 25m
- Merged PRs (30d)
- 310
Description
## Summary
Nemotron-Parse (`nvidia/NVIDIA-Nemotron-Parse-2.0`, `model_type: "nemotron_parse"`, `architectures: ["NemotronParseForConditionalGeneration"]`, also v1.2 with an untied head) is a document OCR / layout model: a C-RADIO ViT-Huge encoder (32 blocks, 1280 wide, patch 16, 8 prefix tokens, learned 128x128 positional grid) over a white-padded 2048x1664 page, a compression neck that shrinks the token count 4x horizontally and appends one summary token, and a 10-layer pre-norm mBART decoder (1024 wide, 16 heads, cross-attention, no positional embeddings) that emits markdown with bounding boxes and class tags. It is encoder-decoder, so it cannot use the decoder-only generation loop; mlxcel already runs one such family (Florence-2) on a dedicated seq2seq worker with a greedy cross-attention decode, and this port reuses that attention, cache, CLI and server plumbing. mlxcel has no `nemotron_parse` arm today.
## Current behavior
- `src/models/detection.rs:376` has `"florence2" => Ok(ModelType::Florence2VLM)` and `:378` `"whisper"`; no `nemotron_parse` arm.
- Florence-2 seq2seq pieces that apply directly: `Florence2Attention::{self_attention, cross_attention}` with the one-shot cross K/V cache (`src/models/florence2/layers.rs:189-255`), `KvCache` / `Florence2LayerCache` (`layers.rs:37-60`), `additive_causal_mask` (`layers.rs:123`), `layer_norm` / `linear` loaders (`layers.rs:106`), the greedy loop `Florence2Model::generate_greedy_with_cancel_and_prompt_len` (`src/models/florence2/model.rs:542-584`), the CLI entry `src/commands/generate_florence2.rs` (dispatched at `src/commands/generate.rs:2431`), and the server worker `run_florence2_worker_loop` (`src/server/florence2_worker.rs:243`, selected at `src/server/model_worker.rs:346` and `:1014`, admission at `src/server/model_provider.rs:130-138`, warmup skip at `src/server/startup.rs:1695`).
- Two things do not carry over: the Florence-2 decoder is post-norm with learned `embed_positions` (`src/models/florence2/decoder.rs:35-128`), whereas this decoder is pre-norm, has a final `layer_norm`, scales embeddings by `sqrt(d_model)`, and has no positional table; and the Florence-2 runtime seeds the decoder with only `decoder_start_token_id` (`model.rs:566`), whereas this model must be seeded with the tokenized task prompt.
- No C-RADIO encoder exists; the nearest ViT pieces are `src/vision/encoders/deepseekocr_sam.rs` (pre-norm blocks with `qkv`/`proj`, `mlp.fc1/fc2`) and `src/vision/encoders/florence2_davit.rs`.
## Expected behavior
```
./target/release/mlxcel generate -m models/Nemotron-Parse-2.0-4bit --image page.png \
-p "" -n 4096
```
prints the page as markdown annotated with `` coordinate tokens and `` tags, stopping at `` (2). Through the server, `/v1/chat/completions` with one `image_url` part and the task prompt as the text part returns the same text (single-stream, batch 1, like Florence-2). Omitting the task prompt is allowed but produces a repetition loop, so the CLI prints a warning and substitutes the default prompt above when `-p` is empty.
### config.json
Top level: `model_type "nemotron_parse"`, `is_encoder_decoder true`, `image_size [2048, 1664]` (height, width), `vocab_size 72256` (v2.0), `max_sequence_length 9000`, `pad_token_id 1`, `bos_token_id 0`, `eos_token_id 2`, `decoder_start_token_id 2`, `tie_word_embeddings true`, `class_token_start_idx 52315`.
`decoder` sub-config (note the key is `decoder`, not `text_config`): `model_type "nemotron_parse_text"`, `d_model 1024`, `decoder_layers 10`, `decoder_attention_heads 16`, `decoder_ffn_dim 4096`, `activation_function "gelu"`, `scale_embedding true`, `add_final_layer_norm true`, `max_position_embeddings 9000` (no table exists; used only as a decode-length bound), `vocab_size 72256`, `decoder_start_token_id null` (use the top-level 2). The `encoder_*` keys in this block are unused.
`encoder` sub-config: a C-RADIO wrapper config with `patch_size 16`, `max_resolution 2048`, `args.model "vit_huge_patch16_224"`, `args.register_multiple 8`, `args.cls_token_per_teacher true`, `args.teachers` of length 4 (clip, siglip, dino_v2, sam). Derived constants (hardcode with config overrides): `hidden_size 1280`, `num_heads 16`, `mlp_ratio 4.0`, `num_layers 32`, `num_cls_tokens 4`, `num_register_tokens 4`, `summary_idxs [0, 1, 2]`, `neck_dim 1024`, `image_mean [0.48145466, 0.4578275, 0.40821073]`, `image_std [0.26862954, 0.26130258, 0.27577711]`.
`preprocessor_config.json`: `final_size [2048, 1664]`, CLIP mean/std, `rescale_factor 1/255`.
### Image processor
1. Convert to RGB. 2. Resize keeping aspect ratio so `height <= 2048` and `width <= 1664` (first clamp height, recompute width; then clamp width, recompute height; `int()` truncation; bilinear). 3. Center-pad with white (255) to exactly 2048x1664 (`pad_top = pad_h // 2`, `pad_left = pad_w // 2`). 4. `pixel / 255`, CHW, normalize with CLIP mean/std. Output `[1, 3, 2048, 1664]`.
### Encoder (C-RADIO ViT-H + neck)
Let `H = 2048, W = 1664`, `hp = 128`, `wp = 104`, `L = hp * wp = 13312`.
1. Patchify: `x.reshape(B, 3, hp, 16, wp, 16).transpose(0, 2, 4, 1, 3, 5).reshape(B, L, 768)`; `x = x @ patch_embed.weight.T` (Linear 768 -> 1280, no bias).
2. Positional grid `pos_embed: [1, 16384, 1280]` viewed as `[128, 128, 1280]`; for `(hp, wp) == (128, 128)` use it whole, otherwise bilinear-interpolate (`align_corners = true`) the 128x128 grid to `(m, m)` with `m = max(hp, wp)` and crop `[:hp, :wp]`. For the native page `m = 128`, so the interpolation is the identity and the crop keeps the first 104 columns of every row. `x = x + pos[None]`.
3. Prepend the 8 prefix tokens `cls_token: [8, 1280]` (4 CLS + 4 registers) -> `[B, 8 + L, 1280]`.
4. 32 pre-norm blocks: `x = x + attn(LN1(x))`, `x = x + mlp(LN2(x))`; LayerNorm eps 1e-6; attention `qkv: Linear(1280, 3840)` with bias reshaped `[B, N, 3, 16, 80]`, `proj: Linear(1280, 1280)`; MLP `fc1 1280 -> 5120`, exact GELU, `fc2 5120 -> 1280`. Full attention over 13320 tokens: use `mlxcel_core::scaled_dot_product_attention` so the `[16, 13320, 13320]` score matrix is never materialized (17x peak-memory difference).
5. Split: `summary = concat(x[:, 0], x[:, 1], x[:, 2])` -> `[B, 3840]`; `features = x[:, 8:]` -> `[B, L, 1280]`.
6. Neck: `f = LN1(conv1(features))` where `conv1` is a 1x1 Conv1d 1280 -> 1024 with bias (a Linear); reshape `[B, hp, wp, 1024]`; `conv2`: Conv2d 1024 -> 1024, kernel `(1, 4)`, stride `(1, 4)`, no bias -> `[B, hp, wp/4, 1024]` -> reshape `[B, hp*wp/4, 1024]` = `[B, 3328, 1024]`; `f = LN2(f)`; `s = LN3(sum_proj(summary))` with `sum_proj: Linear(3840, 1024)`; `encoder_hidden = concat(f, s[:, None]) -> [B, 3329, 1024]`.
### Decoder (pre-norm mBART, no positions)
`embed_scale = sqrt(1024) = 32`. For decoder ids `y: [B, T]`: `h = LN(layernorm_embedding)(shared(y) * 32)`. Each of 10 layers:
```
r = h; h = LN(self_attn_layer_norm)(h); h = r + self_attn(h, causal, self-kv cache)
r = h; h = LN(encoder_attn_layer_norm)(h); h = r + encoder_attn(h, encoder_hidden, cross-kv cache)
r = h; h = LN(final_layer_norm)(h); h = r + fc2(gelu_exact(fc1(h)))
```
then `h = LN(layer_norm)(h)`; `logits = h @ lm_head.weight.T`. All attention projections (`q_proj, k_proj, v_proj, out_proj`) have bias; 16 heads of 64; scale `64^-0.5`. LayerNorm eps is the torch default 1e-5 for every decoder norm. `lm_head.weight` is the tied `shared.weight` on v2.0 (and on the converted v2.0 MLX checkpoints, which store it explicitly), and a real untied tensor on v1.x.
### Prompt seeding and decode
`prompt_ids = tokenize(prompt, add_special_tokens = true)`; when the result is longer than 2 tokens, strip a leading `` (0) and a trailing `` (2) that the tokenizer wrapper added, keeping the specials that are part of the prompt text. For the default prompt the seed is `[2, 0, 50004, 50008, 50001, 50010]`. If the prompt is empty the seed is `[decoder_start_token_id] = [2]`. Run the seed through the decoder in one call (prefill, causal mask, offset 0), take the argmax of the last position, then decode one token per step until `` (2), `max_new_tokens`, or 9000 total decoder positions. Greedy, with the `--repetition-penalty` value from `SamplingOptions` (`src/main.rs:630-632`) applied over the generated ids when it is not 1.0 (the model card's inference script defaults to 1.1); the Florence-2 loop is greedy-only, so the penalty is new to the seq2seq path and must be applied to the `[1, 1, vocab]` logits before the argmax.
### Weight keys
Hub layout (v2.0):
```
encoder.model_encoder.radio_model.model.patch_generator.embedder.weight -> patch_embed.weight [1280, 768]
encoder.model_encoder.radio_model.model.patch_generator.pos_embed -> pos_embed [1, 16384, 1280]
encoder.model_encoder.radio_model.model.patch_generator.cls_token.token -> cls_token [8, 1280]
encoder.model_encoder.radio_model.model.blocks.{0..31}.{norm1,norm2}.{weight,bias}
encoder.model_encoder.radio_model.model.blocks.N.attn.{qkv,proj}.{weight,bias}
encoder.model_encoder.radio_model.model.blocks.N.mlp.{fc1,fc2}.{weight,bias}
encoder.model_encoder.summary_idxs -> drop
encoder.model_encoder.radio_model.input_conditioner.* -> drop (normalization is in the processor)
encoder.conv1.{weight,bias} # [1024, 1280, 1] torch Conv1d -> store as [1024, 1280] Linear
encoder.conv2.weight # [1024, 1024, 1, 4] torch -> MLX [1024, 1, 4, 1024]
encoder.{layer_norm1,layer_norm2,layer_norm3}.{weight,bias}
encoder.sum_proj.{weight,bias}
decoder.embed_tokens.weight -> shared.weight (tied head)
decoder.layers.{0..9}.self_attn.{q,k,v,out}_proj.{weight,bias}
decoder.layers.N.encoder_attn.{q,k,v,out}_proj.{weight,bias}
decoder.layers.N.{self_attn_layer_norm,encoder_attn_layer_norm,final_layer_norm}.{weight,bias}
decoder.layers.N.{fc1,fc2}.{weight,bias}
decoder.{layernorm_embedding,layer_norm}.{weight,bias}
lm_head.weight # v1.x only (untied)
```
Converted MLX layout (`mlx-community/Nemotron-Parse-2.0-{4bit,8bit}`): `vision_tower.patch_embed.weight`, `vision_tower.pos_embed`, `vision_tower.cls_token`, `vision_tower.blocks.N.*`, `vision_tower.neck.{conv1,conv2,layer_norm1..3,sum_proj}.*`, `language_model.model.shared.{weight,scales,biases}`, `language_model.model.decoder.layers.N.*` (quantized linears carry `.scales/.biases` next to `.bias`), `language_model.model.decoder.{layernorm_embedding,layer_norm}.*`, `language_model.lm_head.{weight,scales,biases}`. The vision tower is not quantized in those exports. Conv transposes are shape-gated: `conv1.weight` is torch when `shape[-1] == 1 && shape[1] != 1`; `conv2.weight` is torch when `shape[1] == shape[0] && shape[-1] != shape[0]`.
## Implementation plan
1. **Detection and registry.** `src/models/detection.rs`: `"nemotron_parse" => Ok(ModelType::NemotronParseVLM)`. `src/models/mod.rs`: enum, supported list, description `("Nemotron-Parse (C-RADIO ViT-H + mBART seq2seq OCR)", "Other VLM")`. Tests in `src/models/detection_tests.rs`.
2. **Module layout (`src/models/nemotron_parse/`, new, mirroring `src/models/florence2/`).** `mod.rs` (config structs `NemotronParseConfig { vision: NemotronParseVisionConfig, text: NemotronParseTextConfig, decoder_start_token_id, eos_token_id, bos_token_id, pad_token_id, tie_word_embeddings, quantization }` parsed from the `encoder` / `decoder` sub-objects with the top-level overrides above; `NemotronParseSeqCache` = `Vec`-equivalent), `encoder.rs` (C-RADIO tower + neck), `decoder.rs` (pre-norm mBART), `model.rs` (`NemotronParseModel { encoder, shared: UnifiedEmbedding, decoder, lm_head: Option, config }` with `encode_image`, `decode`, `generate_greedy_with_cancel_and_prompt_len(pixel_values, seed_ids, max_new_tokens, cancel)`), `processor.rs` (resize/pad/normalize + prompt seeding), `runtime.rs` (`NemotronParseVlmModel` wrapping model + processor + tokenizer with `run(image, prompt, max_new_tokens, cancel) -> (text, prompt_tokens, generated_tokens)` and a `LanguageModel` impl that refuses decoder-only calls exactly like `Florence2VlmModel` at `src/models/florence2/runtime.rs:228`). Promote `KvCache`, `Florence2LayerCache`, `additive_causal_mask`, `layer_norm`, `linear`, and `Florence2Attention` to `pub(crate)` exports of `src/models/florence2/layers.rs` (they already are `pub(crate)`) and reuse them; `Florence2Attention` has exactly the `q_proj/k_proj/v_proj/out_proj` with bias this decoder needs.
3. **Encoder (`encoder.rs`).** `RadioEncoder::from_weights(weights, "vision_tower", &cfg)` accepting both key layouts through a `fn canonicalize_keys(WeightMap) -> WeightMap` in `checkpoint.rs` that applies the hub -> `vision_tower.*` / `language_model.*` renames and the two conv transposes. Store `conv1` as `Linear(1280, 1024)` (squeeze the trailing kernel axis). Implement `bilinear_align_corners(grid: [h, w, c], m) -> [m, m, c]` as a pure helper with a unit test on a 2x2 grid. Memory: run the 32 blocks in f16 on Apple Silicon (bf16 weights converted like every other tower), `[1, 13320, 1280]` activations are 34 MB each; the SDPA path keeps attention under a few hundred MB.
4. **Decoder (`decoder.rs`).** `NemotronParseDecoder { layernorm_embedding, layers: Vec, layer_norm, embed_scale: 32.0 }` with `forward(inputs_embeds, encoder_hidden, offset, caches)`; `NemotronParseDecoderLayer` wraps two `Florence2Attention`s and the three LayerNorms in the pre-norm order above. `make_cache()` returns one `Florence2LayerCache` per layer.
5. **Processor (`processor.rs`).** `NemotronParseImageProcessor { final_size: (2048, 1664), mean, std }` with `preprocess(&DynamicImage) -> UniquePtr [1, 3, 2048, 1664]` using `FilterType::Triangle` (bilinear) and white padding. `fn seed_ids(tokenizer, prompt) -> Vec` implementing the BOS/EOS stripping rule; `DEFAULT_TASK_PROMPT` constant.
6. **Loader.** `src/loading/vlm_nemotron_parse.rs` (pattern: `src/loading/vlm_florence2.rs`): read config, load weights with `load_vlm_weights_common`, canonicalize keys, build the model; `lm_head` = explicit tensor when present, else `shared.as_linear`. `LoadedModel::NemotronParseVLM(models::NemotronParseVlmModel)` in `src/loaded_model.rs` and dispatch in `src/loading/mod.rs`.
7. **CLI.** `src/commands/generate_nemotron_parse.rs` (pattern: `generate_florence2.rs`): require exactly one `--image`, use `-p` as the task prompt (warn and use `DEFAULT_TASK_PROMPT` when empty), `-n` as `max_new_tokens`, print the decoded text with `skip_special_tokens = false` so coordinate and class tokens survive, and print token counts. Dispatch from `src/commands/generate.rs` next to the Florence-2 branch (line 2431).
8. **Server.** Generalize `src/server/florence2_worker.rs` into a seq2seq worker that takes a `dyn Seq2SeqImageToText` trait object (`run(image, prompt, max_new_tokens, cancel)`), or add `run_nemotron_parse_worker_loop` beside it with the same request/response handling; select it in `src/server/model_worker.rs` at both `LoadedModel::Florence2VLM` match sites, add `ModelType::NemotronParseVLM` to `uses_single_stream_queue_admission` (`src/server/model_provider.rs:130`) and to the warmup skip (`src/server/startup.rs:1695`). Reject requests with more than one image or with audio/video using the existing `reject_media` / `reject_image_count` helpers. The text part of the user message is the task prompt; when absent use the default.
9. **Tokenizer.** The checkpoint ships `tokenizer.json` (fast tokenizer, `` 0, `` 1, `` 2, 2352 added tokens including `` 50004, `` 50008, `` 50001, `` 50010, `` 50009, the `` / `` coordinate tokens and `` tags from 52315). `MlxcelTokenizer::HuggingFace` loads it as-is; verify that `encode` with `add_special_tokens` wraps with ``/`` so the stripping rule applies.
10. **Docs.** `docs/supported-models.md`: add Nemotron-Parse 2.0 / v1.2 under document OCR with the task-prompt table (`predict_text_in_pic` vs `predict_no_text_in_pic`).
## Validation
(a) Unit tests: `src/models/nemotron_parse/nemotron_parse_tests.rs` (`patchify_order_matches_row_major_grid`, `pos_embed_native_page_is_identity_crop`, `bilinear_align_corners_2x2_to_3x3`, `neck_output_has_3329_rows_for_native_page`, `seed_ids_strip_wrapper_bos_eos_but_keep_prompt_specials`, `seed_ids_default_prompt_is_expected_sequence`, `decoder_prenorm_layer_shape_and_cache_growth`, `conv_transpose_gates_are_idempotent`), `src/models/detection_tests.rs`.
(b) Real checkpoint: `mlx-community/Nemotron-Parse-2.0-4bit` (0.7 GB) and `mlx-community/Nemotron-Parse-2.0-8bit`; the bf16 original `nvidia/NVIDIA-Nemotron-Parse-2.0` (1.5 GB) for the oracle.
```
./target/release/mlxcel generate -m models/Nemotron-Parse-2.0-8bit --image tests/fixtures/test_image.png \
-p "" -n 256
./target/release/mlxcel generate -m models/Nemotron-Parse-2.0-8bit --image /tmp/invoice_page.png \
-p "" -n 4096
```
Acceptance: finite logits; on a rendered text page (generate one with any word processor or `magick -size 1240x1754 xc:white -pointsize 40 -annotate +100+200 "Hello Nemotron" page.png`) the output contains the rendered words inside `...` structure; byte-identical greedy tokens against the checkpoint's own `transformers` implementation on CPU (the model card notes that its published CUDA golden output diverges after the third step cross-hardware, so compare against a local CPU run, not the golden file) for the first 64 tokens on the 8-bit and bf16 checkpoints.
## Acceptance criteria
- [ ] `nemotron_parse` resolves to `ModelType::NemotronParseVLM`
- [ ] Encoder produces `[1, 3329, 1024]` for a 2048x1664 page and loads both hub and converted key layouts
- [ ] Decoder is pre-norm with `sqrt(d_model)` embedding scale, no positional table, tied or untied head by checkpoint
- [ ] Decoder is seeded with the tokenized task prompt (wrapper BOS/EOS stripped); empty prompt falls back to the default with a warning
- [ ] CLI `--image` + `-p` path and `/v1/chat/completions` single-image path both run on the seq2seq worker
- [ ] 64-token greedy parity against a CPU run of the checkpoint's own implementation
- [ ] docs/supported-models.md updated
- [ ] detection table in src/models/detection.rs updated with a test
- [ ] cargo test --workspace --profile test-fast --features metal,accelerate passes
- [ ] cargo clippy --workspace --all-targets -- -D warnings and cargo fmt --all -- --check pass
## Out of scope
- The model card's table-insertion and repetition-stop logits processors and its markdown/LaTeX post-processing scripts.
- Batch size above 1 and multi-image requests.
- Non-native page sizes beyond the resize/pad rule (the encoder's interpolation branch is implemented but only the native grid is validated).
Contributor guide
Research direction
Start with the Florence-2 implementation and the named integration points in src/models/detection.rs, src/models/mod.rs, src/commands/generate.rs, and src/server/model_worker.rs. Then read src/models/detection_tests.rs and the listed Florence-2 layers, model, processor, CLI, and worker files before adding the new src/models/nemotron_parse/ module. Done means the model is registered, CLI and server paths accept the image and task prompt, and generation matches the documented markdown output and stopping behavior.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- backend, computer-vision, machine-learning
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 30/100