Skip to content

[Dinov2] Enable device_map="auto" support #38487

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
SunMarc merged 8 commits into from
Jun 4, 2025
Merged

[Dinov2] Enable device_map="auto" support #38487

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
669f526
Fix: resolve import order and duplicate import (ruff I001, F811)
aryanchauhan31 May 30, 2025
2a0605e
Format: clean up Dinov2 test file with ruff formatter
aryanchauhan31 May 30, 2025
2b67df1
Add _no_split_modules = ['Dinov2Layer'] to enable device_map='auto'
aryanchauhan31 May 30, 2025
3346f36
Revert dinov2_with_registers _no_split_modules to original state
aryanchauhan31 May 30, 2025
1dfb058
Remove redundant device_map test as suggested
aryanchauhan31 Jun 3, 2025
ef13720
Remove unused import after deleting test
aryanchauhan31 Jun 3, 2025
a98f7be
removed import torch and the redundant test function
aryanchauhan31 Jun 3, 2025
a91ba87
Update tests/models/dinov2/test_modeling_dinov2.py
SunMarc Jun 4, 2025
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
2 changes: 1 addition & 1 deletion src/transformers/models/dinov2/modeling_dinov2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -491,7 +491,7 @@ class Dinov2PreTrainedModel(PreTrainedModel):
base_model_prefix = "dinov2"
main_input_name = "pixel_values"
supports_gradient_checkpointing = True
_no_split_modules = ["Dinov2SwiGLUFFN"]
_no_split_modules = ["Dinov2Layer"]
_supports_sdpa = True
_supports_flash_attn_2 = True

Expand Down
2 changes: 1 addition & 1 deletion src/transformers/models/dinov2_with_registers/modeling_dinov2_with_registers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -509,7 +509,7 @@ class Dinov2WithRegistersPreTrainedModel(PreTrainedModel):
base_model_prefix = "dinov2_with_registers"
main_input_name = "pixel_values"
supports_gradient_checkpointing = True
_no_split_modules = ["Dinov2WithRegistersSwiGLUFFN"]
_no_split_modules = ["Dinov2WithRegistersLayer"]
_supports_sdpa = True
_supports_flash_attn_2 = True

Expand Down