Patch: v0.20.1

* Add step reset to free memory

* Check if not Accelerated Optimizer

* Continue

* Another try

* Check the rest

* Try with just check on init

* Change logic based on review

* Update

* Oops very big logic issue!

setup.py

@@ -32,7 +32,7 @@ extras["sagemaker"] = [
 
 setup(
     name="accelerate",
-    version="0.20.0.dev0",
+    version="0.20.1",
     description="Accelerate",
     long_description=open("README.md", "r", encoding="utf-8").read(),
     long_description_content_type="text/markdown",

src/accelerate/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.20.0.dev0"
+__version__ = "0.20.1"
 
 from .accelerator import Accelerator
 from .big_modeling import (

src/accelerate/accelerator.py

@@ -1197,6 +1197,9 @@ class Accelerator:
         if self.distributed_type == DistributedType.FSDP and model_count == 1 and optimizer_present:
             result = self._prepare_fsdp(*result)
 
+        for item in result:
+            setattr(item, "_is_accelerate_prepared", True)
+
         return result if len(result) > 1 else result[0]
 
     def prepare_model(self, model: torch.nn.Module, device_placement: bool = None, evaluation_mode: bool = False):
@@ -1691,6 +1694,11 @@ class Accelerator:
         >>> data_loader = accelerator.prepare_data_loader(data_loader, device_placement=True)
         ```
         """
+        # Ensure we can't double wrap a DataLoader due to `find_batch_size`
+        if getattr(data_loader, "_accelerator_prepared", False):
+            if data_loader not in self._dataloaders:
+                self._dataloaders.append(data_loader)
+            return data_loader
         if device_placement is None:
             device_placement = self.device_placement if self.distributed_type != DistributedType.TPU else False
         prepared_data_loader = prepare_data_loader(
@@ -1729,6 +1737,11 @@ class Accelerator:
         >>> optimizer = accelerator.prepare_optimizer(optimizer, device_placement=True)
         ```
         """
+        # Ensure we can't double wrap an optimizer due to `find_batch_size`
+        if getattr(optimizer, "_accelerator_prepared", False):
+            if optimizer not in self._optimizers:
+                self._optimizers.append(optimizer)
+            return optimizer
         if device_placement is None:
             device_placement = self.device_placement
         optimizer = AcceleratedOptimizer(optimizer, device_placement=device_placement, scaler=self.scaler)
@@ -1756,6 +1769,11 @@ class Accelerator:
         >>> scheduler = accelerator.prepare_scheduler(scheduler)
         ```
         """
+        # Ensure we can't double wrap a scheduler due to `find_batch_size`
+        if getattr(scheduler, "_accelerator_prepared", False):
+            if scheduler not in self._schedulers:
+                self._schedulers.append(scheduler)
+            return scheduler
         # We try to find the optimizer associated with `scheduler`, the default is the full list.
         optimizer = self._optimizers
         for opt in self._optimizers:
@@ -2546,7 +2564,7 @@ class Accelerator:
     def free_memory(self):
         """
         Will release all references to the internal objects stored and call the garbage collector. You should call this
-        method between two trainings with different models/optimizers.
+        method between two trainings with different models/optimizers. Also will reset `Accelerator.step` to 0.
 
         Example:
 
@@ -2565,6 +2583,7 @@ class Accelerator:
         self._models = []
         self._dataloaders = []
         self.deepspeed_engine_wrapped = None
+        self.step = 0
         release_memory()
 
     def clear(self):

src/accelerate/utils/modeling.py

@@ -941,12 +941,11 @@ def load_state_dict(checkpoint_file, device_map=None):
         if device_map is None:
             return safe_load_file(checkpoint_file)
         else:
-            devices = list(set(device_map.values()) - {"disk"})
-
             # if we only have one device we can load everything directly
-            if len(devices) == 1:
-                return safe_load_file(checkpoint_file, device=devices[0])
+            if len(set(device_map.values())) == 1:
+                return safe_load_file(checkpoint_file, device=list(device_map.values())[0])
 
+            devices = list(set(device_map.values()) - {"disk"})
             # cpu device should always exist as fallback option
             if "cpu" not in devices:
                 devices.append("cpu")