PyTorch MPS Backend in MinerU#

Overview#

PyTorch's Metal Performance Shaders (MPS) backend enables GPU acceleration on Apple Silicon (M1/M2/M3) Macs. In MinerU's pipeline backend, MPS is the second-priority device after CUDA, selected automatically when CUDA is unavailable 1. The device selection order is: CUDA → MPS → NPU → GCU → MUSA → MLU → SDAA → CPU.

To override auto-selection, set:

MINERU_DEVICE_MODE=mps # force MPS
MINERU_DEVICE_MODE=cpu # force CPU

MPS-Specific Behaviors vs. CUDA#

1. Operation Fallback (PYTORCH_ENABLE_MPS_FALLBACK)#

pipeline_analyze.py sets os.environ['PYTORCH_ENABLE_MPS_FALLBACK'] = '1' at module import time. This instructs PyTorch to silently fall back to CPU for any MPS-unsupported operations rather than raising an error. The comment translates to "allow MPS to fallback."

Implication: Operations unsupported by MPS (e.g., certain custom CUDA kernels) will silently run on CPU, which can cause unexpected performance degradation during inference.

2. Attention Implementation — Eager Mode Required#

The UnimernetModel (formula recognition) detects the device at init time 2:

• MPS / NPU / MUSA: model loaded with attn_implementation="eager" — FlashAttention is not available on these backends.
• CUDA / CPU: default attention implementation used.

This means the MPS path cannot benefit from FlashAttention's memory efficiency, which matters for long formula sequences.

3. No Pinned Memory or Non-Blocking Transfers#

In UnimernetModel, pinned memory and non-blocking tensor transfers are explicitly CUDA-only:

@staticmethod
def _should_pin_memory(device) -> bool:
    return str(device).startswith("cuda")

On MPS, DataLoader(pin_memory=False) and blocking .to(device) calls are used. This reduces throughput compared to CUDA's asynchronous data transfer path 3.

4. Float16 Inference (Shared with CUDA)#

All non-CPU devices, including MPS, cast models to float16 after loading 4. This is identical behavior to CUDA — MPS benefits from halved activation memory.

Batch Sizing and VRAM Detection#

get_vram() in model_utils.py auto-detects VRAM to compute the batch_ratio multiplier. MPS has no detection branch; the function returns a default of 1 GB for any unrecognized device 5. This means:

• MPS devices always get batch_ratio = 1 (the lowest tier, for devices with < 6 GB).
• Base batch sizes at batch_ratio = 1: layout=1, MFR=16, OCR-det=8 6.
• Apple Silicon's unified memory (often 16–96 GB) is never detected, so MPS runs conservatively small batches by default.

To override: MINERU_HYBRID_BATCH_RATIO=<n> or MINERU_VIRTUAL_VRAM_SIZE=<GB>.

Memory Cleanup#

clean_memory() handles MPS explicitly 7:

elif str(device).startswith("mps"):
    torch.mps.empty_cache()

Unlike CUDA and NPU, there is no availability check before calling torch.mps.empty_cache(). This is safe as long as mps is in the device string, but unlike CUDA/NPU it cannot be guarded against a missing backend.

clean_vram() 8 only fires clean_memory() when the detected VRAM ≤ 8 GB. Since MPS defaults to 1 GB (no detection), clean_vram will fire after each inference stage (layout, MFR, OCR-det), which is the correct behavior for memory-constrained scenarios.

Key Files#

Summary of MPS vs. CUDA Differences#