[Feature Request] 3B Model

[charitarthchugh]

🚀 The feature, motivation and pitch

I have seen that many users have constrained hardware choices/or are converting pdfs at large scales. Would it be beneficial to train a 3B model (based on Qwen2.5-vl)?

Besides the speedup from switching to a smaller model, we can also further tune for speculative decoding methods such as EAGLE which in turn would make inference of the base model faster.

If that sounds reasonable, I would love to help and contribute to creating this.

Alternatives

No response

Additional context

(For me specifically, I have to convert ~1M pdfs and I don't have enough compute to do this in a reasonable timeframe)

[jakep-allenai]

I think one of our team members tried to train the 3B, but training wasn't stable with the 7B hyperparameters. It would be nice to have such a model, and it would likely be faster and cheaper to inference a little bit. I've found that in batch-inference contexts though, speculative decoding will not give the gains that you might be hoping for. And in fact, the best way we can make olmOCR cheaper at inference in general would be to invest into something like a compressed KV cache instead. We could probably cut inference cost 30-40% and preserve most of the quality this way, while a 3B model would be a big step down in quality.

[charitarthchugh]

I think one of our team members tried to train the 3B, but training wasn't stable with the 7B hyperparameters. It would be nice to have such a model, and it would likely be faster and cheaper to inference a little bit.

That's interesting. Would it be worthwhile looking into PEFT for the 3B and doing a benchmark of it?

And in fact, the best way we can make olmOCR cheaper at inference in general would be to invest into something like a compressed KV cache instead. We could probably cut inference cost 30-40% and preserve most of the quality this way, while a 3B model would be a big step down in quality.

It seems like the vLLM v1 engine does not support this KV Cache quantization yet (so did not look further).

[jakep-allenai]

VllM should support it, but every time I tried I have gotten garbled responses from the model. I've used llmcompressor for the quantization. I think it might need some quantization aware training or similar to really get working right.

[charitarthchugh]

So do you think I should try the in-flight quantization strategies supported by vLLM?

[jakep-allenai]

I've tried a bit of them: https://github.com/allenai/olmocr/blob/main/olmocr/train/compress_checkpoint.py
That script supports anything that llm compressor supports, and we use it to make the FP8 version of olmocr.

But when I've tried this config: https://github.com/allenai/olmocr/blob/main/olmocr/train/quantization_configs/qwen2_5vl_w8a8_fp8_kv8.yaml to enable KV cache compression, I always get garbled outputs, no matter what.

If you can figure it out, it would be a huge win, since a single GPU could pack many more concurrent requests.