Merge branch 'huggingface:main' into stateful-dataloader

setup.py

@@ -48,7 +48,7 @@
 
 setup(
     name="accelerate",
-    version="0.32.0.dev0",
+    version="0.33.0.dev0",
     description="Accelerate",
     long_description=open("README.md", encoding="utf-8").read(),
     long_description_content_type="text/markdown",

src/accelerate/__init__.py

@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-__version__ = "0.32.0.dev0"
+__version__ = "0.33.0.dev0"
 
 from .accelerator import Accelerator
 from .big_modeling import (

src/accelerate/utils/memory.py

@@ -26,8 +26,13 @@
 from .imports import is_mlu_available, is_mps_available, is_npu_available, is_xpu_available
 
 
-def clear_device_cache():
-    gc.collect()
+def clear_device_cache(garbage_collection=False):
+    """
+    Clears the device cache by calling `torch.{backend}.empty_cache`. Can also run `gc.collect()`, but do note that
+    this is a *considerable* slowdown and should be used sparingly.
+    """
+    if garbage_collection:
+        gc.collect()
 
     if is_xpu_available():
         torch.xpu.empty_cache()
@@ -67,7 +72,7 @@ def release_memory(*objects):
         objects = list(objects)
     for i in range(len(objects)):
         objects[i] = None
-    clear_device_cache()
+    clear_device_cache(garbage_collection=True)
     return objects
 
 
@@ -123,7 +128,7 @@ def find_executable_batch_size(function: callable = None, starting_batch_size: i
 
     def decorator(*args, **kwargs):
         nonlocal batch_size
-        clear_device_cache()
+        clear_device_cache(garbage_collection=True)
         params = list(inspect.signature(function).parameters.keys())
         # Guard against user error
         if len(params) < (len(args) + 1):
@@ -139,7 +144,7 @@ def decorator(*args, **kwargs):
                 return function(batch_size, *args, **kwargs)
             except Exception as e:
                 if should_reduce_batch_size(e):
-                    clear_device_cache()
+                    clear_device_cache(garbage_collection=True)
                     batch_size //= 2
                 else:
                     raise

tests/test_modeling_utils.py

@@ -748,10 +748,7 @@ def test_load_state_dict(self):
 
             for param, device in device_map.items():
                 device = device if device != "disk" else "cpu"
-                expected_device = (
-                    torch.device(f"{torch_device}:{device}") if isinstance(device, int) else torch.device(device)
-                )
-                assert loaded_state_dict[param].device == expected_device
+                assert loaded_state_dict[param].device == torch.device(device)
 
     def test_convert_file_size(self):
         result = convert_file_size_to_int("0MB")