Add local Transformers inference for CUDA / Apple Silicon / CPU (fixes silent empty output on MPS)

[lucastsui]

Problem

The released modeling_unlimitedocr.py hardcodes CUDA (14 .cuda() calls, 3 torch.autocast("cuda", ...) blocks), so the model cannot run on Apple Silicon (MPS) or CPU at all.

Replacing .cuda() with .to(device) is not sufficient. On the MPS backend, torch.Tensor.masked_scatter_ is silently wrong under two independent conditions, and the image-embedding injection at

https://github.com/baidu/Unlimited-OCR (HF snapshot, modeling_unlimitedocr.py, the inputs_embeds[idx].masked_scatter_(...) line)

hits one of them. The mask at that call is images_seq_mask[idx].unsqueeze(-1), a [T, 1] tensor the op broadcasts across the hidden dim with stride 0, and a stride-0 (broadcast) mask mis-scatters even with a contiguous source. The source at that call is contiguous, since images_in_this_batch comes from torch.cat, which always allocates a fresh contiguous tensor (including for a single-element list). A non-contiguous source is a second, independent trigger of the same op bug. Measured divergence of the op against the CPU reference on MPS (isolated reproduction; float32 / float16 / bfloat16 all identical):

mask layout	contiguous source	non-contiguous source
materialized [T,H], contiguous	0.0000 (ok)	155.0 (wrong)
[T,1] via unsqueeze(-1) (the injection call's exact layout)	192.0 (wrong)	192.0 (wrong)
[T,H] stride-0 expand	192.0 (wrong)	192.0 (wrong)
materialized [T,H], non-contiguous strides	0.0000 (ok)	155.0 (wrong)

An exact mimic of the injection call (cat-built source, [T,1] mask, destination a batch view) diverges 264.0. Either trigger scrambles the visual tokens with no error raised, and the model emits an immediate end-of-sequence token: empty output, silently, for every image.

(An earlier revision of this description attributed the failure at this call site to source non-contiguity. That is a real trigger, but the one firing in this model is the broadcast mask; see the grid discussion in the comments, prompted by kushdab's question. The fix is unaffected either way: positional assignment avoids the op, and with it both triggers.)

What this PR adds (additive only, nothing existing is touched)

• patch_model_for_local.py — patches a local snapshot for device-agnostic use: replaces the injection with explicit positional assignment (mathematically identical; correct on CUDA, MPS, and CPU) and makes every hardcoded CUDA call follow the model's device. It asserts on the exact released code (1 injection fix + 19 counted device edits) so it fails loudly instead of mis-patching if the file changes, and it is idempotent.
• infer_transformers.py — single-process local inference via 🤗 Transformers on cuda / mps / cpu, mirroring infer.py's --image_dir / --pdf interface and the gundam / base image modes.
• tests/test_local_patch.py — unit tests that the positional-assignment replacement is exactly equivalent to masked_scatter_ semantics for contiguous and non-contiguous sources (CPU reference), and correct on MPS when available. The tests assert our replacement is right rather than that the upstream op is wrong, so they stay green if PyTorch fixes the underlying issue. The matrix includes the injection call's exact layout: [B,T,H] batch-view destination, [T,1] broadcast mask, cat-built source.
• README — a short "Local inference with Transformers" section.

hf download baidu/Unlimited-OCR --local-dir ./Unlimited-OCR-local
python patch_model_for_local.py ./Unlimited-OCR-local
python infer_transformers.py --model_dir ./Unlimited-OCR-local --image_dir ./my_pages

Validation

• pytest tests/test_local_patch.py: 5/5 pass (including the MPS cases, fp32 and bf16).
• End-to-end on an Apple M4 Max (48 GB, PyTorch 2.10, bfloat16): patched a pristine snapshot and transcribed a page of the Unlimited-OCR paper correctly on MPS (645 tokens, ~35 tok/s, base mode), where the unpatched port returns empty output. We have also run the patched model across a full 14-page document with stable results.
• On CUDA the device edits are behavioral no-ops and the positional-assignment injection computes the identical result (covered by the unit test); I could not run your GitHub Actions on CUDA hardware myself.

Notes

• The one-line positional-assignment fix could equally be adopted directly in the HF-hosted modeling_unlimitedocr.py (it is correct on CUDA too); happy to open that as a Hub PR instead or in addition, whichever you prefer.
• We also have opt-in decode-speed work for the MPS path (fused SDPA attention + a lean single-token MoE path, measured 1.28x decode at 0.0066% character deviation, and int8 weight-only quantization, 1.71x smaller) that we can offer as follow-up PRs if there is interest — kept out of this PR to keep the review surface small.

Open to maintainer direction on scope and placement.

[kushdab]

Diagnosis matches exactly what we found in #18 and documented in our HF Hub fix (Discussion #5): the masked_scatter_ call at

inputs_embeds[idx].masked_scatter_(images_seq_mask[idx].unsqueeze(-1).cuda(), images_in_this_batch)

is the actual point of failure on MPS, not just the .cuda() hardcode. Good to see this independently confirmed with a divergence measurement — the 0.0 vs 155.0 table is exactly the kind of evidence this repo's Apple Silicon threads (#18, #48, #49) have been missing.

Comparing the two fix mechanisms

Our HF Hub PR keeps masked_scatter_ and fixes the mask's device/broadcast shape so the op itself behaves correctly. This PR instead replaces the op entirely with explicit positional assignment. Both are mathematically valid fixes for the same root cause, but they carry different risk profiles:

• Positional assignment (this PR): sidesteps the MPS masked_scatter_ bug completely — correct regardless of whether PyTorch ever fixes the underlying op. Cost: it duplicates masking logic that lives elsewhere in the injection path, so if images_seq_mask construction changes upstream, this patch's assumptions need re-verification (your test_local_patch.py assertion on "1 injection fix + 19 counted device edits" is the right defense here — it'll fail loudly rather than silently mis-patch, which is exactly correct).
• Fixed masked_scatter_ (HF Hub 可视化模块是怎么做的 #5): smaller diff against upstream, easier to review/merge, but stays exposed if there's a second latent MPS bug in the same op family.

Given your quantitative validation is stronger (isolated divergence test + full 14-page document run), I'd lean toward this PR's approach being the safer merge candidate if the maintainers pick one.

One thing worth cross-checking

PR #57 (also just opened, Mac support) independently found that torch.autocast(device_type="mps", dtype=bfloat16) has a separate PyTorch/MPS bug that causes bbox coordinate loops (168168168...) after 5-10 tokens, and their fix is to disable autocast entirely on MPS rather than just retargeting its device string. Your infer_transformers.py — does it call torch.autocast anywhere in the MPS path, or run without it? If it uses autocast, worth testing whether #57's bbox-loop failure mode shows up in your 14-page run too (35 tok/s bf16 without loops suggests you may already be autocast-free, but flagging in case it's silent/intermittent like the masked_scatter_ bug was).

Overall: solid, well-tested contribution. The idempotent, fail-loud patch script design is the right call for a script that mutates a downloaded model snapshot.

[lucastsui]

To answer the autocast question directly, the MPS path runs autocast-free by construction. infer_transformers.py itself never calls torch.autocast, and patch_model_for_local.py rewrites the three hardcoded torch.autocast("cuda", ...) sites in modeling_unlimitedocr.py to torch.autocast(self.device.type, dtype=torch.bfloat16, enabled=(self.device.type == "cuda")), so on MPS the context manager is constructed but disabled. Weights and inputs are already bf16, so autocast only ever mattered for CUDA mixed-precision behavior, and gating it rather than deleting it keeps the CUDA path identical to upstream. That means the bbox-loop failure mode from #57 cannot trigger here, which matches the observed behavior of the 14-page run staying loop-free end to end.

Since you flagged #57, I also checked its other cross-platform claim, that transformers 4.57 silently ignores logits_processor entries that do not subclass LogitsProcessor. This matters for this PR because upstream's SlidingWindowNoRepeatNgramProcessor is a plain class and infer_transformers.py passes it on every run with no_repeat_ngram_size=35 and ngram_window=128. On transformers 4.57.1 the claim does not reproduce. GenerationMixin._merge_criteria_processor_list appends custom processors unconditionally, with no isinstance filter on the custom list, and LogitsProcessorList.call duck-types, so a two-argument call(input_ids, scores) takes the plain call branch.

import torch
from transformers.generation.logits_process import LogitsProcessorList
from transformers.generation.utils import GenerationMixin

class Plain:  # deliberately NOT a LogitsProcessor subclass
    def __init__(self): self.hits = 0
    def __call__(self, input_ids, scores):
        self.hits += 1
        return scores

p = Plain()
LogitsProcessorList([p])(torch.zeros(1, 5, dtype=torch.long), torch.zeros(1, 10))
assert p.hits == 1                      # applied, not ignored
merged = GenerationMixin._merge_criteria_processor_list(None, LogitsProcessorList(), [p])
assert p in merged                      # survives generate()'s merge

Both assertions pass on 4.57.1. So the ngram guard is active in this PR's runs as-is, which is likely one more reason the long run stays clean. Subclassing LogitsProcessor is harmless hygiene, but before treating the silent-ignore as a repo bug it would be good to know which transformers version #57 observed it on.

One more observation from reading #57 with your mechanism comparison in mind. It does not identify the injection-site bug at all. The only masked_scatter_ occurrence in its 6.2k-line diff is the unmodified upstream line, which its blanket rewrite of .cuda() calls to .to('mps') touches only incidentally. That may well land the same mask-device fix as your HF #5, which would explain their correct output, but it seems worth confirming before the two PRs are treated as equivalent on correctness.

On the two fix mechanisms, agreed with your read, and I would weight one factor more heavily than diff size. Positional assignment also survives any second latent bug in the same op family, and the MPS track record here with masked_scatter_ and autocast suggests that is a live risk.

[kushdab]

Your enabled=(self.device.type == "cuda") gating is actually the cleanest of the three approaches now that the root cause of #57's 100100100... loop is clear — just traced it down over there: #57 (comment)

Short version: modeling_deepseekv2.py's attention forward path has a fallback target_dtype = torch.get_autocast_gpu_dtype() that fires whenever query_states.dtype == torch.float32 after a LayerNorm cast-back. get_autocast_gpu_dtype() is the legacy CUDA-only accessor — it ignores whatever device_type your active torch.autocast(...) context is actually using. So if that fallback branch is reached while an enabled MPS autocast context is active, it silently pulls CUDA's cached dtype (defaults to fp16) instead of your actual bf16 MPS context, casting query_states to the wrong dtype mid-attention — plausible mechanism for the MLA+MoE collapse #57 is seeing when autocast is retargeted in forward() rather than removed.

Your gate sidesteps this cleanly: with enabled=False on MPS, torch.is_autocast_enabled() correctly reports False inside forward(), so that buggy fallback branch is never reached regardless of what's cached for CUDA. That's a better answer than either "delete autocast" (#57's current fix, loses nothing functionally but diverges further from upstream) or "retarget everywhere" (#57's original attempt, hits the bug) — you get upstream-identical CUDA behavior and correctness on MPS from one line, without needing to know about the get_autocast_gpu_dtype() trap at all. Worth flagging back to #57 if you want the credit for the cleaner fix pattern.

Also, nice catch on #57's diff not actually touching the masked_scatter_ injection line — I'd taken their "13/13 detections correct" test result as implicit confirmation the injection path was fine, but you're right that a blanket .cuda() → .to('mps') regex wouldn't fix the broadcast-mask issue on its own if the injection semantics are otherwise untouched. Given you're now also getting correct results, do you know if masked_scatter_ behaves correctly on .to('mps') alone once the mask's shape is already right, or does it specifically need the mask to be non-contiguous to trigger the divergence you measured (155.0)? If their test image happens not to exercise the non-contiguous path, "correct output" wouldn't tell us either way — might be worth them re-running your isolated reproduction case specifically.

[lucastsui]

Your masked_scatter_ question caught a real error in this PR's description, so taking that first. The description attributed the failure at the injection site to the source tensor's layout. For this call site that attribution was wrong. The 155.0 divergence is a real trigger, but it is not the one firing in Unlimited-OCR.

On the upstream line the source is contiguous and the mask is broadcast. images_in_this_batch is the output of torch.cat, which always allocates a fresh contiguous tensor, including for a single-element list (verified on 2.10). The mask is images_seq_mask[idx].unsqueeze(-1), shape [T, 1], which the op broadcasts across the hidden dim with stride 0. The full mask-layout by source-layout grid against the CPU reference (M4 Max, torch 2.10, fp32 and bf16 identical):

mask layout	contiguous source	non-contiguous source
materialized [T,H], contiguous	0.0000	155.0
[T,1] via unsqueeze(-1) (upstream's exact layout)	192.0	192.0
[T,H] stride-0 expand	192.0	192.0
materialized [T,H], non-contiguous strides	0.0000	155.0

An exact mimic of the real call (cat-built source, [T,1] mask, destination a batch view) diverges 264.0. So there are two independent triggers. A stride-0 broadcast mask is sufficient on its own, and a non-contiguous source is sufficient on its own. Mask non-contiguity as such is not a trigger, since a materialized mask with flipped strides is fine. The trigger firing in this model is the broadcast mask, and the PR's isolated test measured the other one. I have amended the PR description accordingly and will add the broadcast-mask case to test_local_patch.py's equivalence matrix.

That answers your question directly. .to('mps') alone on the upstream line keeps the [T,1] broadcast mask and stays wrong regardless of the source, and "mask shape already right" fixes this call site only if the mask is materialized to [T,H] real memory. A bare .expand(...) is equally wrong (row 3), so it is worth checking which of the two your HF #5 change produces. .expand(...).contiguous() or constructing the full-shape mask directly is correct here because upstream's source is contiguous, though it stays exposed to the source trigger if upstream ever hands the op a view. For #57 the grid removes the escape hatch we were both reaching for. Every image routes through this line, so there is no test image that fails to exercise it, and a bare device retarget scrambles every injection. If their 13/13 result is real, something else in the 6.2k-line diff is altering the call, or their test path does not reach this line. Running the grid against their patched snapshot would settle it, and that seems like the thing worth flagging over there rather than the autocast pattern.

Grid repro (runs on any Mac with MPS)

import torch
T, H, n = 16, 32, 6  # band rows 3:9

def source(kind, dtype):
    s = (torch.arange(n * H, dtype=torch.float32).reshape(n, H) + 1).to(dtype)
    if kind == "contig": return s
    s = s.t().contiguous().t()
    assert not s.is_contiguous()
    return s

def mask(kind, dev):
    band = torch.zeros(T, dtype=torch.bool, device=dev); band[3:3+n] = True
    return {"materialized":   band.unsqueeze(-1).expand(T, H).contiguous(),
            "unsqueeze[T,1]": band.unsqueeze(-1),                      # upstream exact
            "expand-stride0": band.unsqueeze(-1).expand(T, H),
            "mat-noncontig":  band.unsqueeze(0).expand(H, T).contiguous().t()}[kind]

def cell(mk, sk, dtype):
    out = {}
    for dev in ("cpu", "mps"):
        dst = torch.zeros(T, H, dtype=dtype, device=dev)
        dst.masked_scatter_(mask(mk, dev), source(sk, dtype).to(dev))
        out[dev] = dst.cpu().float()
    return (out["cpu"] - out["mps"]).abs().max().item()

for dtype in (torch.float32, torch.bfloat16):
    for mk in ("materialized", "unsqueeze[T,1]", "expand-stride0", "mat-noncontig"):
        print(dtype, f"{mk:15s}", [cell(mk, sk, dtype) for sk in ("contig", "noncontig")])

# exact real-call layout: cat-built (contiguous) source, [T,1] mask, batch-view dst
ref = {}
for dev in ("cpu", "mps"):
    emb = torch.zeros(2, T, H, dtype=torch.bfloat16, device=dev)
    sm = torch.zeros(2, T, dtype=torch.bool, device=dev); sm[0, 3:3+n] = True
    src = torch.cat([torch.arange(4*H, dtype=torch.float32).reshape(4, H) + 1,
                     torch.arange(2*H, dtype=torch.float32).reshape(2, H) + 200]
                    ).to(torch.bfloat16).to(dev)
    assert src.is_contiguous()
    emb[0].masked_scatter_(sm[0].unsqueeze(-1), src)
    ref[dev] = emb.cpu().float()
print("real-call divergence:", (ref["cpu"] - ref["mps"]).abs().max().item())  # 264.0

On the autocast trace, the proposed mechanism cannot fire under a retargeted MPS context, because the branch's guard is the same legacy CUDA-scoped accessor as its payload. Measured on torch 2.10, inside an enabled torch.autocast("mps", dtype=torch.bfloat16) context, torch.is_autocast_enabled() returns False, torch.is_autocast_enabled("mps") returns True, and torch.get_autocast_gpu_dtype() returns float16 with a deprecation message pointing to get_autocast_dtype('cuda') as its replacement. The elif torch.is_autocast_enabled(): one line above the get_autocast_gpu_dtype() call is therefore False whenever only MPS autocast is active. The branch is reachable only from an enabled CUDA autocast context, where pulling CUDA's cached dtype is the intended behavior, so the mispairing you describe, an enabled MPS context feeding a CUDA dtype lookup, is exactly the state the guard excludes.

Two further blocks make it moot for this checkpoint regardless. The fallback lives only in DeepseekV2FlashAttention2.forward, and with use_mla=false the decoder looks up mha_* keys in ATTENTION_CLASSES, where mha_eager is the only entry (mha_flash_attention_2 is commented out), so the hosting class is never instantiated. The branch also requires query states arriving as fp32, which a bf16 inference pipeline does not produce. So I would not credit the enabled= gate with dodging this particular trap, since nothing in this model can reach it. The gate's value stays what it was, a CUDA path byte-identical to upstream plus the guarantee that torch.is_autocast_enabled("mps") reads False for any code that consults the per-device accessor. Which leaves #57's 100100... loop still without a mechanism. Whatever it is, it is not this branch under an MPS-retargeted autocast.

[kushdab]

@lucastsui that grid is thorough and it caught something real — thanks for running it rather than trusting the PR description (mine included).

You asked whether HF Hub #5 materializes the mask or leaves it as a stride-0 broadcast. Checked the actual patch — it's the latter:

# HF Hub Discussion #5 (kushdab), the fix as posted:
_mask = images_seq_mask[idx].unsqueeze(-1).expand_as(inputs_embeds[idx]).to(inputs_embeds.device)
inputs_embeds[idx].masked_scatter_(_mask, images_in_this_batch)

expand_as is implemented via expand, so this is exactly your row 3 (expand-stride0) — a [T,1]-derived stride-0 view over [T,H], not a materialized [T,H] mask. Per your grid, that's the layout that diverges regardless of source contiguity (192.0 either way). So the HF Hub #5 fix removes the shape-mismatch crash (the mask is now broadcastable, so masked_scatter_ no longer raises RuntimeError: wrong number of elements) but does not remove the numerical MPS masked_scatter_ bug you found — it's still exposed to exactly the trigger you isolated. It happened to test clean because whatever image count/shape I validated against didn't hit the divergence, not because the fix is correct.

The actual fix per your grid is .expand_as(...).contiguous() (materializes the [T,H] mask into real memory) rather than the bare expand_as. I'll get that correction posted on the HF Hub discussion — flagging it here first since you're the one who found it and #56 is the PR that should carry the real fix upstream anyway.

On the autocast side: agreed with your read, and separately confirmed with Konsn666 on #57 that the elif torch.is_autocast_enabled(): get_autocast_gpu_dtype() branch lives in DeepseekV2FlashAttention2, which this checkpoint's use_mla=false config never instantiates (mha_eager → SlidingWindowLlamaAttention, mha_flash_attention_2 is commented out of ATTENTION_CLASSES). So that branch is unreachable here regardless of the MPS/CUDA autocast-guard question — the 100100100... loop in #57 still needs a different mechanism, most likely op-level under MPS autocast rather than this code path.