diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 63ea4a71b..e2248fcfa 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -25,7 +25,7 @@
 import torch
 
 from .checkpointing import load_accelerator_state, load_custom_state, save_accelerator_state, save_custom_state
-from .data_loader import prepare_data_loader
+from .data_loader import DataLoaderDispatcher, prepare_data_loader
 from .logging import get_logger
 from .optimizer import AcceleratedOptimizer
 from .scheduler import AcceleratedScheduler
@@ -87,6 +87,9 @@
         megatron_lm_prepare_scheduler,
     )
 
+if is_torch_version(">", "1.10.0"):
+    from torch.distributed.algorithms.join import Join
+
 
 if is_tpu_available(check_device=False):
     import torch_xla.distributed.xla_multiprocessing as xmp
@@ -353,6 +356,7 @@ def __init__(
         self._optimizers = []
         self._models = []
         self._schedulers = []
+        self._dataloaders = []
         self._custom_objects = []
 
         # RNG Types
@@ -608,6 +612,93 @@ def accumulate(self, model):
         with context(model):
             yield
 
+    @contextmanager
+    def join_uneven_inputs(self, joinables, even_batches=None):
+        """
+        A context manager that facilitates distributed training or evaluation on uneven inputs, which acts as a wrapper
+        around `torch.distributed.algorithms.join`. This is useful when the total batch size does not evenly divide the
+        length of the dataset.
+
+        Args:
+            joinables (`List[torch.distributed.algorithms.Joinable]`):
+                A list of models or optimizers that subclass `torch.distributed.algorithms.Joinable`. Most commonly, a
+                PyTorch Module that was prepared with `Accelerator.prepare` for DistributedDataParallel training.
+            even_batches (`bool`, *optional*)
+                If set, this will override the value of `even_batches` set in the `Accelerator`. If it is not provided,
+                the default `Accelerator` value wil be used.
+
+        <Tip warning={true}>
+
+        `join_uneven_inputs` is only supported for Distributed Data Parallel training on multiple GPUs. For any other
+        configuration, this method will have no effect.
+
+        </Tip>
+
+        <Tip warning={true}>
+
+        Overidding `even_batches` will not affect iterable-style data loaders.
+
+        </Tip>
+
+        Example:
+
+        ```python
+        >>> from accelerate import Accelerator
+
+        >>> accelerator = Accelerator(even_batches=True)
+        >>> ddp_model, optimizer, dataloader = accelerator.prepare(model, optimizer, dataloader)
+
+        >>> with accelerator.join_uneven_inputs([ddp_model], even_batches=False):
+        ...     for input, output in dataloader:
+        ...         outputs = model(input)
+        ...         loss = loss_func(outputs)
+        ...         loss.backward()
+        ...         optimizer.step()
+        ...         optimizer.zero_grad()
+        ```
+
+        """
+        if is_torch_version("<", "1.10.0"):
+            raise ValueError(f"Joining uneven inputs requires PyTorch >= 1.10.0, You have {torch.__version__}.")
+
+        if self.distributed_type == DistributedType.MULTI_GPU:
+            dl_even_batches_values = []
+
+            if even_batches is not None:
+                iterable_dl_seen = False
+                # override value in batch sampler for map-style datasets
+                for dl_idx, dl in enumerate(self._dataloaders):
+                    if isinstance(dl, DataLoaderDispatcher):
+                        iterable_dl_seen = True
+                        continue
+                    dl_even_batches_values.append((dl_idx, dl.batch_sampler.even_batches))
+                    dl.batch_sampler.even_batches = even_batches
+
+                if iterable_dl_seen:
+                    warnings.warn(
+                        "Overridding even_batches is only supported for map-style datasets, yet some dataloaders given were iterable"
+                    )
+            else:
+                even_batches = self.even_batches
+
+            enable_join = False if even_batches else True
+            try:
+                with Join(joinables, enable=enable_join, throw_on_early_termination=False):
+                    yield
+            finally:
+                # reset any batch samplers that have been modified
+                for dl_idx, even_batches_value in dl_even_batches_values:
+                    self._dataloaders[dl_idx].batch_sampler.even_batches = even_batches_value
+        else:
+            # Even when disabled, Join expects models to subclass Joinable, so skip entirely for single process runs
+            if self.distributed_type != DistributedType.NO:
+                warnings.warn(
+                    "Joining uneven inputs is only supported for multi-GPU training, as a result `join_uneven_inputs` will have no effect."
+                )
+
+            with contextlib.nullcontext(joinables):
+                yield
+
     def print(self, *args, **kwargs):
         """
         Use in replacement of `print()` to only print once per server.
@@ -1117,7 +1208,7 @@ def prepare_data_loader(self, data_loader: torch.utils.data.DataLoader, device_p
         """
         if device_placement is None:
             device_placement = self.device_placement if self.distributed_type != DistributedType.TPU else False
-        return prepare_data_loader(
+        prepared_data_loader = prepare_data_loader(
             data_loader,
             self.device,
             num_processes=self.num_processes,
@@ -1128,6 +1219,8 @@ def prepare_data_loader(self, data_loader: torch.utils.data.DataLoader, device_p
             dispatch_batches=self.dispatch_batches,
             even_batches=self.even_batches,
         )
+        self._dataloaders.append(prepared_data_loader)
+        return prepared_data_loader
 
     def prepare_optimizer(self, optimizer: torch.optim.Optimizer, device_placement=None):
         """
@@ -1611,6 +1704,7 @@ def free_memory(self):
         self._schedulers = []
         self._optimizers = []
         self._models = []
+        self._dataloaders = []
         self.deepspeed_engine_wrapped = None
         gc.collect()
         torch.cuda.empty_cache()
diff --git a/src/accelerate/test_utils/scripts/test_distributed_data_loop.py b/src/accelerate/test_utils/scripts/test_distributed_data_loop.py
index eaf7c6a34..6576e0358 100644
--- a/src/accelerate/test_utils/scripts/test_distributed_data_loop.py
+++ b/src/accelerate/test_utils/scripts/test_distributed_data_loop.py
@@ -14,12 +14,25 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+
+import warnings
 from typing import List
+from unittest.mock import Mock
 
 import torch
-from torch.utils.data import DataLoader, TensorDataset
+from torch.utils.data import DataLoader, IterableDataset, TensorDataset
 
 from accelerate.accelerator import Accelerator
+from accelerate.utils.dataclasses import DistributedType
+
+
+class DummyIterableDataset(IterableDataset):
+    def __init__(self, data):
+        self.data = data
+
+    def __iter__(self):
+        for element in self.data:
+            yield element
 
 
 def create_accelerator(even_batches=True):
@@ -28,11 +41,14 @@ def create_accelerator(even_batches=True):
     return accelerator
 
 
-def create_dataloader(accelerator: Accelerator, dataset_size: int, batch_size: int):
+def create_dataloader(accelerator: Accelerator, dataset_size: int, batch_size: int, iterable: bool = False):
     """
     Create a simple DataLoader to use during the test cases
     """
-    dataset = TensorDataset(torch.as_tensor(range(dataset_size)))
+    if iterable:
+        dataset = DummyIterableDataset(torch.as_tensor(range(dataset_size)))
+    else:
+        dataset = TensorDataset(torch.as_tensor(range(dataset_size)))
 
     dl = DataLoader(dataset, batch_size=batch_size)
     dl = accelerator.prepare(dl)
@@ -103,7 +119,96 @@ def test_can_disable_even_batches():
     )
 
 
-if __name__ == "__main__":
+def test_can_join_uneven_inputs():
+    accelerator = create_accelerator(even_batches=False)
+
+    model = torch.nn.Linear(1, 1)
+    ddp_model = accelerator.prepare(model)
+
+    dl = create_dataloader(accelerator, dataset_size=3, batch_size=1)
+
+    batch_idxs = []
+    with accelerator.join_uneven_inputs([ddp_model]):
+        for batch_idx, batch in enumerate(dl):
+            output = ddp_model(batch[0].float())
+            loss = output.sum()
+            loss.backward()
+            batch_idxs.append(batch_idx)
+
+    accelerator.wait_for_everyone()
+
+    if accelerator.process_index == 0:
+        assert batch_idxs == [0, 1]
+    elif accelerator.process_index == 1:
+        assert batch_idxs == [0]
+
+
+def test_join_raises_warning_for_non_ddp_distributed(accelerator):
+
+    with warnings.catch_warnings(record=True) as w:
+        with accelerator.join_uneven_inputs([Mock()]):
+            pass
+
+        assert issubclass(w[-1].category, UserWarning)
+        assert "only supported for multi-GPU" in str(w[-1].message)
+
+
+def test_join_can_override_even_batches():
+    default_even_batches = True
+    overridden_even_batches = False
+    accelerator = create_accelerator(even_batches=default_even_batches)
+    model = torch.nn.Linear(1, 1)
+    ddp_model = accelerator.prepare(model)
+    train_dl = create_dataloader(accelerator, dataset_size=3, batch_size=1)
+    valid_dl = create_dataloader(accelerator, dataset_size=3, batch_size=1)
+
+    with accelerator.join_uneven_inputs([ddp_model], even_batches=overridden_even_batches):
+        train_dl_overridden_value = train_dl.batch_sampler.even_batches
+        valid_dl_overridden_value = valid_dl.batch_sampler.even_batches
+
+    assert train_dl_overridden_value == overridden_even_batches
+    assert valid_dl_overridden_value == overridden_even_batches
+    assert train_dl.batch_sampler.even_batches == default_even_batches
+    assert valid_dl.batch_sampler.even_batches == default_even_batches
+
+
+def test_join_can_override_for_mixed_type_dataloaders():
+    default_even_batches = True
+    overridden_even_batches = False
+    accelerator = create_accelerator(even_batches=default_even_batches)
+    model = torch.nn.Linear(1, 1)
+    ddp_model = accelerator.prepare(model)
+    create_dataloader(accelerator, dataset_size=3, batch_size=1, iterable=True)
+    batch_dl = create_dataloader(accelerator, dataset_size=3, batch_size=1)
+
+    with warnings.catch_warnings():
+        warnings.filterwarnings("ignore")
+        try:
+            with accelerator.join_uneven_inputs([ddp_model], even_batches=overridden_even_batches):
+                batch_dl_overridden_value = batch_dl.batch_sampler.even_batches
+        except AttributeError:
+            # ensure attribute error is not raised when processing iterable dl
+            raise AssertionError
+
+    assert batch_dl_overridden_value == overridden_even_batches
+    assert batch_dl.batch_sampler.even_batches == default_even_batches
+
+
+def test_join_raises_warning_for_iterable_when_overriding_even_batches():
+    accelerator = create_accelerator()
+    model = torch.nn.Linear(1, 1)
+    ddp_model = accelerator.prepare(model)
+    create_dataloader(accelerator, dataset_size=3, batch_size=1, iterable=True)
+
+    with warnings.catch_warnings(record=True) as w:
+        with accelerator.join_uneven_inputs([ddp_model], even_batches=False):
+            pass
+
+        assert issubclass(w[-1].category, UserWarning)
+        assert "only supported for map-style datasets" in str(w[-1].message)
+
+
+def main():
     accelerator = create_accelerator()
 
     accelerator.print("Test that even_batches variable ensures uniform batches across processes")
@@ -111,3 +216,25 @@ def test_can_disable_even_batches():
 
     accelerator.print("Run tests with even_batches disabled")
     test_can_disable_even_batches()
+
+    accelerator.print("Test joining uneven inputs")
+    test_can_join_uneven_inputs()
+
+    accelerator.print("Test overriding even_batches when joining uneven inputs")
+    test_join_can_override_even_batches()
+
+    accelerator.print("Test overriding even_batches for mixed dataloader types")
+    test_join_can_override_for_mixed_type_dataloaders()
+
+    accelerator.print("Test overriding even_batches raises a warning for iterable dataloaders")
+    test_join_raises_warning_for_iterable_when_overriding_even_batches()
+
+    accelerator.print("Test join with non DDP distributed raises warning")
+    original_state = accelerator.state.distributed_type
+    accelerator.state.distributed_type = DistributedType.FSDP
+    test_join_raises_warning_for_non_ddp_distributed(accelerator)
+    accelerator.state.distributed_type = original_state
+
+
+if __name__ == "__main__":
+    main()
diff --git a/tests/test_accelerator.py b/tests/test_accelerator.py
new file mode 100644
index 000000000..19d6c1655
--- /dev/null
+++ b/tests/test_accelerator.py
@@ -0,0 +1,51 @@
+import unittest
+
+import torch
+from torch.utils.data import DataLoader, TensorDataset
+
+from accelerate.accelerator import Accelerator
+from accelerate.state import AcceleratorState
+
+
+def create_components():
+    model = torch.nn.Linear(2, 4)
+    optimizer = torch.optim.AdamW(model.parameters(), lr=1.0)
+    scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer, max_lr=0.01, steps_per_epoch=2, epochs=1)
+    train_dl = DataLoader(TensorDataset(torch.tensor([1, 2, 3])))
+    valid_dl = DataLoader(TensorDataset(torch.tensor([4, 5, 6])))
+
+    return model, optimizer, scheduler, train_dl, valid_dl
+
+
+class AcceleratorTester(unittest.TestCase):
+    def test_prepared_objects_are_referenced(self):
+        accelerator = Accelerator()
+        model, optimizer, scheduler, train_dl, valid_dl = create_components()
+
+        (
+            prepared_model,
+            prepared_optimizer,
+            prepared_scheduler,
+            prepared_train_dl,
+            prepared_valid_dl,
+        ) = accelerator.prepare(model, optimizer, scheduler, train_dl, valid_dl)
+
+        self.assertTrue(prepared_model in accelerator._models)
+        self.assertTrue(prepared_optimizer in accelerator._optimizers)
+        self.assertTrue(prepared_scheduler in accelerator._schedulers)
+        self.assertTrue(prepared_train_dl in accelerator._dataloaders)
+        self.assertTrue(prepared_valid_dl in accelerator._dataloaders)
+        AcceleratorState._reset_state()
+
+    def test_free_memory_dereferences_prepared_components(self):
+        accelerator = Accelerator()
+        model, optimizer, scheduler, train_dl, valid_dl = create_components()
+        accelerator.prepare(model, optimizer, scheduler, train_dl, valid_dl)
+
+        accelerator.free_memory()
+
+        self.assertTrue(len(accelerator._models) == 0)
+        self.assertTrue(len(accelerator._optimizers) == 0)
+        self.assertTrue(len(accelerator._schedulers) == 0)
+        self.assertTrue(len(accelerator._dataloaders) == 0)
+        AcceleratorState._reset_state()


Review comments: body: Maybe only throw an error if there are no dataloader with that attribute? I can imagine having an iterable dataset for training but a standard dataset for evaluation.
diff_hunk: @@ -608,6 +612,41 @@ def accumulate(self, model):
         with context(model):
             yield
 
+    @contextmanager
+    def join_uneven_inputs(self, joinables, even_batches=None):
+        """
+        A context manager that facilitates distributed training on uneven inputs, which acts as a wrapper around `torch.distributed.algorithms.join`.
+        """
+        if is_torch_version("<", "1.10.0"):
+            raise ValueError(f"Joining uneven inputs requires PyTorch >= 1.10.0, You have {torch.__version__}.")
+
+        if self.distributed_type == DistributedType.NO:
+            # Even when disabled, Join expects models to subclass Joinable, so skip entirely for single process runs
+            with contextlib.nullcontext(joinables):
+                yield
+
+        elif self.distributed_type == DistributedType.MULTI_GPU:
+            dl_even_batches_values = []
+
+            if even_batches is not None:
+                for dl in self._dataloaders:
+                    if not hasattr(dl, "even_batches"):
+                        raise ValueError("Overridding even_batches is not supported for iterable-style datasets")
body: I'd just issue a warning and exit.
diff_hunk: @@ -608,6 +612,41 @@ def accumulate(self, model):
         with context(model):
             yield
 
+    @contextmanager
+    def join_uneven_inputs(self, joinables, even_batches=None):
+        """
+        A context manager that facilitates distributed training on uneven inputs, which acts as a wrapper around `torch.distributed.algorithms.join`.
+        """
+        if is_torch_version("<", "1.10.0"):
+            raise ValueError(f"Joining uneven inputs requires PyTorch >= 1.10.0, You have {torch.__version__}.")
+
+        if self.distributed_type == DistributedType.NO:
+            # Even when disabled, Join expects models to subclass Joinable, so skip entirely for single process runs
+            with contextlib.nullcontext(joinables):
+                yield
+
+        elif self.distributed_type == DistributedType.MULTI_GPU:
+            dl_even_batches_values = []
+
+            if even_batches is not None:
+                for dl in self._dataloaders:
+                    if not hasattr(dl, "even_batches"):
+                        raise ValueError("Overridding even_batches is not supported for iterable-style datasets")
+                    dl_even_batches_values.append(dl.even_batches)
+                    dl.even_batches = even_batches
+            else:
+                even_batches = self.even_batches
+
+            enable_join = False if even_batches else True
+            try:
+                with Join(joinables, enable=enable_join, throw_on_early_termination=False):
+                    yield
+            finally:
+                for dl, even_batches_value in zip(self._dataloaders, dl_even_batches_values):
+                    dl.even_batches = even_batches_value
+        else:
+            raise ValueError("Joining uneven inputs is only supported for DistributedDataParallel training")
body: This test will need two processes, so probably needs some decorator/special launcher.
diff_hunk: @@ -103,11 +117,78 @@ def test_can_disable_even_batches():
     )
 
 
-if __name__ == "__main__":
+def test_can_join_uneven_inputs():
+    accelerator = create_accelerator(even_batches=False)
+
+    model = torch.nn.Linear(1, 1)
+    ddp_model = accelerator.prepare(model)
+
+    dl = create_dataloader(accelerator, dataset_size=3, batch_size=1)
+
+    batch_idxs = []
+    with accelerator.join_uneven_inputs([ddp_model]):
+        for batch_idx, batch in enumerate(dl):
+            output = ddp_model(batch[0].float())
+            loss = output.sum()
+            loss.backward()
+            batch_idxs.append(batch_idx)
+
+    accelerator.wait_for_everyone()
+
+    if accelerator.process_index == 0:
body: You are correct. I think that this should be taken care of here, from my previous PR: https://github.com/huggingface/accelerate/blob/main/tests/test_multigpu.py#L49
diff_hunk: @@ -103,11 +117,78 @@ def test_can_disable_even_batches():
     )
 
 
-if __name__ == "__main__":
+def test_can_join_uneven_inputs():
+    accelerator = create_accelerator(even_batches=False)
+
+    model = torch.nn.Linear(1, 1)
+    ddp_model = accelerator.prepare(model)
+
+    dl = create_dataloader(accelerator, dataset_size=3, batch_size=1)
+
+    batch_idxs = []
+    with accelerator.join_uneven_inputs([ddp_model]):
+        for batch_idx, batch in enumerate(dl):
+            output = ddp_model(batch[0].float())
+            loss = output.sum()
+            loss.backward()
+            batch_idxs.append(batch_idx)
+
+    accelerator.wait_for_everyone()
+
+    if accelerator.process_index == 0:
body: Good spot. Based on the other feedback, I have just raised a warning for this rather than an exception
diff_hunk: @@ -608,6 +612,41 @@ def accumulate(self, model):
         with context(model):
             yield
 
+    @contextmanager
+    def join_uneven_inputs(self, joinables, even_batches=None):
+        """
+        A context manager that facilitates distributed training on uneven inputs, which acts as a wrapper around `torch.distributed.algorithms.join`.
+        """
+        if is_torch_version("<", "1.10.0"):
+            raise ValueError(f"Joining uneven inputs requires PyTorch >= 1.10.0, You have {torch.__version__}.")
+
+        if self.distributed_type == DistributedType.NO:
+            # Even when disabled, Join expects models to subclass Joinable, so skip entirely for single process runs
+            with contextlib.nullcontext(joinables):
+                yield
+
+        elif self.distributed_type == DistributedType.MULTI_GPU:
+            dl_even_batches_values = []
+
+            if even_batches is not None:
+                for dl in self._dataloaders:
+                    if not hasattr(dl, "even_batches"):
+                        raise ValueError("Overridding even_batches is not supported for iterable-style datasets")
body: Done
diff_hunk: @@ -608,6 +612,41 @@ def accumulate(self, model):
         with context(model):
             yield
 
+    @contextmanager
+    def join_uneven_inputs(self, joinables, even_batches=None):
+        """
+        A context manager that facilitates distributed training on uneven inputs, which acts as a wrapper around `torch.distributed.algorithms.join`.
+        """
+        if is_torch_version("<", "1.10.0"):
+            raise ValueError(f"Joining uneven inputs requires PyTorch >= 1.10.0, You have {torch.__version__}.")
+
+        if self.distributed_type == DistributedType.NO:
+            # Even when disabled, Join expects models to subclass Joinable, so skip entirely for single process runs
+            with contextlib.nullcontext(joinables):
+                yield
+
+        elif self.distributed_type == DistributedType.MULTI_GPU:
+            dl_even_batches_values = []
+
+            if even_batches is not None:
+                for dl in self._dataloaders:
+                    if not hasattr(dl, "even_batches"):
+                        raise ValueError("Overridding even_batches is not supported for iterable-style datasets")
+                    dl_even_batches_values.append(dl.even_batches)
+                    dl.even_batches = even_batches
+            else:
+                even_batches = self.even_batches
+
+            enable_join = False if even_batches else True
+            try:
+                with Join(joinables, enable=enable_join, throw_on_early_termination=False):
+                    yield
+            finally:
+                for dl, even_batches_value in zip(self._dataloaders, dl_even_batches_values):
+                    dl.even_batches = even_batches_value
+        else:
+            raise ValueError("Joining uneven inputs is only supported for DistributedDataParallel training")
body: Let's add a warning here maybe? Longterm, `DataLoaderDispatcher` will also support the `even_batches` kwarg.
diff_hunk: @@ -608,6 +612,91 @@ def accumulate(self, model):
         with context(model):
             yield
 
+    @contextmanager
+    def join_uneven_inputs(self, joinables, even_batches=None):
+        """
+        A context manager that facilitates distributed training or evaluation on uneven inputs, which acts as a wrapper
+        around `torch.distributed.algorithms.join`. This is useful when the total batch size does not evenly divide the
+        length of the dataset.
+
+        Args:
+            joinables (`List[torch.distributed.algorithms.Joinable]`):
+                A list of models or optimizers that subclass `torch.distributed.algorithms.Joinable`. Most commonly, a
+                PyTorch Module that was prepared with `Accelerator.prepare` for DistributedDataParallel training.
+            even_batches (`bool`, *optional*)
+                If set, this will override the value of `even_batches` set in the `Accelerator`. If it is not provided,
+                the default `Accelerator` value wil be used.
+
+        <Tip warning={true}>
+
+        `join_uneven_inputs` is only supported for Distributed Data Parallel training on multiple GPUs. For any other
+        configuration, this method will have no effect.
+
+        </Tip>
+
+        <Tip warning={true}>
+
+        Overidding `even_batches` will not affect iterable-style data loaders.
+
+        </Tip>
+
+        Example:
+
+        ```python
+        >>> from accelerate import Accelerator
+
+        >>> accelerator = Accelerator(even_batches=True)
+        >>> ddp_model, optimizer, dataloader = accelerator.prepare(model, optimizer, dataloader)
+
+        >>> with accelerator.join_uneven_inputs([ddp_model], even_batches=False):
+        ...     for input, output in dataloader:
+        ...         outputs = model(input)
+        ...         loss = loss_func(outputs)
+        ...         loss.backward()
+        ...         optimizer.step()
+        ...         optimizer.zero_grad()
+        ```
+
+        """
+        if is_torch_version("<", "1.10.0"):
+            raise ValueError(f"Joining uneven inputs requires PyTorch >= 1.10.0, You have {torch.__version__}.")
+
+        if self.distributed_type == DistributedType.MULTI_GPU:
+            dl_even_batches_values = []
+
+            if even_batches is not None:
+                # override value in batch sampler for map-style datasets
+                for dl_idx, dl in enumerate(self._dataloaders):
+                    if isinstance(dl, DataLoaderDispatcher):
+                        continue
body: ```suggestion
                    "Joining uneven inputs is only supported for multi-GPU training, as a result `join_uneven_inputs` will have no effect."
```
diff_hunk: @@ -608,6 +612,91 @@ def accumulate(self, model):
         with context(model):
             yield
 
+    @contextmanager
+    def join_uneven_inputs(self, joinables, even_batches=None):
+        """
+        A context manager that facilitates distributed training or evaluation on uneven inputs, which acts as a wrapper
+        around `torch.distributed.algorithms.join`. This is useful when the total batch size does not evenly divide the
+        length of the dataset.
+
+        Args:
+            joinables (`List[torch.distributed.algorithms.Joinable]`):
+                A list of models or optimizers that subclass `torch.distributed.algorithms.Joinable`. Most commonly, a
+                PyTorch Module that was prepared with `Accelerator.prepare` for DistributedDataParallel training.
+            even_batches (`bool`, *optional*)
+                If set, this will override the value of `even_batches` set in the `Accelerator`. If it is not provided,
+                the default `Accelerator` value wil be used.
+
+        <Tip warning={true}>
+
+        `join_uneven_inputs` is only supported for Distributed Data Parallel training on multiple GPUs. For any other
+        configuration, this method will have no effect.
+
+        </Tip>
+
+        <Tip warning={true}>
+
+        Overidding `even_batches` will not affect iterable-style data loaders.
+
+        </Tip>
+
+        Example:
+
+        ```python
+        >>> from accelerate import Accelerator
+
+        >>> accelerator = Accelerator(even_batches=True)
+        >>> ddp_model, optimizer, dataloader = accelerator.prepare(model, optimizer, dataloader)
+
+        >>> with accelerator.join_uneven_inputs([ddp_model], even_batches=False):
+        ...     for input, output in dataloader:
+        ...         outputs = model(input)
+        ...         loss = loss_func(outputs)
+        ...         loss.backward()
+        ...         optimizer.step()
+        ...         optimizer.zero_grad()
+        ```
+
+        """
+        if is_torch_version("<", "1.10.0"):
+            raise ValueError(f"Joining uneven inputs requires PyTorch >= 1.10.0, You have {torch.__version__}.")
+
+        if self.distributed_type == DistributedType.MULTI_GPU:
+            dl_even_batches_values = []
+
+            if even_batches is not None:
+                # override value in batch sampler for map-style datasets
+                for dl_idx, dl in enumerate(self._dataloaders):
+                    if isinstance(dl, DataLoaderDispatcher):
+                        continue
+                    dl_even_batches_values.append((dl_idx, dl.batch_sampler.even_batches))
+                    dl.batch_sampler.even_batches = even_batches
+
+                if len(dl_even_batches_values) == 0:
+                    warnings.warn(
+                        "Overridding even_batches is only supported for map-style datasets, yet all dataloaders given were iterable"
+                    )
+            else:
+                even_batches = self.even_batches
+
+            enable_join = False if even_batches else True
+            try:
+                with Join(joinables, enable=enable_join, throw_on_early_termination=False):
+                    yield
+            finally:
+                # reset any batch samplers that have been modified
+                for dl_idx, even_batches_value in dl_even_batches_values:
+                    self._dataloaders[dl_idx].batch_sampler.even_batches = even_batches_value
+        else:
+            # Even when disabled, Join expects models to subclass Joinable, so skip entirely for single process runs
+            if self.distributed_type != DistributedType.NO:
+                warnings.warn(
+                    "Joining uneven inputs is only supported for DistributedDataParallel training, join_unenven_inputs has no effect."
body: Sure, I could change the condition on the warning that occurs if only `DataLoaderDispacher`s are provided, so that it warns if at least `DataLoaderDispacher` is seen. 
diff_hunk: @@ -608,6 +612,91 @@ def accumulate(self, model):
         with context(model):
             yield
 
+    @contextmanager
+    def join_uneven_inputs(self, joinables, even_batches=None):
+        """
+        A context manager that facilitates distributed training or evaluation on uneven inputs, which acts as a wrapper
+        around `torch.distributed.algorithms.join`. This is useful when the total batch size does not evenly divide the
+        length of the dataset.
+
+        Args:
+            joinables (`List[torch.distributed.algorithms.Joinable]`):
+                A list of models or optimizers that subclass `torch.distributed.algorithms.Joinable`. Most commonly, a
+                PyTorch Module that was prepared with `Accelerator.prepare` for DistributedDataParallel training.
+            even_batches (`bool`, *optional*)
+                If set, this will override the value of `even_batches` set in the `Accelerator`. If it is not provided,
+                the default `Accelerator` value wil be used.
+
+        <Tip warning={true}>
+
+        `join_uneven_inputs` is only supported for Distributed Data Parallel training on multiple GPUs. For any other
+        configuration, this method will have no effect.
+
+        </Tip>
+
+        <Tip warning={true}>
+
+        Overidding `even_batches` will not affect iterable-style data loaders.
+
+        </Tip>
+
+        Example:
+
+        ```python
+        >>> from accelerate import Accelerator
+
+        >>> accelerator = Accelerator(even_batches=True)
+        >>> ddp_model, optimizer, dataloader = accelerator.prepare(model, optimizer, dataloader)
+
+        >>> with accelerator.join_uneven_inputs([ddp_model], even_batches=False):
+        ...     for input, output in dataloader:
+        ...         outputs = model(input)
+        ...         loss = loss_func(outputs)
+        ...         loss.backward()
+        ...         optimizer.step()
+        ...         optimizer.zero_grad()
+        ```
+
+        """
+        if is_torch_version("<", "1.10.0"):
+            raise ValueError(f"Joining uneven inputs requires PyTorch >= 1.10.0, You have {torch.__version__}.")
+
+        if self.distributed_type == DistributedType.MULTI_GPU:
+            dl_even_batches_values = []
+
+            if even_batches is not None:
+                # override value in batch sampler for map-style datasets
+                for dl_idx, dl in enumerate(self._dataloaders):
+                    if isinstance(dl, DataLoaderDispatcher):
+                        continue
body: I have made the change now
diff_hunk: @@ -608,6 +612,91 @@ def accumulate(self, model):
         with context(model):
             yield
 
+    @contextmanager
+    def join_uneven_inputs(self, joinables, even_batches=None):
+        """
+        A context manager that facilitates distributed training or evaluation on uneven inputs, which acts as a wrapper
+        around `torch.distributed.algorithms.join`. This is useful when the total batch size does not evenly divide the
+        length of the dataset.
+
+        Args:
+            joinables (`List[torch.distributed.algorithms.Joinable]`):
+                A list of models or optimizers that subclass `torch.distributed.algorithms.Joinable`. Most commonly, a
+                PyTorch Module that was prepared with `Accelerator.prepare` for DistributedDataParallel training.
+            even_batches (`bool`, *optional*)
+                If set, this will override the value of `even_batches` set in the `Accelerator`. If it is not provided,
+                the default `Accelerator` value wil be used.
+
+        <Tip warning={true}>
+
+        `join_uneven_inputs` is only supported for Distributed Data Parallel training on multiple GPUs. For any other
+        configuration, this method will have no effect.
+
+        </Tip>
+
+        <Tip warning={true}>
+
+        Overidding `even_batches` will not affect iterable-style data loaders.
+
+        </Tip>
+
+        Example:
+
+        ```python
+        >>> from accelerate import Accelerator
+
+        >>> accelerator = Accelerator(even_batches=True)
+        >>> ddp_model, optimizer, dataloader = accelerator.prepare(model, optimizer, dataloader)
+
+        >>> with accelerator.join_uneven_inputs([ddp_model], even_batches=False):
+        ...     for input, output in dataloader:
+        ...         outputs = model(input)
+        ...         loss = loss_func(outputs)
+        ...         loss.backward()
+        ...         optimizer.step()
+        ...         optimizer.zero_grad()
+        ```
+
+        """
+        if is_torch_version("<", "1.10.0"):
+            raise ValueError(f"Joining uneven inputs requires PyTorch >= 1.10.0, You have {torch.__version__}.")
+
+        if self.distributed_type == DistributedType.MULTI_GPU:
+            dl_even_batches_values = []
+
+            if even_batches is not None:
+                # override value in batch sampler for map-style datasets
+                for dl_idx, dl in enumerate(self._dataloaders):
+                    if isinstance(dl, DataLoaderDispatcher):
+                        continue

Discussion: body: The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/accelerate/pr_820). All of your documentation changes will be reflected on that endpoint.
body: @sgugger @muellerzr Is this the sort of thing that you had in mind? If so, I can write some docs and perhaps a small example.

One thing that I was thinking about whilst working on this though is that, although uneven batches is likely only going to be needed during evaluation, the approach that we have taken with setting a variable at the accelerator level is that it will apply to both training and validation. So an example would look something like:

```
accelerator = Accelerator(even_batches=True)

model, optimizer, train_dataloader, val_dataloader = create_components()

ddp_model, optimizer, train_dataloader = accelerator.prepare(model, optimizer, train_dataloader)

# training loop as normal
...


accelerator.even_batches = False

val_dataloader = accelerator.prepare(val_dataloader)

# run eval loop
with accelerator.join_uneven_inputs([ddp_model]):
    ...  
```
It's a bit clunky, but as it is an advanced use case, it may be acceptable.

I'm interested to hear any thoughts that you have!
body: @Chris-hughes10 @sgugger what do you think about something like this. Here's an MVP example showing the basic behavior:

```python
from contextlib import contextmanager

class A:
    "A mock version of what we would call the dataloader"
    def __init__(self, a:int):
        self.a = a

class B:
    "A mock version of what we would call the Accelerator"
    @contextmanager
    def do_the_thing(self, a):
        old_a = a.a
        a.a = 2
        try:
            yield
        finally:
            a.a = old_a

a = A(1)
b = B()

with b.do_the_thing(a):
    print(a.a)
print(a.a)
```
You'd find that `a.a` prints 2 in the first and 1 in the second. We can apply that similarly here since we expect if the user is calling this that they will be training/evaluating/doing things on uneven inputs, leading to an API such as:

```python
with accelerator.join_uneven_inputs([model], dataloader=eval_dataloader):
   ...
```
Inside this ctx manager `a.a` would be `eval_dataloader.even_batches`, and if it wasn't disabled originally we can print for the user that using this context manager will temporarily change this value or something. This let's us get rid of the explicit of setting in the accelerator and having to prepare later, so the only code change is the context manager. 

Lmk what you both think

Otherwise I think you've done a great job with it! Super small API design nit of my own 😄 

(specifically this would do `eval_dataloader.even_batches=False` during the `a.a` bits in the context manager example)

It could probably go as far as accept a tuple of models and dataloaders and Accelerate could figure out which one is which and apply them rather than having to pass in a list 🤔 Decisions decisions :) 

If wanted I could do some of this in a follow up PR, as this PR is great already
body: Agreed, maybe remove some of the magic with an `override` arg as well (that defaults to `False`) and a very explicit docstring as middle ground. (Also would love others opinions 😄 )
body: Hi both, I have given this some thought and have a couple of ideas. 

My first reaction was to agree with @sgugger, changing the `even_batches` value in the dataloader seems a little too magical, but it definitely makes the particular use case smoother. I think the main irk that I have with this is that it is changing a value on an unrelated object; if it was changing the value on the accelerator itself, I think it would be quite clean. Additionally, it seems a bit like a side effect, as it isn't really part of the same operation as the join.

One thought that I had, is that we could expose a method on the dataloader to make it easy to toggle this behaviour. Perhaps something like:
```
accelerator = Accelerator()

model, optimizer, train_dataloader, val_dataloader = create_components()

ddp_model, optimizer, train_dataloader, val_dataloader = accelerator.prepare(model, optimizer, train_dataloader, val_dataloader)

# training loop as normal
...


val_dataloader.ensure_even_batches(enabled=False)

# run eval loop
with accelerator.join_uneven_inputs([ddp_model]):
    ...  
```
There is still an extra step, but then it is very explicit what is going on. Conversely, this does mean that the flag in the accelerator is a little redundant. Also, it probably means that, in most cases, setting this flag and the context manager will have to be used together.

As an alternative, if we would like to keep everything accelerator-centric, we could potentially have the accelerator maintain a reference to any prepared dataloaders. This way, if we change the setting on the accelerator, we could change the value on all dataloaders. This could look something like:

```
accelerator = Accelerator(even_batches=True)

model, optimizer, train_dataloader, val_dataloader = create_components()

ddp_model, optimizer, train_dataloader, val_dataloader = accelerator.prepare(model, optimizer, train_dataloader, val_dataloader)

# training loop as normal
...

# run eval loop
with accelerator.join_uneven_inputs([ddp_model], even_batches=False):
    ...  
```
Personally, I think that this is quite clean, but it does mean adding additional logic to the accelerator, and perhaps the actual difference between this solution and the one proposed by @muellerzr is a bit academic.

EDIT: cc @pacman100 , I think we commented essentially at the same time. What do you think about the second approach proposed here?
body: Yes, I like the second approach proposed above ✨
body: @Chris-hughes10 I too like that solution, as it keeps in line with how we design bits like this such as the gradient accumulation wrapper 😄 It also keeps the dataloader preparation the same as optimizer and model in the sense of pointers. This just means that there would also need to be a change to `accelerator.free_memory()` to also free the dataloader references. 

I also think it should just be an `Accelerator` only attribute, there wouldn't be a need to put this in `AcceleratorState` since we deem that too magical already 😃 
body: Works for me too!
body: It looks like there are failing tests that are unrelated to this PR. Perhaps https://github.com/huggingface/accelerate/pull/825 will solve the problem.
body: @Chris-hughes10 correct, you can ignore the "checkpoint" and "rest" failures here. 
body: I have made some more updates, so I think it should be feature complete now.

Unfortunately, I am getting a couple of failing tests on the CI, which look like they occur in the bf16 section of the test script, comparing training on CPU and GPU. I am unable to reproduce locally as I don't have a device with bf16 available, and the rest of the tests in that script seem to pass. 

It is quite strange though, because looking at what it is failing, it doesn't really interact with anything in this PR other than the accelerator maintaining references to the dataloader; which doesn't seem to be a problem for the other training configurations in the script. After noticing that the accelerator state is cleared multiple times in the test - which does not include the dataloaders, as discussed above - I also tried adding resetting the accelerator manually before the bf16 section, using `accelerator.free_memory`, but this did not make any difference. 

Honesty, I am stumped at this point. Do you guys have any ideas on what could be the issue?
body: @Chris-hughes10 I'll be able to dig into this further tomorrow and see if I can't root out what's going on, so don't stress on it today 😃 
body: Hello @Chris-hughes10 , I went over the PR and tests. I found the cause of the tests failing. In the below screenshot please find the changes required in `test_accelerator.py` to resolve failing tests.

<img width="764" alt="Screenshot 2022-11-10 at 7 28 32 PM" src="https://user-images.githubusercontent.com/13534540/201111269-367a3afa-5957-4b3e-8288-bbdd11107393.png">

body: > Hello @Chris-hughes10 , I went over the PR and tests. I found the cause of the tests failing. In the below screenshot please find the changes required in `test_accelerator.py` to resolve failing tests.
> 


Awesome, thanks for this! I didn't anticipate that the state would carry over outside of each test, but I suppose it makes sense for a singleton. I will make the change.
body: The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/accelerate/pr_820). All of your documentation changes will be reflected on that endpoint.
body: There seems to be a problem with building the PR documentation, is this something that I can fix?
body: @Chris-hughes10 nope don't worry about that, it's a problem on our side :)
body: Thanks a bunch @Chris-hughes10, excellent work! I'll be keeping an eye on our CI that runs multi-gpu tests post this merge to see if anything is up and I'll let you know if there is and I can't seem to figure out what may be failing quickly :) 
body: Thanks for helping to get this merged @muellerzr @sgugger @pacman100! 
body: The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/accelerate/pr_820). All of your documentation changes will be reflected on that endpoint.