Skip to content

fix: cache warmup RuntimeError on mps #46239

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ArthurZucker merged 1 commit into from
Jun 1, 2026
Merged

fix: cache warmup RuntimeError on mps #46239

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions src/transformers/modeling_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5138,6 +5138,12 @@ def caching_allocator_warmup(model: PreTrainedModel, expanded_device_map: dict,
# Note that we use an absolute value instead of device proportion here, as a 8GiB device could still allocate too much
# if using e.g. 90% of device size, while a 140GiB device would allocate too little
byte_count = min(byte_count, total_device_memory - 1.2 * 1024**3)
elif device.type == "mps":
# Skip warmup on MPS: there is a limit of the maximum size a single buffer can have on MPS,
# which from testing seems to be about 2/3 of the total device memory (tested on apple silicon).
# This causes the warmup function to return a `RuntimeError: Invalid buffer size: XX.XX GiB`.
# NOTE: not tested on intel macs
continue
Comment on lines +5142 to +5146

@ArthurZucker ArthurZucker May 27, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

please provide a repro, it does not fail for me AFAIK! Loading a big mixtral to max capa!
Also no skip maybe reduce allocate + bench speed loss please

@McPatate McPatate May 28, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We won't need pre-allocation with safetensors on mps once safetensors/safetensors#767 is merged. We allocate the mtlbuffers, fill them with pread and then hand them 0-copy to torch with dlpack. So as we don't go through torch's allocation stack, it's going to become unnecessary, at least for mps.

As we discussed by message, you in fact cannot allocate a buffer of size over 58gb out of your 96 available. It'd be interesting to see what total_byte_count's value is when you load your Mixtral model.

@McPatate McPatate May 28, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added more details in the PR desc

@McPatate McPatate May 28, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The only reason I can see your load not crash is because you don't set device_map="mps", which skips the cache warmup function altogether.

# We divide by 2 here as we allocate in fp16
_ = torch.empty(int(byte_count // 2), dtype=torch.float16, device=device, requires_grad=False)

Expand Down
Loading