How olmOCR Manages GPU Memory with Large Documents and vLLM

olmOCR processes PDFs one page at a time through a vLLM inference server, capping each request at 8,000 tokens to ensure the GPU's KV-cache never exceeds available VRAM.

The olmOCR pipeline from the Allen Institute for AI (AI2) is designed to process millions of PDF pages efficiently without triggering out-of-memory errors. By chunking documents at the page level and explicitly controlling vLLM's memory allocation parameters, the system maintains stable GPU performance even when handling heterogeneous document collections with varying page sizes.

Per-Page Request Construction

The foundation of olmOCR's memory management strategy lies in single-page processing. Instead of feeding entire documents to the language model, the pipeline constructs individual requests for each rendered page image.

The build_page_query Implementation

In olmocr/pipeline.py, the build_page_query function (lines 6-45) creates inference requests with a hardcoded MAX_TOKENS = 8000 limit. This design ensures that the LLM never processes more than one page and a short prompt simultaneously, keeping the attention matrices and KV-cache allocation constant regardless of the source document's total length.

By restricting the input to a single rendered image plus minimal prompt text, the GPU memory required for attention computations remains predictable and bounded.

vLLM Server Configuration for Memory Control

The vllm_server_task function in olmocr/pipeline.py (lines 7-33) launches the inference server with specific flags that directly govern VRAM consumption:

• --gpu-memory-utilization <float> – Sets the fraction of GPU VRAM allocated for model weights and KV cache
• --max-model-len <int> – Caps the maximum token length the server will accept
• --tensor-parallel-size and --data-parallel-size – Distribute the model across multiple GPUs when available

Example launch configuration generated by the pipeline:

vllm serve allenai/olmOCR-2-7B-1025-FP8 \
    --port 8000 \
    --tensor-parallel-size 1 \
    --data-parallel-size 1 \
    --gpu-memory-utilization 0.85 \
    --max-model-len 16384 \
    --limit-mm-per-prompt '{"video": 0}'

This explicit allocation prevents vLLM from automatically consuming all available GPU memory, leaving headroom for CUDA overhead and transient allocations.

GPU Availability and Cleanup Mechanisms

Before processing begins, olmOCR validates hardware availability and clears residual memory to prevent allocation conflicts.

Pre-Processing GPU Checks

The pipeline imports check_torch_gpu_available from olmocr.check (referenced at lines 32-35 of pipeline.py) to verify CUDA device availability. If no GPU is detected, the system gracefully falls back to CPU execution using identical code paths, albeit with reduced throughput.

Checkpoint Memory Management

When loading training checkpoints, olmocr/train/prepare_checkpoint.py explicitly calls torch.cuda.empty_cache() at lines 373-375. This frees any leftover VRAM from previous operations before inference workers initialize, eliminating out-of-memory errors caused by residual tensors from checkpoint loading.

Dynamic Request Management

The pipeline monitors vLLM's internal queue state to prevent GPU overload during high-throughput processing.

Queue Monitoring and Back-off

The process_page function tracks the vllm_queued_requests variable, which the server-log parser updates at lines 76-80 of pipeline.py. When the queue empties, the pipeline fires remaining retries in parallel, but each retry still transmits only a single-page request. This mechanism ensures that GPU load never spikes beyond the per-request memory limit, even when retrying failed pages.

To inspect server-side memory usage programmatically:

if match := re.search(r"Running: (\d+)", line):
    current_running = int(match.group(1))
    logger.info(f"vllm running req: {current_running} queue req: {vllm_queued_requests}")

Summary

• olmOCR processes documents one page at a time, keeping each request under 8,000 tokens to bound GPU memory usage
• vLLM server flags (--gpu-memory-utilization, --max-model-len) explicitly control VRAM allocation for model weights and KV cache
• GPU validation occurs before processing via check_torch_gpu_available, with automatic CPU fallback
• Memory cleanup via torch.cuda.empty_cache() in prepare_checkpoint.py prevents residual allocation conflicts
• Dynamic queue monitoring ensures parallel retries never exceed per-page memory limits

Frequently Asked Questions

How does olmOCR handle documents longer than the context window?

olmOCR never processes full documents in a single inference call. The build_page_query function in pipeline.py constructs individual requests for each page image with a hard limit of 8,000 tokens, ensuring that even multi-thousand-page documents are handled as a series of bounded, independent operations that fit comfortably within GPU memory constraints.

What GPU memory utilization setting should I use for large batch processing?

The default --gpu-memory-utilization 0.85 (85%) works well for most single-GPU deployments, leaving 15% of VRAM for CUDA overhead and system operations. For large batch processing with multiple workers, you may reduce this value to 0.80 or lower to accommodate concurrent request buffering, as specified in the pipeline's vLLM launch configuration at lines 7-33 of pipeline.py.

Can olmOCR run on CPU if GPU memory is insufficient?

Yes. The check_torch_gpu_available function validates CUDA device presence before processing begins. If no GPU is detected or available, the pipeline automatically falls back to CPU execution. While this eliminates GPU memory constraints entirely, processing speed decreases significantly compared to GPU-accelerated inference.

How does the pipeline prevent out-of-memory errors during retries?

The process_page function monitors the vllm_queued_requests counter (updated in the log parser at lines 76-80 of pipeline.py) to track server load. When retrying failed pages, the pipeline waits for the queue to empty before issuing parallel requests, ensuring that each retry still adheres to the single-page, 8,000-token memory limit rather than accumulating unbounded requests in GPU memory.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}