Apple Silicon (MPS): empty output — masked_scatter with a broadcast mask scatters incorrectly on Metal

Unlimited-OCR: Empty Output on Apple Silicon (MPS / Metal)

Summary

On Apple Silicon using MPS / Metal, model.infer(...) returns an empty string for every input.

The model loads and runs without any error, but it generates 0 tokens and terminates immediately with EOS.

Root Cause

The image-embedding injection in modeling_unlimitedocr.py uses Tensor.masked_scatter_ with a broadcast mask.

PyTorch’s MPS backend appears to scatter broadcast masks incorrectly. As a result, the image embeddings are written to the wrong token positions. The model therefore never receives a coherent image representation and emits EOS immediately.

Environment

• Apple M-series GPU

• macOS

• device="mps"

• torch==2.10.0

• transformers==4.57.1

• Python 3.12

• Reproduced with:

  • bfloat16
  • float16
  • float32

Because all tested dtypes are affected, this does not appear to be a precision issue.

Location in the Code

The issue occurs in modeling_unlimitedocr.py, during image injection inside the model forward pass:

inputs_embeds[idx].masked_scatter_(
    images_seq_mask[idx].unsqueeze(-1),   # mask shape: (seq, 1), broadcast to (seq, hidden)
    images_in_this_batch,
)

The call to .unsqueeze(-1) creates a mask of shape (seq, 1), which masked_scatter then broadcasts to (seq, hidden).

That broadcast path is incorrect on MPS.

Minimal Reproduction

No model is required to reproduce the issue:

import torch

mask = torch.tensor([True, False, True, False]).unsqueeze(-1)  # shape: (4, 1), broadcast
src = torch.arange(1, 7).float()                               # 2 true rows × 3 cols

for dev in ("cpu", "mps"):
    x = torch.zeros(4, 3, device=dev)
    print(dev, x.masked_scatter(mask.to(dev), src.to(dev)).flatten().tolist())

Expected CPU result:

cpu -> [1, 2, 3, 0, 0, 0, 4, 5, 6, 0, 0, 0]

Observed MPS result:

mps -> [1, 2, 3, 4, 5, 6, 0, 0, 0, 0, 0, 0]

Confirmed Workaround

With an explicitly expanded and contiguous mask, MPS behaves correctly:

mask_full = mask.expand(4, 3).contiguous()

Correct MPS result:

mps -> [1, 2, 3, 0, 0, 0, 4, 5, 6, 0, 0, 0]

Suggested Fix

A robust model-side fix is to expand the mask to the target shape before scattering:

m = (
    images_seq_mask[idx]
    .unsqueeze(-1)
    .to(inputs_embeds.device)
    .expand_as(inputs_embeds[idx])
    .contiguous()
)

inputs_embeds[idx] = inputs_embeds[idx].masked_scatter(
    m,
    images_in_this_batch,
)

This avoids relying on the broken broadcast-mask behavior in PyTorch’s MPS backend.

Impact

This is ultimately a PyTorch MPS bug in masked_scatter with a broadcast mask.

However, guarding against it in the model makes Unlimited-OCR work on Apple Silicon today. With the proposed change, the model produces correct, grounded output on M-series GPUs.

This has been verified end-to-end on German legal PDFs. :::