From b9c17c5dc021ed17484adfdf2ae5614984adfc6e Mon Sep 17 00:00:00 2001 From: Aryan Chauhan Date: Wed, 4 Jun 2025 21:12:40 +0530 Subject: [PATCH] [Dinov2] Enable device_map="auto" support (#38487) * 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> --- src/transformers/models/dinov2/modeling_dinov2.py | 2 +- .../dinov2_with_registers/modeling_dinov2_with_registers.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/transformers/models/dinov2/modeling_dinov2.py b/src/transformers/models/dinov2/modeling_dinov2.py index 30bf56d3524..69911e40bc3 100644 --- a/src/transformers/models/dinov2/modeling_dinov2.py +++ b/src/transformers/models/dinov2/modeling_dinov2.py @@ -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 _supports_flex_attn = True diff --git a/src/transformers/models/dinov2_with_registers/modeling_dinov2_with_registers.py b/src/transformers/models/dinov2_with_registers/modeling_dinov2_with_registers.py index 82d8169c38d..db2d2b01d13 100644 --- a/src/transformers/models/dinov2_with_registers/modeling_dinov2_with_registers.py +++ b/src/transformers/models/dinov2_with_registers/modeling_dinov2_with_registers.py @@ -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 _supports_flex_attn = True