Skip to content
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.

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

HFShim: support MPS device #328

Merged
merged 6 commits into from
Aug 1, 2022
Merged

Conversation

danieldk
Copy link
Contributor

@danieldk danieldk commented Jun 1, 2022

Before this change, two devices and map locations were supported:

  • CUDA: cuda:N
  • CPU: cpu

This change adds support for other devices like Metal Performance Shader (MPS) devices by mapping to the active Torch device.

Before this change, two devices and map locations were supported:

- CUDA: `cuda:N`
- CPU: `cpu`

This change adds support for Metal Performance Shader (MPS) devices.  In
the case of MPS, the map location for loading a model is `cpu`, so we
use:

- CUDA: device: `cuda:N`, map location: `cuda:N`
- MPS: device: `mps`, map location: `cpu`
- CPU: device: `cpu`, map location: `cpu`
@danieldk
Copy link
Contributor Author

I think we also need a dev version of spaCy to get the right Thinc version for this PR.

@adrianeboyd
Copy link
Contributor

Can this be done soon after the v3.4.0 release instead? It's not going to break anything in spacy-transformers or the trf models, right?

@danieldk
Copy link
Contributor Author

Can this be done soon after the v3.4.0 release instead? It's not going to break anything in spacy-transformers or the trf models, right?

That would work. It shouldn't break anything, just moving models to the right device in a more generic manner.

@danieldk danieldk marked this pull request as ready for review July 22, 2022 07:30
@danieldk
Copy link
Contributor Author

Set the minimum spaCy version to 3.4.0 and merged in master. This is now ready for review.

@adrianeboyd
Copy link
Contributor

@explosion-bot please test_gpu

@explosion-bot
Copy link
Collaborator

explosion-bot commented Jul 27, 2022

🪁 Successfully triggered build on Buildkite

URL: https://buildkite.com/explosion-ai/spacy-transformers-gpu-test-suite/builds/19

@danieldk danieldk requested a review from svlandeg August 1, 2022 07:30
@svlandeg svlandeg merged commit 848c54c into explosion:master Aug 1, 2022
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

4 participants