ragarwal/chandra-ocr-2-nvfp4-mtp

Chandra OCR 2, NVFP4 + fine-tuned MTP

A quantized version of datalab-to/chandra-ocr-2 for fast serving with vLLM, with a reference client and benchmark harness.

• Quantization (ModelOpt): NVFP4 weights and activations (W4A4) in the MLP layers, FP8 in the attention projections, the rest in BF16. MSE calibration, then quantization-aware distillation against the BF16 model.
• Speculative decoding: the model's built-in MTP head, fine-tuned on in-domain document pages. Use 3 speculative tokens.
• Size: 6.0 GB, against 10.6 GB for BF16.

Results

2x NVIDIA B200, vLLM 0.33.0, the server and client settings below. BF16 is the original model as released, without MTP; its scores are the mean (range) of 3 runs.

	BF16	This model
Throughput, output tokens/s (sustained, 640 concurrent requests)	24,742	38,891 (1.57x)
omni_parse_bench	88.6 (88.5-88.7)	88.5
olmOCR-Bench	84.3	83.9
Private test set (tolerant pass rate)	91.4 (91.3-91.6)	91.3

Server

Image: docker.io/vllm/vllm-openai:v0.33.0 (CUDA 13: NVIDIA driver 580 or newer).

docker run -d --gpus all --ipc=host --network=host \
  -v ~/.cache/huggingface:/root/.cache/huggingface \
  docker.io/vllm/vllm-openai:v0.33.0 \
  ragarwal/chandra-ocr-2-nvfp4-mtp --served-model-name chandra --port 8000 \
  --limit-mm-per-prompt '{"image": 1}' --max-model-len 24576 --gpu-memory-utilization 0.85 \
  --mm-processor-kwargs '{"min_pixels": 3136, "max_pixels": 6291456}' \
  --data-parallel-size 2 --max-num-seqs 256 --max-num-batched-tokens 65536 \
  --api-server-count 4 --renderer-num-workers 8 --performance-mode throughput \
  --speculative-config '{"method": "mtp", "num_speculative_tokens": 3}'

Set --data-parallel-size to the number of GPUs. Full speed needs Blackwell GPUs (FP4 tensor cores); on older GPUs such as A100, vLLM falls back to weight-only kernels, which is untested.

Client

ocr.py is the reference client. It follows Chandra's own pipeline:

• PDF pages rendered at 150 dpi (at least 1024 px on the short side), then Chandra's image scaling;
• Chandra's layout prompt, greedy decoding (temperature 0, top_p 0.1, max_tokens 12384);
• Chandra's retry rule: a reply that repeats itself or is cut off is retried up to 6 times at temperature 0.2-0.8 and top_p 0.95, with fixed seeds so reruns send identical requests.

It writes markdown and the raw HTML for each page.

pip install -r requirements.txt
python ocr.py report.pdf scans/*.png --out ocr_out --url http://localhost:8000/v1

--loop-stop adds vLLM's repetition detector, which ends runaway loops early (a reply ending in 64 copies of a 1-64 token pattern). That saves about 11% of generated tokens. It was not used for the scores above.

Reproducing the results

The public benchmarks are run with benchmark.py, which sends every page through ocr.py, and scored with each benchmark's own scorer. The private test set is not public.

hf download datalab-to/omni_parse_bench --repo-type dataset --revision d70182b9f901d1660a27e60bf6be6d6e6ea83013 --local-dir omni_parse_bench
hf download allenai/olmOCR-bench --repo-type dataset --revision 54a96a6fb6a2bd3b297e59869491db4d3625b711 --local-dir olmOCR-bench

# With the server above running:
python benchmark.py omni   --data omni_parse_bench --out runs/omni
python benchmark.py olmocr --data olmOCR-bench     --out runs/olmocr

opb grade runs/omni --dataset omni_parse_bench        # headline: runs/omni/unit_tests/summary.json

# olmOCR-Bench scorer, in a separate environment:
pip install -r requirements-olmocr.txt && playwright install chromium
python -m olmocr.bench.benchmark --dir runs/olmocr/bench --force

Expect small run-to-run differences: outputs depend slightly on how requests are batched, even with greedy decoding. BF16 ranged 0.2 points on omni_parse_bench over 3 runs.

How it was made

All training used a private set of document pages (PhD theses, many languages) rendered at 150 dpi. No benchmark pages were used. Targets are the BF16 model's own outputs, so no ground truth was needed.

1. Post-training quantization (ModelOpt).
   • MLP layers: NVFP4 weights and activations.
   • Attention projections: FP8.
   • Vision encoder, embeddings, lm_head and MTP head: BF16. KV cache is not quantized.
   • Scales from MSE calibration on 128 pages.
2. Quantization-aware distillation.
   • The quantized model is trained to match the BF16 model's token distributions (KL divergence, teacher-forced on the BF16 outputs) on 12,000 pages.
   • Text-model weights only; vision encoder, embeddings and lm_head frozen.
   • 1,500 updates at lr 1e-6 (AdamW), about 2 hours on 4 B200s.
   • At lr 1e-5 the model drifted away from BF16 instead.
   • On held-out pages, KL per token fell 41% and top-1 agreement with BF16 rose from 99.22% to 99.33%.
3. MTP head fine-tuning (speculators).
   • Chandra's built-in 1-layer MTP head trained on the BF16 model's hidden states and outputs for about 23,600 pages.
   • 3 epochs, AdamW, lr 3e-5, 3 speculative steps, on one GPU; then stitched into the checkpoint.
   • Acceptance at the first two draft positions rose from 93% and 77% (stock head) to 96% and 87% (measured on an NVFP4 W4A16 build with 2 speculative tokens).

License

Derived from datalab-to/chandra-ocr-2 and distributed under its license (modified OpenRAIL-M, included as LICENSE), including its use restrictions. Changes: quantized as described above, and the MTP head fine-tuned.