fix MPS errors#3818
Merged
Merged
Conversation
tomaarsen
reviewed
Jun 12, 2026
tomaarsen
left a comment
tomaarsen
left a comment
Member
There was a problem hiding this comment.
Overall this looks good! I made a small comment to shrink the new comment a bit. Thanks for verifying that this works on your hardware 🤗
tomaarsen
enabled auto-merge (squash)
June 12, 2026 10:04
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Three spots in the library still hardcode CUDA, even though MPS is listed as a supported device in
get_device_name(). This PR fixes them.What changes
sparse_encoder/model.py:SparseEncoder.sparsity()callsto_sparse_csr(),which has no kernel on
SparseMPS. Wrapped intry/except NotImplementedErrorwith a CPU fallback. The result is a Python float so it's bit-identical, and the fallback covers any other backend missing the kernel too.sentence_transformer/fit_mixin.pyandcross_encoder/fit_mixin.py:fit(use_amp=True)was hardcoded totorch.cuda.amp.GradScaler()(and tofrom torch.cuda.amp import autocastin the SentenceTransformer side). Now both files dispatch ondevice.type: cuda and npu keep their existing namespaces, everything else (mps, xpu, cpu) goes throughtorch.amp.*.Testing
Ran the small examples on Apple Silicon end to end against this branch, without
PYTORCH_ENABLE_MPS_FALLBACK. All passed:naver/splade-cocondenser-ensembledistil)computing_embeddings.pyclustering/agglomerative.py,clustering/kmeans.pysemantic-search/semantic_search.pycross_encoder/applications/cross_encoder_usage.pyI also tried the heavier retrieval examples (Quora 100k with pytorch top-k, HNSWlib, Annoy). Those work too.