Support NVIDIA Blackwell architecture

[hongbo-miao]

Experiment 1

I try to run on GeForce RTX 5090 (NVIDIA Blackwell) using uv.
olmocr is using an old sglang, so I try to install sglang separately like sglang[blackwell]:

[project]
name = "hm-olmocr"
version = "1.0.0"
requires-python = "~=3.11.0"
dependencies = [
  "olmocr==0.1.71",
  "sglang[blackwell]==0.4.7",
]

[tool.uv]
package = false
required-version = ">=0.6.0"

To reproduce, you can run:

curl -o olmocr-sample.pdf https://olmocr.allenai.org/papers/olmocr_3pg_sample.pdf
uv run python -m olmocr.pipeline ./localworkspace --markdown --pdfs olmocr-sample.pdf

However, I got error:

uv sync
  × No solution found when resolving dependencies:
  ╰─▶ Because olmocr==0.1.71 depends on transformers==4.46.2 and sglang[blackwell]==0.4.7 depends on transformers==4.52.3, we can conclude that olmocr==0.1.71 and sglang[blackwell]==0.4.7 are incompatible.
      And because your project depends on olmocr==0.1.71 and sglang[blackwell]==0.4.7, we can conclude that your project's requirements are unsatisfiable.
error: Recipe `uv-install-dependencies` failed on line 5 with exit code 1

It would be great to support NVIDIA Blackwell architecture, thanks! ☺️

Experiment 2 (Docker way)

I also tried Docker way

docker run -it --gpus all \
  -v /path/to/your/local/files:/local_files \
  --name olmocr_container \
  alleninstituteforai/olmocr:v0.1.71 /bin/bash

However, you can see in the middle, PyTorch had issue.
In this case PyTorch needs to be new version and set https://download.pytorch.org/whl/cu128, more info can be found here.

Here is full log: https://gist.github.com/hongbo-miao/fe51beaa5faa2477ddb72c42e1914d96

[jakep-allenai]

Yeah, can you try with 0.1.75 which just came out? We have a CUDA 12.8 based docker image now, and a new VLLM version instead of sglang.

[hongbo-miao]

Thank you @jakep-allenai !

Experiment 3 (Upgrade to olmocr==0.1.75)

I tried

[project]
name = "hm-olmocr"
version = "1.0.0"
requires-python = "~=3.11.0"
dependencies = [
  "olmocr==0.1.75",
  "torch==2.7.1",
]

[tool.uv]
package = false
required-version = ">=0.6.0"

[tool.uv.sources]
torch = [
  { index = "pytorch-cu128" },
]

[[tool.uv.index]]
name = "pytorch-cu128"
url = "https://download.pytorch.org/whl/cu128"
explicit = true

torch has to use https://download.pytorch.org/whl/cu128 to support Blackwell architecture.

curl -o olmocr-sample.pdf https://olmocr.allenai.org/papers/olmocr_3pg_sample.pdf
uv run python -m olmocr.pipeline ./localworkspace --markdown --pdfs olmocr-sample.pdf

failed with log.

But I think it is getting close! ☺️

I believe right now need wait vLLM fully support Blackwell architecture, I only succeed in Docker way: vllm-project/vllm#14452 (comment)

Experiment 4 (Add vllm explicitly)

[project]
name = "hm-olmocr"
version = "1.0.0"
requires-python = "~=3.11.0"
dependencies = [
  "olmocr==0.1.75",
  "torch==2.7.0",
  "vllm==0.9.1",
]

[tool.uv]
package = false
required-version = ">=0.6.0"

[tool.uv.sources]
torch = [
  { index = "pytorch-cu128" },
]

[[tool.uv.index]]
name = "pytorch-cu128"
url = "https://download.pytorch.org/whl/cu128"
explicit = true

curl -o olmocr-sample.pdf https://olmocr.allenai.org/papers/olmocr_3pg_sample.pdf
uv run python -m olmocr.pipeline ./localworkspace --markdown --pdfs olmocr-sample.pdf

The error looks different, here is full log.

[jakep-allenai]

Yeah, I've gotten an experiment to run on blackwell, but it was really slow (much slower than H100), so we are also still waiting for support to become better. However, it did run at least when I tried. What device is this on?

[jakep-allenai]

Ahh I see RTX 5090, haven't tried that myself yet

[hongbo-miao]

Hi @jakep-allenai yeah, this is RTX 5090. I updated my previous two posts to make it more clear for how to reproduce. Basically just

curl -o olmocr-sample.pdf https://olmocr.allenai.org/papers/olmocr_3pg_sample.pdf
uv run python -m olmocr.pipeline ./localworkspace --markdown --pdfs olmocr-sample.pdf

[jakep-allenai]

Curious if you try our FP8 version? https://huggingface.co/allenai/olmOCR-7B-0225-preview-FP8

Just add --model allenai/olmOCR-7B-0225-preview-FP8

[hongbo-miao]

Experiment 5 (using --model allenai/olmOCR-7B-0225-preview-FP8)

[project]
name = "hm-olmocr"
version = "1.0.0"
requires-python = "~=3.11.0"
dependencies = [
  "olmocr==0.1.75",
  "torch==2.7.0",
  "vllm==0.9.1",
]

[tool.uv]
package = false
required-version = ">=0.6.0"

[tool.uv.sources]
torch = [
  { index = "pytorch-cu128" },
]

[[tool.uv.index]]
name = "pytorch-cu128"
url = "https://download.pytorch.org/whl/cu128"
explicit = true

curl -o olmocr-sample.pdf https://olmocr.allenai.org/papers/olmocr_3pg_sample.pdf
uv run python -m olmocr.pipeline ./localworkspace --markdown --pdfs olmocr-sample.pdf  --model allenai/olmOCR-7B-0225-preview-FP8

This time failed with log.

Note inside, it has

SystemError: PY_SSIZE_T_CLEAN macro must be defined for '#' formats

I recently met same error when use Unsloth unslothai/unsloth#1856 on RTX 5090, and I haven't figured out the way to use Unsloth as well.