-
Notifications
You must be signed in to change notification settings - Fork 399
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
Allow users to choose devices for models #264
Conversation
This addresses a problem that pops up when you've got several GPUs in play. When you set the device-map to 'auto', the model loads onto multiple GPUs just fine. But then, when you try to do 'model.to(device)', it tries to cram everything onto just one GPU. |
I've updated this branch with some formatting fixes and test/docstring changes from @rlouf. |
@@ -114,7 +114,7 @@ def test_transformers_integration_choice(): | |||
|
|||
def test_transformers_integration_with_pad_token(): | |||
model_name = "hf-internal-testing/tiny-random-XLMRobertaXLForCausalLM" | |||
model = models.transformers(model_name, device="cpu") | |||
model = models.transformers(model_name, device="meta") |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This was necessary in order to get around some "meta device" issues related to accelerate
.
Thank you @BramVanroy ! |
Currently, there are some issues when loading a transformers model because it explicitly calls
.to(device)
. The default is "cpu" (which is already the default in transformers to begin with). So I recommend to removedevice
as an argument completely and let the user decide how they want to load the model. That allows things like device_map and loading inkbit
. Internally inside theTransformers
class nothing changes, but we setself.device
tomodel.device
to ensure that everything stays compatible in the rest of the library.closes #238