[Dinov2] Enable device_map="auto" support#38487
Conversation
|
Hi @amyeroberts 👋 I've added Would really appreciate it if you could take a look when you get a chance. |
|
Hi @aryanchauhan31, I don't see the |
|
Hi @Rocketknight1 , Thanks for the heads up. |
|
Thanks! cc @SunMarc since I think this falls under |
SunMarc
left a comment
There was a problem hiding this comment.
Thanks ! This should indeed solve the issue. Just a nit
|
|
||
|
|
||
| @require_torch_multi_gpu | ||
| @slow | ||
| class Dinov2ModelDeviceMapTest(unittest.TestCase): | ||
| def test_model_parallelism(self): | ||
| model = Dinov2Model.from_pretrained("facebook/dinov2-base", device_map="auto") | ||
| self.assertTrue(hasattr(model, "hf_device_map")) | ||
| self.assertIsInstance(model.hf_device_map, dict) | ||
|
|
||
| dummy_input = torch.randn(1, 3, 224, 224).to(model.device) | ||
| outputs = model(pixel_values=dummy_input) | ||
| self.assertIsNotNone(outputs.last_hidden_state) |
There was a problem hiding this comment.
we already have tests for device_map called test_model_parallelism in the common tests we defined. We can remove this
There was a problem hiding this comment.
Hi @SunMarc, I'll go ahead and remove the redundant test function as suggested.
Thanks a lot @Rocketknight1 for pointing out the earlier issues and helping me clean things up! 🙌
| import torch | ||
|
|
||
| from transformers import Dinov2Config, Dinov2Model |
There was a problem hiding this comment.
Hi @SunMarc — just a heads-up: Dinov2Config, Dinov2Model, and torch are still being used in the other test functions in this file, so I’ve kept those imports.
Let me know if there's anything else you'd like cleaned up.
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
* Fix: resolve import order and duplicate import (ruff I001, F811) * Format: clean up Dinov2 test file with ruff formatter * Add _no_split_modules = ['Dinov2Layer'] to enable device_map='auto' * Revert dinov2_with_registers _no_split_modules to original state * Remove redundant device_map test as suggested * Remove unused import after deleting test * removed import torch and the redundant test function * Update tests/models/dinov2/test_modeling_dinov2.py --------- Co-authored-by: Marc Sun <57196510+SunMarc@users.noreply.github.com>
This PR adds support for
device_map="auto"to the Dinov2 model by defining_no_split_modules = ["Dinov2Layer"], which enables inference across multiple devices using Accelerate and Transformers.✔️ Summary
_no_split_modules = ["Dinov2Layer"]toDinov2PreTrainedModeltest_model_parallelismfor multi-GPUdevice_map="auto"behavior using a dummy input🔬 Why this matters
Large models like Dinov2 can now be used efficiently on limited memory setups using Transformers'
device_mapfeature.Closes #29786