Patch: v0.20.3

* add possibility to pass nonetype objects

* adds nice test

.github/workflows/build_documentation.yml

@@ -15,3 +15,4 @@ jobs:
       package: accelerate
     secrets:
       token: ${{ secrets.HUGGINGFACE_PUSH }}
+      hf_token: ${{ secrets.HF_DOC_BUILD_PUSH }}

.github/workflows/delete_doc_comment.yml

@@ -1,13 +1,14 @@
-name: Delete dev documentation
+name: Delete doc comment
 
 on:
-  pull_request:
-    types: [ closed ]
+  workflow_run:
+    workflows: ["Delete doc comment trigger"]
+    types:
+      - completed
 
 
 jobs:
   delete:
     uses: huggingface/doc-builder/.github/workflows/delete_doc_comment.yml@main
-    with:
-      pr_number: ${{ github.event.number }}
-      package: accelerate
+    secrets:
+      comment_bot_token: ${{ secrets.COMMENT_BOT_TOKEN }}

.github/workflows/delete_doc_comment_trigger.yml

@@ -0,0 +1,12 @@
+name: Delete doc comment trigger
+
+on:
+  pull_request:
+    types: [ closed ]
+
+
+jobs:
+  delete:
+    uses: huggingface/doc-builder/.github/workflows/delete_doc_comment_trigger.yml@main
+    with:
+      pr_number: ${{ github.event.number }}

.github/workflows/upload_pr_documentation.yml

@@ -0,0 +1,16 @@
+name: Upload PR Documentation
+
+on:
+  workflow_run:
+    workflows: ["Build PR Documentation"]
+    types:
+      - completed
+
+jobs:
+  build:
+    uses: huggingface/doc-builder/.github/workflows/upload_pr_documentation.yml@main
+    with:
+      package_name: accelerate
+    secrets:
+      hf_token: ${{ secrets.HF_DOC_BUILD_PUSH }}
+      comment_bot_token: ${{ secrets.COMMENT_BOT_TOKEN }}

setup.py

@@ -32,7 +32,7 @@ extras["sagemaker"] = [
 
 setup(
     name="accelerate",
-    version="0.20.0.dev0",
+    version="0.20.3",
     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.3"
 
 from .accelerator import Accelerator
 from .big_modeling import (

src/accelerate/accelerator.py

@@ -1197,6 +1197,10 @@ class Accelerator:
         if self.distributed_type == DistributedType.FSDP and model_count == 1 and optimizer_present:
             result = self._prepare_fsdp(*result)
 
+        for item in result:
+            if item is not None:
+                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 +1695,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, "_is_accelerate_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 +1738,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, "_is_accelerate_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 +1770,11 @@ class Accelerator:
         >>> scheduler = accelerator.prepare_scheduler(scheduler)
         ```
         """
+        # Ensure we can't double wrap a scheduler due to `find_batch_size`
+        if getattr(scheduler, "_is_accelerate_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 +2565,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 +2584,7 @@ class Accelerator:
         self._models = []
         self._dataloaders = []
         self.deepspeed_engine_wrapped = None
+        self.step = 0
         release_memory()
 
     def clear(self):

src/accelerate/data_loader.py

@@ -316,6 +316,10 @@ class DataLoaderStateMixin:
         cls.end_of_dataloader = False
         cls.remainder = -1
 
+    def reset(self):
+        self.end_of_dataloader = False
+        self.remainder = -1
+
 
 class DataLoaderShard(DataLoader, DataLoaderStateMixin):
     """
@@ -361,6 +365,7 @@ class DataLoaderShard(DataLoader, DataLoaderStateMixin):
     def __iter__(self):
         if self.rng_types is not None:
             synchronize_rng_states(self.rng_types, self.synchronized_generator)
+        self.reset()
         self.gradient_state._add_dataloader(self)
         # We can safely pass because the default is -1
         with suppress(Exception):

src/accelerate/state.py

@@ -903,7 +903,9 @@ class GradientState:
             self.sync_gradients = True
             self.active_dataloader = None
             self.dataloader_references = [None]
-            self.plugin_kwargs = gradient_accumulation_plugin.to_kwargs()
+            self.plugin_kwargs = (
+                gradient_accumulation_plugin.to_kwargs() if gradient_accumulation_plugin is not None else {}
+            )
 
         # Plugin args are different and can be updated
         if gradient_accumulation_plugin is not None and self.plugin_kwargs != gradient_accumulation_plugin.to_kwargs():

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")

tests/test_accelerator.py

@@ -175,6 +175,18 @@ class AcceleratorTester(AccelerateTestCase):
             # mode.class_name is NOT loaded from config
             self.assertTrue(model.class_name != model.__class__.__name__)
 
+    def test_accelerator_none(self):
+        """Just test that passing None to accelerator.prepare() works."""
+        accelerator = Accelerator()
+        model, optimizer, scheduler, train_dl, valid_dl = create_components()
+        dummy_obj = None
+
+        # This should work
+        model, optimizer, scheduler, train_dl, valid_dl, dummy_obj = accelerator.prepare(
+            model, optimizer, scheduler, train_dl, valid_dl, dummy_obj
+        )
+        self.assertTrue(dummy_obj is None)
+
     @slow
     def test_accelerator_bnb(self):
         """Tests that the accelerator can be used with the BNB library."""

tests/test_data_loader.py

@@ -19,6 +19,7 @@ from torch.utils.data import BatchSampler, DataLoader, IterableDataset
 
 from accelerate.data_loader import (
     BatchSamplerShard,
+    DataLoaderShard,
     IterableDatasetShard,
     SkipBatchSampler,
     SkipDataLoader,
@@ -374,3 +375,12 @@ class DataLoaderTester(unittest.TestCase):
         dataloader = DataLoader(list(range(16)), batch_size=4)
         new_dataloader = skip_first_batches(dataloader, num_batches=2)
         self.assertListEqual([t.tolist() for t in new_dataloader], [[8, 9, 10, 11], [12, 13, 14, 15]])
+
+    def test_end_of_dataloader(self):
+        dataloader = DataLoaderShard(list(range(16)), batch_size=4)
+        for idx, _ in enumerate(dataloader):
+            self.assertEqual(dataloader.end_of_dataloader, idx == 3)
+
+        # Test it also works on the second iteration
+        for idx, _ in enumerate(dataloader):
+            self.assertEqual(dataloader.end_of_dataloader, idx == 3)