diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index b841bc809..1254c19e9 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -22,6 +22,7 @@
 
 from packaging import version
 
+from .checkpointing import load_accelerator_state, save_accelerator_state
 from .data_loader import prepare_data_loader
 from .kwargs_handlers import DistributedDataParallelKwargs, GradScalerKwargs, InitProcessGroupKwargs, KwargsHandler
 from .optimizer import AcceleratedOptimizer
@@ -40,6 +41,7 @@
 
 if is_deepspeed_available():
     import deepspeed
+
     from .deepspeed_utils import DeepSpeedEngineWrapper, DeepSpeedOptimizerWrapper
 
 import logging
@@ -560,6 +562,36 @@ def save(self, obj, f):
         """
         save(obj, f)
 
+    def save_state(self, output_dir: str):
+        """
+        Saves the current states of the model, optimizer, scaler, and RNG generators.
+
+        Args:
+            output_dir (:obj:`str` or :obj:`os.PathLike`):
+                The name of the folder to save all relevant weights and states.
+        """
+        # Check if folder exists
+        output_dir = os.path.expanduser(output_dir)
+        os.makedirs(output_dir, exist_ok=True)
+        logger.info(f"Saving current state to {output_dir}")
+        weights = [self.get_state_dict(m) for m in self._models]
+        return save_accelerator_state(output_dir, weights, self._optimizers, self.state.process_index, self.scaler)
+
+    def load_state(self, input_dir: str):
+        """
+        Loads the current states of the model, optimizer, scaler, and RNG generators.
+
+        Args:
+            input_dir (:obj:`str` or :obj:`os.PathLike`):
+                The name of the folder all relevant weights and states were saved in.
+        """
+        # Check if folder exists
+        input_dir = os.path.expanduser(input_dir)
+        if not os.path.isdir(input_dir):
+            raise ValueError(f"Tried to find {input_dir} but folder does not exist")
+        logger.info(f"Loading states from {input_dir}")
+        load_accelerator_state(input_dir, self._models, self._optimizers, self.state.process_index, self.scaler)
+
     def free_memory(self):
         """
         Will release all references to the internal objects stored and call the garbage collector. You should call this
diff --git a/src/accelerate/checkpointing.py b/src/accelerate/checkpointing.py
new file mode 100644
index 000000000..37d68b90c
--- /dev/null
+++ b/src/accelerate/checkpointing.py
@@ -0,0 +1,134 @@
+# Copyright 2022 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+import os
+import random
+from typing import List
+
+import numpy as np
+import torch
+from torch.cuda.amp import GradScaler
+
+from .state import is_tpu_available
+from .utils import MODEL_NAME, OPTIMIZER_NAME, RNG_STATE_NAME, SCALER_NAME, save
+
+
+if is_tpu_available():
+    import torch_xla.core.xla_model as xm
+
+import logging
+
+
+logger = logging.getLogger(__name__)
+
+
+def save_accelerator_state(
+    output_dir: str, model_states: List[dict], optimizers: list, process_index: int, scaler: GradScaler = None
+):
+    """
+    Saves the current states of the models, optimizers, scaler, and RNG generators to a given directory.
+
+    Args:
+        output_dir (:obj:`str` or :obj:`os.PathLike`):
+            The name of the folder to save all relevant weights and states.
+        model_states (:obj:`List[torch.nn.Module]`):
+            A list of model states
+        optimizers (:obj:`List[torch.optim.Optimizer]`):
+            A list of optimizer instances
+        process_index (:obj:`int`):
+            The current process index in the Accelerator state
+        scaler (:obj:`torch.cuda.amp.GradScaler`, `optional`):
+            An optional gradient scaler instance to save
+    """
+    # Model states
+    for i, state in enumerate(model_states):
+        weights_name = f"{MODEL_NAME}.bin" if i == 0 else f"{MODEL_NAME}_{i}.bin"
+        output_model_file = os.path.join(output_dir, weights_name)
+        save(state, output_model_file)
+        logger.info(f"Model weights saved in {output_model_file}")
+    # Optimizer states
+    for i, opt in enumerate(optimizers):
+        state = opt.state_dict()
+        optimizer_name = f"{OPTIMIZER_NAME}.bin" if i == 0 else f"{OPTIMIZER_NAME}_{i}.bin"
+        output_optimizer_file = os.path.join(output_dir, optimizer_name)
+        save(state, output_optimizer_file)
+        logger.info(f"Optimizer state saved in {output_optimizer_file}")
+    # GradScaler state
+    if scaler is not None:
+        state = scaler.state_dict()
+        output_scaler_file = os.path.join(output_dir, SCALER_NAME)
+        torch.save(state, output_scaler_file)
+        logger.info(f"Gradient scaler state saved in {output_scaler_file}")
+    # Random number generator states
+    states = {}
+    states_name = f"{RNG_STATE_NAME}_{process_index}.pkl"
+    states["random_state"] = random.getstate()
+    states["numpy_random_seed"] = np.random.get_state()
+    states["torch_manual_seed"] = torch.get_rng_state()
+    states["torch_cuda_manual_seed"] = torch.cuda.get_rng_state_all()
+    # ^^ safe to call this function even if cuda is not available
+    if is_tpu_available():
+        states["xm_seed"] = torch.tensor(xm.get_rng_state())
+    output_states_file = os.path.join(output_dir, states_name)
+    torch.save(states, output_states_file)
+    logger.info(f"Random states saved in {output_states_file}")
+    return output_dir
+
+
+def load_accelerator_state(input_dir, models, optimizers, process_index, scaler=None):
+    """
+    Loads states of the models, optimizers, scaler, and RNG generators from a given directory.
+
+    Args:
+        input_dir (:obj:`str` or :obj:`os.PathLike`):
+            The name of the folder to load all relevant weights and states.
+        model_stmodelsates (:obj:`List[torch.nn.Module]`):
+            A list of model instances
+        optimizers (:obj:`List[torch.optim.Optimizer]`):
+            A list of optimizer instances
+        process_index (:obj:`int`):
+            The current process index in the Accelerator state
+        scaler (:obj:`torch.cuda.amp.GradScaler`, `optional`):
+            An optional `GradScaler` instance to load
+    """
+    # Model states
+    for i, model in enumerate(models):
+        weights_name = f"{MODEL_NAME}.bin" if i == 0 else f"{MODEL_NAME}_{i}.bin"
+        input_model_file = os.path.join(input_dir, weights_name)
+        models[i].load_state_dict(torch.load(input_model_file))
+    logger.info("All model weights loaded successfully")
+
+    # Optimizer states
+    for i, opt in enumerate(optimizers):
+        optimizer_name = f"{OPTIMIZER_NAME}.bin" if i == 0 else f"{OPTIMIZER_NAME}_{i}.bin"
+        input_optimizer_file = os.path.join(input_dir, optimizer_name)
+        optimizers[i].load_state_dict(torch.load(input_optimizer_file))
+    logger.info("All optimizer states loaded successfully")
+
+    # GradScaler state
+    if scaler is not None:
+        input_scaler_file = os.path.join(input_dir, SCALER_NAME)
+        scaler.load_state_dict(torch.load(input_scaler_file))
+        logger.info("GradScaler state loaded successfully")
+
+    # Random states
+    states = torch.load(os.path.join(input_dir, f"{RNG_STATE_NAME}_{process_index}.pkl"))
+    random.setstate(states["random_state"])
+    np.random.set_state(states["numpy_random_seed"])
+    torch.set_rng_state(states["torch_manual_seed"])
+    torch.cuda.set_rng_state_all(states["torch_cuda_manual_seed"])
+    # ^^ safe to call this function even if cuda is not available
+    if is_tpu_available():
+        xm.set_rng_state(states["xm_seed"])
+    logger.info("All random states loaded successfully")
diff --git a/src/accelerate/utils.py b/src/accelerate/utils.py
index 3af792f1f..abd5152ae 100644
--- a/src/accelerate/utils.py
+++ b/src/accelerate/utils.py
@@ -43,6 +43,11 @@ def is_sagemaker_available():
 if is_deepspeed_available():
     from deepspeed import DeepSpeedEngine
 
+SCALER_NAME = "scaler.pt"
+MODEL_NAME = "pytorch_model"
+RNG_STATE_NAME = "random_states"
+OPTIMIZER_NAME = "optimizer"
+
 
 class RNGType(Enum):
     TORCH = "torch"
diff --git a/tests/test_state_checkpointing.py b/tests/test_state_checkpointing.py
new file mode 100644
index 000000000..a74dcb724
--- /dev/null
+++ b/tests/test_state_checkpointing.py
@@ -0,0 +1,125 @@
+# Copyright 2022 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+import logging
+import os
+import random
+import tempfile
+import unittest
+
+import torch
+from torch import nn
+from torch.utils.data import DataLoader, TensorDataset
+
+from accelerate import Accelerator
+from accelerate.utils import set_seed
+
+
+logger = logging.getLogger(__name__)
+
+
+def dummy_dataloaders(a=2, b=3, batch_size=16, n_train_batches: int = 10, n_valid_batches: int = 2):
+    "Generates a tuple of dummy DataLoaders to test with"
+
+    def get_dataset(n_batches):
+        x = torch.randn(batch_size * n_batches, 1)
+        return TensorDataset(x, a * x + b + 0.1 * torch.randn(batch_size * n_batches, 1))
+
+    train_dataset = get_dataset(n_train_batches)
+    valid_dataset = get_dataset(n_valid_batches)
+    train_dataloader = DataLoader(train_dataset, shuffle=True, batch_size=batch_size, num_workers=4)
+    valid_dataloader = DataLoader(valid_dataset, shuffle=False, batch_size=batch_size, num_workers=4)
+    return (train_dataloader, valid_dataloader)
+
+
+def train(num_epochs, model, dataloader, optimizer, accelerator):
+    "Trains for `num_epochs`"
+    rands = []
+    for epoch in range(num_epochs):
+        # Train quickly
+        model.train()
+        for step, batch in enumerate(dataloader):
+            x, y = batch
+            outputs = model(x)
+            loss = torch.nn.functional.mse_loss(outputs, y)
+            accelerator.backward(loss)
+            optimizer.step()
+            optimizer.zero_grad()
+        rands.append(random.random())  # Introduce some randomness
+    return rands
+
+
+class DummyModel(nn.Module):
+    "Simple model to do y=mx+b"
+
+    def __init__(self):
+        super().__init__()
+        self.a = nn.Parameter(torch.randn(1))
+        self.b = nn.Parameter(torch.randn(1))
+
+    def forward(self, x):
+        return x * self.a + self.b
+
+
+class CheckpointTest(unittest.TestCase):
+    def test_can_resume_training(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            set_seed(42)
+            model = DummyModel()
+            optimizer = torch.optim.Adam(params=model.parameters(), lr=1e-3)
+            train_dataloader, valid_dataloader = dummy_dataloaders()
+            # Train baseline
+            accelerator = Accelerator()
+            model, optimizer, train_dataloader, valid_dataloader = accelerator.prepare(
+                model, optimizer, train_dataloader, valid_dataloader
+            )
+            # Save initial
+            initial = os.path.join(tmpdir, "initial")
+            accelerator.save_state(initial)
+            (a, b) = model.a.item(), model.b.item()
+            opt_state = optimizer.state_dict()
+            ground_truth_rands = train(3, model, train_dataloader, optimizer, accelerator)
+            (a1, b1) = model.a.item(), model.b.item()
+            opt_state1 = optimizer.state_dict()
+
+            # Train partially
+            set_seed(42)
+            model = DummyModel()
+            optimizer = torch.optim.Adam(params=model.parameters(), lr=1e-3)
+            train_dataloader, valid_dataloader = dummy_dataloaders()
+            accelerator = Accelerator()
+            model, optimizer, train_dataloader, valid_dataloader = accelerator.prepare(
+                model, optimizer, train_dataloader, valid_dataloader
+            )
+            accelerator.load_state(initial)
+            (a2, b2) = model.a.item(), model.b.item()
+            opt_state2 = optimizer.state_dict()
+            self.assertEqual(a, a2)
+            self.assertEqual(b, b2)
+            self.assertEqual(opt_state, opt_state2)
+
+            test_rands = train(2, model, train_dataloader, optimizer, accelerator)
+            # Save everything
+            checkpoint = os.path.join(tmpdir, "checkpoint")
+            accelerator.save_state(checkpoint)
+
+            # Load everything back in and make sure all states work
+            accelerator.load_state(checkpoint)
+            test_rands += train(1, model, train_dataloader, optimizer, accelerator)
+            (a3, b3) = model.a.item(), model.b.item()
+            opt_state3 = optimizer.state_dict()
+            self.assertEqual(a1, a3)
+            self.assertEqual(b1, b3)
+            self.assertEqual(opt_state1, opt_state3)
+            self.assertEqual(ground_truth_rands, test_rands)


Review comments: body: As much as possible, I would prefer to have everything in another module and have `accelerator` be a thin wrapper. I think we could have a new `checkpointing.py` file where a function that does all of this (taking `models`, `optimizers`, `scaler`). Same for the loading.
diff_hunk: @@ -560,6 +566,109 @@ def save(self, obj, f):
         """
         save(obj, f)
 
+    def save_state(self, output_dir: str):
body: ```suggestion
            if i != 0:
                weights_name += f"_{i}"
```
Is it a typo? I think it's for every model **but** the first one.
diff_hunk: @@ -560,6 +566,109 @@ def save(self, obj, f):
         """
         save(obj, f)
 
+    def save_state(self, output_dir: str):
+        """
+        Saves the current states of the model, optimizer, scalar, and RNG generators to `folder_name`.
+
+        Args:
+            output_dir (:obj:`str` or :obj:`os.PathLike`):
+                The name of the folder to save all relevant weights and states.
+        """
+        # Check if folder exists
+        output_dir = os.path.expanduser(output_dir)
+        os.makedirs(output_dir, exist_ok=True)
+        logger.info(f"Saving current state to {output_dir}")
+
+        # Model states
+        for i, model in enumerate(self._models):
+            state = self.get_state_dict(model)
+            weights_name = "pytorch_model"
+            if i == 0:
+                weights_name += f"_{i}"
body: Should be `accelerator.save` here (otherwise won't work for TPUs).
diff_hunk: @@ -560,6 +566,109 @@ def save(self, obj, f):
         """
         save(obj, f)
 
+    def save_state(self, output_dir: str):
+        """
+        Saves the current states of the model, optimizer, scalar, and RNG generators to `folder_name`.
+
+        Args:
+            output_dir (:obj:`str` or :obj:`os.PathLike`):
+                The name of the folder to save all relevant weights and states.
+        """
+        # Check if folder exists
+        output_dir = os.path.expanduser(output_dir)
+        os.makedirs(output_dir, exist_ok=True)
+        logger.info(f"Saving current state to {output_dir}")
+
+        # Model states
+        for i, model in enumerate(self._models):
+            state = self.get_state_dict(model)
+            weights_name = "pytorch_model"
+            if i == 0:
+                weights_name += f"_{i}"
+            weights_name += ".bin"
+            output_model_file = os.path.join(output_dir, weights_name)
+            torch.save(state, output_model_file)
body: Same here:
```suggestion
            accelerator.save(state, output_optimizer_file)
```
diff_hunk: @@ -560,6 +566,109 @@ def save(self, obj, f):
         """
         save(obj, f)
 
+    def save_state(self, output_dir: str):
+        """
+        Saves the current states of the model, optimizer, scalar, and RNG generators to `folder_name`.
+
+        Args:
+            output_dir (:obj:`str` or :obj:`os.PathLike`):
+                The name of the folder to save all relevant weights and states.
+        """
+        # Check if folder exists
+        output_dir = os.path.expanduser(output_dir)
+        os.makedirs(output_dir, exist_ok=True)
+        logger.info(f"Saving current state to {output_dir}")
+
+        # Model states
+        for i, model in enumerate(self._models):
+            state = self.get_state_dict(model)
+            weights_name = "pytorch_model"
+            if i == 0:
+                weights_name += f"_{i}"
+            weights_name += ".bin"
+            output_model_file = os.path.join(output_dir, weights_name)
+            torch.save(state, output_model_file)
+            logger.info(f"Model weights saved in {output_model_file}")
+        # Optimizer states
+        for i, opt in enumerate(self._optimizers):
+            state = opt.state_dict()
+            optimizer_name = "optimizer"
+            if i == 0:
+                optimizer_name += f"_{i}"
+            optimizer_name += ".pt"
+            output_optimizer_file = os.path.join(output_dir, optimizer_name)
+            torch.save(state, output_optimizer_file)
body: ```suggestion
            if i != 0:
                optimizer_name += f"_{i}"
```
As above!
diff_hunk: @@ -560,6 +566,109 @@ def save(self, obj, f):
         """
         save(obj, f)
 
+    def save_state(self, output_dir: str):
+        """
+        Saves the current states of the model, optimizer, scalar, and RNG generators to `folder_name`.
+
+        Args:
+            output_dir (:obj:`str` or :obj:`os.PathLike`):
+                The name of the folder to save all relevant weights and states.
+        """
+        # Check if folder exists
+        output_dir = os.path.expanduser(output_dir)
+        os.makedirs(output_dir, exist_ok=True)
+        logger.info(f"Saving current state to {output_dir}")
+
+        # Model states
+        for i, model in enumerate(self._models):
+            state = self.get_state_dict(model)
+            weights_name = "pytorch_model"
+            if i == 0:
+                weights_name += f"_{i}"
+            weights_name += ".bin"
+            output_model_file = os.path.join(output_dir, weights_name)
+            torch.save(state, output_model_file)
+            logger.info(f"Model weights saved in {output_model_file}")
+        # Optimizer states
+        for i, opt in enumerate(self._optimizers):
+            state = opt.state_dict()
+            optimizer_name = "optimizer"
+            if i == 0:
+                optimizer_name += f"_{i}"
body: ```suggestion
            scaler_name = "scaler.pt"
```
Let's use the same name as in the `Trainer`.
diff_hunk: @@ -560,6 +566,109 @@ def save(self, obj, f):
         """
         save(obj, f)
 
+    def save_state(self, output_dir: str):
+        """
+        Saves the current states of the model, optimizer, scalar, and RNG generators to `folder_name`.
+
+        Args:
+            output_dir (:obj:`str` or :obj:`os.PathLike`):
+                The name of the folder to save all relevant weights and states.
+        """
+        # Check if folder exists
+        output_dir = os.path.expanduser(output_dir)
+        os.makedirs(output_dir, exist_ok=True)
+        logger.info(f"Saving current state to {output_dir}")
+
+        # Model states
+        for i, model in enumerate(self._models):
+            state = self.get_state_dict(model)
+            weights_name = "pytorch_model"
+            if i == 0:
+                weights_name += f"_{i}"
+            weights_name += ".bin"
+            output_model_file = os.path.join(output_dir, weights_name)
+            torch.save(state, output_model_file)
+            logger.info(f"Model weights saved in {output_model_file}")
+        # Optimizer states
+        for i, opt in enumerate(self._optimizers):
+            state = opt.state_dict()
+            optimizer_name = "optimizer"
+            if i == 0:
+                optimizer_name += f"_{i}"
+            optimizer_name += ".pt"
+            output_optimizer_file = os.path.join(output_dir, optimizer_name)
+            torch.save(state, output_optimizer_file)
+            logger.info(f"Optimizer state saved in {output_optimizer_file}")
+        # GradScalar state
+        if self.scaler is not None:
+            state = self.scaler.state_dict()
+            scalar_name = "scaler.bin"
body: ```suggestion
            logger.info(f"Gradient scaler state saved in {output_scaler_file}")
```
diff_hunk: @@ -560,6 +566,109 @@ def save(self, obj, f):
         """
         save(obj, f)
 
+    def save_state(self, output_dir: str):
+        """
+        Saves the current states of the model, optimizer, scalar, and RNG generators to `folder_name`.
+
+        Args:
+            output_dir (:obj:`str` or :obj:`os.PathLike`):
+                The name of the folder to save all relevant weights and states.
+        """
+        # Check if folder exists
+        output_dir = os.path.expanduser(output_dir)
+        os.makedirs(output_dir, exist_ok=True)
+        logger.info(f"Saving current state to {output_dir}")
+
+        # Model states
+        for i, model in enumerate(self._models):
+            state = self.get_state_dict(model)
+            weights_name = "pytorch_model"
+            if i == 0:
+                weights_name += f"_{i}"
+            weights_name += ".bin"
+            output_model_file = os.path.join(output_dir, weights_name)
+            torch.save(state, output_model_file)
+            logger.info(f"Model weights saved in {output_model_file}")
+        # Optimizer states
+        for i, opt in enumerate(self._optimizers):
+            state = opt.state_dict()
+            optimizer_name = "optimizer"
+            if i == 0:
+                optimizer_name += f"_{i}"
+            optimizer_name += ".pt"
+            output_optimizer_file = os.path.join(output_dir, optimizer_name)
+            torch.save(state, output_optimizer_file)
+            logger.info(f"Optimizer state saved in {output_optimizer_file}")
+        # GradScalar state
+        if self.scaler is not None:
+            state = self.scaler.state_dict()
+            scalar_name = "scaler.bin"
+            output_scalar_file = os.path.join(output_dir, scalar_name)
+            torch.save(state, output_scalar_file)
+            logger.info(f"GradScalar state saved in {output_scalar_file}")
body: ```suggestion
            output_scaler_file = os.path.join(output_dir, scaler_name)
            torch.save(state, output_scaler_file)
```
diff_hunk: @@ -560,6 +566,109 @@ def save(self, obj, f):
         """
         save(obj, f)
 
+    def save_state(self, output_dir: str):
+        """
+        Saves the current states of the model, optimizer, scalar, and RNG generators to `folder_name`.
+
+        Args:
+            output_dir (:obj:`str` or :obj:`os.PathLike`):
+                The name of the folder to save all relevant weights and states.
+        """
+        # Check if folder exists
+        output_dir = os.path.expanduser(output_dir)
+        os.makedirs(output_dir, exist_ok=True)
+        logger.info(f"Saving current state to {output_dir}")
+
+        # Model states
+        for i, model in enumerate(self._models):
+            state = self.get_state_dict(model)
+            weights_name = "pytorch_model"
+            if i == 0:
+                weights_name += f"_{i}"
+            weights_name += ".bin"
+            output_model_file = os.path.join(output_dir, weights_name)
+            torch.save(state, output_model_file)
+            logger.info(f"Model weights saved in {output_model_file}")
+        # Optimizer states
+        for i, opt in enumerate(self._optimizers):
+            state = opt.state_dict()
+            optimizer_name = "optimizer"
+            if i == 0:
+                optimizer_name += f"_{i}"
+            optimizer_name += ".pt"
+            output_optimizer_file = os.path.join(output_dir, optimizer_name)
+            torch.save(state, output_optimizer_file)
+            logger.info(f"Optimizer state saved in {output_optimizer_file}")
+        # GradScalar state
+        if self.scaler is not None:
+            state = self.scaler.state_dict()
+            scalar_name = "scaler.bin"
+            output_scalar_file = os.path.join(output_dir, scalar_name)
+            torch.save(state, output_scalar_file)
body: I realize we are missing the learning rate scheduler -> we should thus add it as an optional argument to the `save_state` function, or add a `register_for_checkpoints` method to the Accelerator object that would take objects not passed to the `prepare` method but the user still wants to save.
diff_hunk: @@ -560,6 +566,109 @@ def save(self, obj, f):
         """
         save(obj, f)
 
+    def save_state(self, output_dir: str):
+        """
+        Saves the current states of the model, optimizer, scalar, and RNG generators to `folder_name`.
+
+        Args:
+            output_dir (:obj:`str` or :obj:`os.PathLike`):
+                The name of the folder to save all relevant weights and states.
+        """
+        # Check if folder exists
+        output_dir = os.path.expanduser(output_dir)
+        os.makedirs(output_dir, exist_ok=True)
+        logger.info(f"Saving current state to {output_dir}")
+
+        # Model states
+        for i, model in enumerate(self._models):
+            state = self.get_state_dict(model)
+            weights_name = "pytorch_model"
+            if i == 0:
+                weights_name += f"_{i}"
+            weights_name += ".bin"
+            output_model_file = os.path.join(output_dir, weights_name)
+            torch.save(state, output_model_file)
+            logger.info(f"Model weights saved in {output_model_file}")
+        # Optimizer states
+        for i, opt in enumerate(self._optimizers):
+            state = opt.state_dict()
+            optimizer_name = "optimizer"
+            if i == 0:
+                optimizer_name += f"_{i}"
+            optimizer_name += ".pt"
+            output_optimizer_file = os.path.join(output_dir, optimizer_name)
+            torch.save(state, output_optimizer_file)
+            logger.info(f"Optimizer state saved in {output_optimizer_file}")
+        # GradScalar state
body: ```suggestion
        # Random number generator states
```
diff_hunk: @@ -560,6 +566,109 @@ def save(self, obj, f):
         """
         save(obj, f)
 
+    def save_state(self, output_dir: str):
+        """
+        Saves the current states of the model, optimizer, scalar, and RNG generators to `folder_name`.
+
+        Args:
+            output_dir (:obj:`str` or :obj:`os.PathLike`):
+                The name of the folder to save all relevant weights and states.
+        """
+        # Check if folder exists
+        output_dir = os.path.expanduser(output_dir)
+        os.makedirs(output_dir, exist_ok=True)
+        logger.info(f"Saving current state to {output_dir}")
+
+        # Model states
+        for i, model in enumerate(self._models):
+            state = self.get_state_dict(model)
+            weights_name = "pytorch_model"
+            if i == 0:
+                weights_name += f"_{i}"
+            weights_name += ".bin"
+            output_model_file = os.path.join(output_dir, weights_name)
+            torch.save(state, output_model_file)
+            logger.info(f"Model weights saved in {output_model_file}")
+        # Optimizer states
+        for i, opt in enumerate(self._optimizers):
+            state = opt.state_dict()
+            optimizer_name = "optimizer"
+            if i == 0:
+                optimizer_name += f"_{i}"
+            optimizer_name += ".pt"
+            output_optimizer_file = os.path.join(output_dir, optimizer_name)
+            torch.save(state, output_optimizer_file)
+            logger.info(f"Optimizer state saved in {output_optimizer_file}")
+        # GradScalar state
+        if self.scaler is not None:
+            state = self.scaler.state_dict()
+            scalar_name = "scaler.bin"
+            output_scalar_file = os.path.join(output_dir, scalar_name)
+            torch.save(state, output_scalar_file)
+            logger.info(f"GradScalar state saved in {output_scalar_file}")
+        # Random states
body: We use `rng_state.pth` in the Trainer, so we should do so here as well for consistency across libs.

Also, we need to save the states on each process as they might be different (in `rng_state_{i}.pth` for process i)
diff_hunk: @@ -560,6 +566,109 @@ def save(self, obj, f):
         """
         save(obj, f)
 
+    def save_state(self, output_dir: str):
+        """
+        Saves the current states of the model, optimizer, scalar, and RNG generators to `folder_name`.
+
+        Args:
+            output_dir (:obj:`str` or :obj:`os.PathLike`):
+                The name of the folder to save all relevant weights and states.
+        """
+        # Check if folder exists
+        output_dir = os.path.expanduser(output_dir)
+        os.makedirs(output_dir, exist_ok=True)
+        logger.info(f"Saving current state to {output_dir}")
+
+        # Model states
+        for i, model in enumerate(self._models):
+            state = self.get_state_dict(model)
+            weights_name = "pytorch_model"
+            if i == 0:
+                weights_name += f"_{i}"
+            weights_name += ".bin"
+            output_model_file = os.path.join(output_dir, weights_name)
+            torch.save(state, output_model_file)
+            logger.info(f"Model weights saved in {output_model_file}")
+        # Optimizer states
+        for i, opt in enumerate(self._optimizers):
+            state = opt.state_dict()
+            optimizer_name = "optimizer"
+            if i == 0:
+                optimizer_name += f"_{i}"
+            optimizer_name += ".pt"
+            output_optimizer_file = os.path.join(output_dir, optimizer_name)
+            torch.save(state, output_optimizer_file)
+            logger.info(f"Optimizer state saved in {output_optimizer_file}")
+        # GradScalar state
+        if self.scaler is not None:
+            state = self.scaler.state_dict()
+            scalar_name = "scaler.bin"
+            output_scalar_file = os.path.join(output_dir, scalar_name)
+            torch.save(state, output_scalar_file)
+            logger.info(f"GradScalar state saved in {output_scalar_file}")
+        # Random states
+        states = {}
+        states_name = "random_states.pkl"
body: ```suggestion
            if i != 0:
                weights_name += f"_{i}"
```
diff_hunk: @@ -560,6 +566,109 @@ def save(self, obj, f):
         """
         save(obj, f)
 
+    def save_state(self, output_dir: str):
+        """
+        Saves the current states of the model, optimizer, scalar, and RNG generators to `folder_name`.
+
+        Args:
+            output_dir (:obj:`str` or :obj:`os.PathLike`):
+                The name of the folder to save all relevant weights and states.
+        """
+        # Check if folder exists
+        output_dir = os.path.expanduser(output_dir)
+        os.makedirs(output_dir, exist_ok=True)
+        logger.info(f"Saving current state to {output_dir}")
+
+        # Model states
+        for i, model in enumerate(self._models):
+            state = self.get_state_dict(model)
+            weights_name = "pytorch_model"
+            if i == 0:
+                weights_name += f"_{i}"
+            weights_name += ".bin"
+            output_model_file = os.path.join(output_dir, weights_name)
+            torch.save(state, output_model_file)
+            logger.info(f"Model weights saved in {output_model_file}")
+        # Optimizer states
+        for i, opt in enumerate(self._optimizers):
+            state = opt.state_dict()
+            optimizer_name = "optimizer"
+            if i == 0:
+                optimizer_name += f"_{i}"
+            optimizer_name += ".pt"
+            output_optimizer_file = os.path.join(output_dir, optimizer_name)
+            torch.save(state, output_optimizer_file)
+            logger.info(f"Optimizer state saved in {output_optimizer_file}")
+        # GradScalar state
+        if self.scaler is not None:
+            state = self.scaler.state_dict()
+            scalar_name = "scaler.bin"
+            output_scalar_file = os.path.join(output_dir, scalar_name)
+            torch.save(state, output_scalar_file)
+            logger.info(f"GradScalar state saved in {output_scalar_file}")
+        # Random states
+        states = {}
+        states_name = "random_states.pkl"
+        states["random_state"] = random.getstate()
+        states["numpy_random_seed"] = np.random.get_state()
+        states["torch_manual_seed"] = torch.get_rng_state()
+        states["torch_cuda_manual_seed"] = torch.cuda.get_rng_state_all()
+        # ^^ safe to call this function even if cuda is not available
+        if is_tpu_available():
+            states["xm_seed"] = torch.tensor(xm.get_rng_state())
+        output_states_file = os.path.join(output_dir, states_name)
+        torch.save(states, output_states_file)
+        logger.info(f"Random states saved in {output_states_file}")
+
+    def load_state(self, input_dir: str):
+        """
+        Loads the current states of the model, optimizer, scalar, and RNG generators from `folder_name`.
+
+        Args:
+            input_dir (:obj:`str` or :obj:`os.PathLike`):
+                The name of the folder all relevant weights and states were saved in.
+        """
+        # Check if folder exists
+        input_dir = os.path.expanduser(input_dir)
+        if not os.path.isdir(input_dir):
+            raise ValueError(f"Tried to find {input_dir} but folder does not exist")
+        logger.info(f"Loading states from {input_dir}")
+
+        # Model states
+        for i, model in enumerate(self._models):
+            weights_name = "pytorch_model"
+            if i == 0:
+                weights_name += f"_{i}"
body: ```suggestion
            if i != 0:
                optimizer_name += f"_{i}"
```
diff_hunk: @@ -560,6 +566,109 @@ def save(self, obj, f):
         """
         save(obj, f)
 
+    def save_state(self, output_dir: str):
+        """
+        Saves the current states of the model, optimizer, scalar, and RNG generators to `folder_name`.
+
+        Args:
+            output_dir (:obj:`str` or :obj:`os.PathLike`):
+                The name of the folder to save all relevant weights and states.
+        """
+        # Check if folder exists
+        output_dir = os.path.expanduser(output_dir)
+        os.makedirs(output_dir, exist_ok=True)
+        logger.info(f"Saving current state to {output_dir}")
+
+        # Model states
+        for i, model in enumerate(self._models):
+            state = self.get_state_dict(model)
+            weights_name = "pytorch_model"
+            if i == 0:
+                weights_name += f"_{i}"
+            weights_name += ".bin"
+            output_model_file = os.path.join(output_dir, weights_name)
+            torch.save(state, output_model_file)
+            logger.info(f"Model weights saved in {output_model_file}")
+        # Optimizer states
+        for i, opt in enumerate(self._optimizers):
+            state = opt.state_dict()
+            optimizer_name = "optimizer"
+            if i == 0:
+                optimizer_name += f"_{i}"
+            optimizer_name += ".pt"
+            output_optimizer_file = os.path.join(output_dir, optimizer_name)
+            torch.save(state, output_optimizer_file)
+            logger.info(f"Optimizer state saved in {output_optimizer_file}")
+        # GradScalar state
+        if self.scaler is not None:
+            state = self.scaler.state_dict()
+            scalar_name = "scaler.bin"
+            output_scalar_file = os.path.join(output_dir, scalar_name)
+            torch.save(state, output_scalar_file)
+            logger.info(f"GradScalar state saved in {output_scalar_file}")
+        # Random states
+        states = {}
+        states_name = "random_states.pkl"
+        states["random_state"] = random.getstate()
+        states["numpy_random_seed"] = np.random.get_state()
+        states["torch_manual_seed"] = torch.get_rng_state()
+        states["torch_cuda_manual_seed"] = torch.cuda.get_rng_state_all()
+        # ^^ safe to call this function even if cuda is not available
+        if is_tpu_available():
+            states["xm_seed"] = torch.tensor(xm.get_rng_state())
+        output_states_file = os.path.join(output_dir, states_name)
+        torch.save(states, output_states_file)
+        logger.info(f"Random states saved in {output_states_file}")
+
+    def load_state(self, input_dir: str):
+        """
+        Loads the current states of the model, optimizer, scalar, and RNG generators from `folder_name`.
+
+        Args:
+            input_dir (:obj:`str` or :obj:`os.PathLike`):
+                The name of the folder all relevant weights and states were saved in.
+        """
+        # Check if folder exists
+        input_dir = os.path.expanduser(input_dir)
+        if not os.path.isdir(input_dir):
+            raise ValueError(f"Tried to find {input_dir} but folder does not exist")
+        logger.info(f"Loading states from {input_dir}")
+
+        # Model states
+        for i, model in enumerate(self._models):
+            weights_name = "pytorch_model"
+            if i == 0:
+                weights_name += f"_{i}"
+            weights_name += ".bin"
+            input_model_file = os.path.join(input_dir, weights_name)
+            self._models[i].load_state_dict(torch.load(input_model_file))
+        logger.info("All model weights loaded successfully")
+        # Optimizer states
+        for i, opt in enumerate(self._optimizers):
+            optimizer_name = "optimizer"
+            if i == 0:
+                optimizer_name += f"_{i}"
body: The names should probably be constants, so we only have to change them in one place ;-)
diff_hunk: @@ -560,6 +566,109 @@ def save(self, obj, f):
         """
         save(obj, f)
 
+    def save_state(self, output_dir: str):
+        """
+        Saves the current states of the model, optimizer, scalar, and RNG generators to `folder_name`.
+
+        Args:
+            output_dir (:obj:`str` or :obj:`os.PathLike`):
+                The name of the folder to save all relevant weights and states.
+        """
+        # Check if folder exists
+        output_dir = os.path.expanduser(output_dir)
+        os.makedirs(output_dir, exist_ok=True)
+        logger.info(f"Saving current state to {output_dir}")
+
+        # Model states
+        for i, model in enumerate(self._models):
+            state = self.get_state_dict(model)
+            weights_name = "pytorch_model"
+            if i == 0:
+                weights_name += f"_{i}"
+            weights_name += ".bin"
+            output_model_file = os.path.join(output_dir, weights_name)
+            torch.save(state, output_model_file)
+            logger.info(f"Model weights saved in {output_model_file}")
+        # Optimizer states
+        for i, opt in enumerate(self._optimizers):
+            state = opt.state_dict()
+            optimizer_name = "optimizer"
+            if i == 0:
+                optimizer_name += f"_{i}"
+            optimizer_name += ".pt"
+            output_optimizer_file = os.path.join(output_dir, optimizer_name)
+            torch.save(state, output_optimizer_file)
+            logger.info(f"Optimizer state saved in {output_optimizer_file}")
+        # GradScalar state
+        if self.scaler is not None:
+            state = self.scaler.state_dict()
+            scalar_name = "scaler.bin"
+            output_scalar_file = os.path.join(output_dir, scalar_name)
+            torch.save(state, output_scalar_file)
+            logger.info(f"GradScalar state saved in {output_scalar_file}")
+        # Random states
+        states = {}
+        states_name = "random_states.pkl"
+        states["random_state"] = random.getstate()
+        states["numpy_random_seed"] = np.random.get_state()
+        states["torch_manual_seed"] = torch.get_rng_state()
+        states["torch_cuda_manual_seed"] = torch.cuda.get_rng_state_all()
+        # ^^ safe to call this function even if cuda is not available
+        if is_tpu_available():
+            states["xm_seed"] = torch.tensor(xm.get_rng_state())
+        output_states_file = os.path.join(output_dir, states_name)
+        torch.save(states, output_states_file)
+        logger.info(f"Random states saved in {output_states_file}")
+
+    def load_state(self, input_dir: str):
+        """
+        Loads the current states of the model, optimizer, scalar, and RNG generators from `folder_name`.
+
+        Args:
+            input_dir (:obj:`str` or :obj:`os.PathLike`):
+                The name of the folder all relevant weights and states were saved in.
+        """
+        # Check if folder exists
+        input_dir = os.path.expanduser(input_dir)
+        if not os.path.isdir(input_dir):
+            raise ValueError(f"Tried to find {input_dir} but folder does not exist")
+        logger.info(f"Loading states from {input_dir}")
+
+        # Model states
+        for i, model in enumerate(self._models):
+            weights_name = "pytorch_model"
+            if i == 0:
+                weights_name += f"_{i}"
+            weights_name += ".bin"
+            input_model_file = os.path.join(input_dir, weights_name)
+            self._models[i].load_state_dict(torch.load(input_model_file))
+        logger.info("All model weights loaded successfully")
+        # Optimizer states
+        for i, opt in enumerate(self._optimizers):
+            optimizer_name = "optimizer"
+            if i == 0:
+                optimizer_name += f"_{i}"
+            optimizer_name += ".pt"
+            input_optimizer_file = os.path.join(input_dir, optimizer_name)
+            self._optimizers[i].load_state_dict(torch.load(input_optimizer_file))
+        logger.info("All optimizer states loaded successfully")
+        # GradScalar state
+        if self.scaler is not None:
+            scaler_name = "scaler.bin"
body: Same ehre + one for each process.
diff_hunk: @@ -560,6 +566,109 @@ def save(self, obj, f):
         """
         save(obj, f)
 
+    def save_state(self, output_dir: str):
+        """
+        Saves the current states of the model, optimizer, scalar, and RNG generators to `folder_name`.
+
+        Args:
+            output_dir (:obj:`str` or :obj:`os.PathLike`):
+                The name of the folder to save all relevant weights and states.
+        """
+        # Check if folder exists
+        output_dir = os.path.expanduser(output_dir)
+        os.makedirs(output_dir, exist_ok=True)
+        logger.info(f"Saving current state to {output_dir}")
+
+        # Model states
+        for i, model in enumerate(self._models):
+            state = self.get_state_dict(model)
+            weights_name = "pytorch_model"
+            if i == 0:
+                weights_name += f"_{i}"
+            weights_name += ".bin"
+            output_model_file = os.path.join(output_dir, weights_name)
+            torch.save(state, output_model_file)
+            logger.info(f"Model weights saved in {output_model_file}")
+        # Optimizer states
+        for i, opt in enumerate(self._optimizers):
+            state = opt.state_dict()
+            optimizer_name = "optimizer"
+            if i == 0:
+                optimizer_name += f"_{i}"
+            optimizer_name += ".pt"
+            output_optimizer_file = os.path.join(output_dir, optimizer_name)
+            torch.save(state, output_optimizer_file)
+            logger.info(f"Optimizer state saved in {output_optimizer_file}")
+        # GradScalar state
+        if self.scaler is not None:
+            state = self.scaler.state_dict()
+            scalar_name = "scaler.bin"
+            output_scalar_file = os.path.join(output_dir, scalar_name)
+            torch.save(state, output_scalar_file)
+            logger.info(f"GradScalar state saved in {output_scalar_file}")
+        # Random states
+        states = {}
+        states_name = "random_states.pkl"
+        states["random_state"] = random.getstate()
+        states["numpy_random_seed"] = np.random.get_state()
+        states["torch_manual_seed"] = torch.get_rng_state()
+        states["torch_cuda_manual_seed"] = torch.cuda.get_rng_state_all()
+        # ^^ safe to call this function even if cuda is not available
+        if is_tpu_available():
+            states["xm_seed"] = torch.tensor(xm.get_rng_state())
+        output_states_file = os.path.join(output_dir, states_name)
+        torch.save(states, output_states_file)
+        logger.info(f"Random states saved in {output_states_file}")
+
+    def load_state(self, input_dir: str):
+        """
+        Loads the current states of the model, optimizer, scalar, and RNG generators from `folder_name`.
+
+        Args:
+            input_dir (:obj:`str` or :obj:`os.PathLike`):
+                The name of the folder all relevant weights and states were saved in.
+        """
+        # Check if folder exists
+        input_dir = os.path.expanduser(input_dir)
+        if not os.path.isdir(input_dir):
+            raise ValueError(f"Tried to find {input_dir} but folder does not exist")
+        logger.info(f"Loading states from {input_dir}")
+
+        # Model states
+        for i, model in enumerate(self._models):
+            weights_name = "pytorch_model"
+            if i == 0:
+                weights_name += f"_{i}"
+            weights_name += ".bin"
+            input_model_file = os.path.join(input_dir, weights_name)
+            self._models[i].load_state_dict(torch.load(input_model_file))
+        logger.info("All model weights loaded successfully")
+        # Optimizer states
+        for i, opt in enumerate(self._optimizers):
+            optimizer_name = "optimizer"
+            if i == 0:
+                optimizer_name += f"_{i}"
+            optimizer_name += ".pt"
+            input_optimizer_file = os.path.join(input_dir, optimizer_name)
+            self._optimizers[i].load_state_dict(torch.load(input_optimizer_file))
+        logger.info("All optimizer states loaded successfully")
+        # GradScalar state
+        if self.scaler is not None:
+            scaler_name = "scaler.bin"
+            input_scaler_file = os.path.join(input_dir, scaler_name)
+            self.scaler.load_state_dict(torch.load(input_scaler_file))
+            logger.info("GradScalar state loaded successfully")
+        # Random states
+        states = torch.load(os.path.join(input_dir, "random_states.pkl"))
body: We could wrap all of this in some train function that does the training.
diff_hunk: @@ -0,0 +1,148 @@
+# Copyright 2022 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+import argparse
+import os
+import tempfile
+
+import torch
+from torch import nn
+from torch.utils.data import DataLoader, TensorDataset
+
+from accelerate import Accelerator
+from accelerate.utils import set_seed
+
+
+def dummy_dataloaders(a=2, b=3, batch_size=16, n_train_batches: int = 10, n_valid_batches: int = 2):
+    "Generates a tuple of dummy DataLoaders to test with"
+
+    def get_dataset(n_batches):
+        x = torch.randn(batch_size * n_batches, 1)
+        return TensorDataset(x, a * x + b + 0.1 * torch.randn(batch_size * n_batches, 1))
+
+    train_dataset = get_dataset(n_train_batches)
+    valid_dataset = get_dataset(n_valid_batches)
+    train_dataloader = DataLoader(train_dataset, shuffle=True, batch_size=batch_size, num_workers=4)
+    valid_dataloader = DataLoader(valid_dataset, shuffle=False, batch_size=batch_size, num_workers=4)
+    return (train_dataloader, valid_dataloader)
+
+
+class DummyModel(nn.Module):
+    "Simple model to do y=mx+b"
+
+    def __init__(self):
+        super().__init__()
+        self.a = nn.Parameter(torch.randn(1))
+        self.b = nn.Parameter(torch.randn(1))
+
+    def forward(self, x):
+        return x * self.a + self.b
+
+
+def test_can_resume_training(args):
+    with tempfile.TemporaryDirectory() as tmpdir:
+        set_seed(42)
+        model = DummyModel()
+        optimizer = torch.optim.Adam(params=model.parameters(), lr=1e-3)
+        train_dataloader, valid_dataloader = dummy_dataloaders()
+        # Train baseline
+        accelerator = Accelerator(
+            fp16=args.fp16, cpu=args.cpu, mixed_precision=args.mixed_precision, device_placement=True
+        )
+        model, optimizer, train_dataloader, valid_dataloader = accelerator.prepare(
+            model, optimizer, train_dataloader, valid_dataloader
+        )
+        # Save initial
+        initial = os.path.join(tmpdir, "initial")
+        accelerator.save_state(initial)
+        model_unwrapped = accelerator.unwrap_model(model)
+        (a, b) = model_unwrapped.a.item(), model_unwrapped.b.item()
+        opt_state = optimizer.state_dict()
+        for epoch in range(3):
+            # Train quickly
+            model.train()
+            for step, batch in enumerate(train_dataloader):
+                x, y = batch
+                outputs = model(x)
+                loss = torch.nn.functional.mse_loss(outputs, y)
+                accelerator.backward(loss)
+                optimizer.step()
+                optimizer.zero_grad()
+
+        model_unwrapped = accelerator.unwrap_model(model)
+        (a1, b1) = model_unwrapped.a.item(), model_unwrapped.b.item()
+        opt_state1 = optimizer.state_dict()
+
+        # Train partially
body: Nit: we should use `unittest` here with `self.assertEqual` and `self.assertEqualDict` instead of those asserts.
diff_hunk: @@ -0,0 +1,148 @@
+# Copyright 2022 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+import argparse
+import os
+import tempfile
+
+import torch
+from torch import nn
+from torch.utils.data import DataLoader, TensorDataset
+
+from accelerate import Accelerator
+from accelerate.utils import set_seed
+
+
+def dummy_dataloaders(a=2, b=3, batch_size=16, n_train_batches: int = 10, n_valid_batches: int = 2):
+    "Generates a tuple of dummy DataLoaders to test with"
+
+    def get_dataset(n_batches):
+        x = torch.randn(batch_size * n_batches, 1)
+        return TensorDataset(x, a * x + b + 0.1 * torch.randn(batch_size * n_batches, 1))
+
+    train_dataset = get_dataset(n_train_batches)
+    valid_dataset = get_dataset(n_valid_batches)
+    train_dataloader = DataLoader(train_dataset, shuffle=True, batch_size=batch_size, num_workers=4)
+    valid_dataloader = DataLoader(valid_dataset, shuffle=False, batch_size=batch_size, num_workers=4)
+    return (train_dataloader, valid_dataloader)
+
+
+class DummyModel(nn.Module):
+    "Simple model to do y=mx+b"
+
+    def __init__(self):
+        super().__init__()
+        self.a = nn.Parameter(torch.randn(1))
+        self.b = nn.Parameter(torch.randn(1))
+
+    def forward(self, x):
+        return x * self.a + self.b
+
+
+def test_can_resume_training(args):
+    with tempfile.TemporaryDirectory() as tmpdir:
+        set_seed(42)
+        model = DummyModel()
+        optimizer = torch.optim.Adam(params=model.parameters(), lr=1e-3)
+        train_dataloader, valid_dataloader = dummy_dataloaders()
+        # Train baseline
+        accelerator = Accelerator(
+            fp16=args.fp16, cpu=args.cpu, mixed_precision=args.mixed_precision, device_placement=True
+        )
+        model, optimizer, train_dataloader, valid_dataloader = accelerator.prepare(
+            model, optimizer, train_dataloader, valid_dataloader
+        )
+        # Save initial
+        initial = os.path.join(tmpdir, "initial")
+        accelerator.save_state(initial)
+        model_unwrapped = accelerator.unwrap_model(model)
+        (a, b) = model_unwrapped.a.item(), model_unwrapped.b.item()
+        opt_state = optimizer.state_dict()
+        for epoch in range(3):
+            # Train quickly
+            model.train()
+            for step, batch in enumerate(train_dataloader):
+                x, y = batch
+                outputs = model(x)
+                loss = torch.nn.functional.mse_loss(outputs, y)
+                accelerator.backward(loss)
+                optimizer.step()
+                optimizer.zero_grad()
+
+        model_unwrapped = accelerator.unwrap_model(model)
+        (a1, b1) = model_unwrapped.a.item(), model_unwrapped.b.item()
+        opt_state1 = optimizer.state_dict()
+
+        # Train partially
+        accelerator.load_state(initial)
+        model_unwrapped = accelerator.unwrap_model(model)
+        (a2, b2) = model_unwrapped.a.item(), model_unwrapped.b.item()
+        opt_state2 = optimizer.state_dict()
+        assert a == a2
+        assert b == b2
+        assert opt_state == opt_state2
body: Same as above.
diff_hunk: @@ -0,0 +1,148 @@
+# Copyright 2022 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+import argparse
+import os
+import tempfile
+
+import torch
+from torch import nn
+from torch.utils.data import DataLoader, TensorDataset
+
+from accelerate import Accelerator
+from accelerate.utils import set_seed
+
+
+def dummy_dataloaders(a=2, b=3, batch_size=16, n_train_batches: int = 10, n_valid_batches: int = 2):
+    "Generates a tuple of dummy DataLoaders to test with"
+
+    def get_dataset(n_batches):
+        x = torch.randn(batch_size * n_batches, 1)
+        return TensorDataset(x, a * x + b + 0.1 * torch.randn(batch_size * n_batches, 1))
+
+    train_dataset = get_dataset(n_train_batches)
+    valid_dataset = get_dataset(n_valid_batches)
+    train_dataloader = DataLoader(train_dataset, shuffle=True, batch_size=batch_size, num_workers=4)
+    valid_dataloader = DataLoader(valid_dataset, shuffle=False, batch_size=batch_size, num_workers=4)
+    return (train_dataloader, valid_dataloader)
+
+
+class DummyModel(nn.Module):
+    "Simple model to do y=mx+b"
+
+    def __init__(self):
+        super().__init__()
+        self.a = nn.Parameter(torch.randn(1))
+        self.b = nn.Parameter(torch.randn(1))
+
+    def forward(self, x):
+        return x * self.a + self.b
+
+
+def test_can_resume_training(args):
+    with tempfile.TemporaryDirectory() as tmpdir:
+        set_seed(42)
+        model = DummyModel()
+        optimizer = torch.optim.Adam(params=model.parameters(), lr=1e-3)
+        train_dataloader, valid_dataloader = dummy_dataloaders()
+        # Train baseline
+        accelerator = Accelerator(
+            fp16=args.fp16, cpu=args.cpu, mixed_precision=args.mixed_precision, device_placement=True
+        )
+        model, optimizer, train_dataloader, valid_dataloader = accelerator.prepare(
+            model, optimizer, train_dataloader, valid_dataloader
+        )
+        # Save initial
+        initial = os.path.join(tmpdir, "initial")
+        accelerator.save_state(initial)
+        model_unwrapped = accelerator.unwrap_model(model)
+        (a, b) = model_unwrapped.a.item(), model_unwrapped.b.item()
+        opt_state = optimizer.state_dict()
+        for epoch in range(3):
+            # Train quickly
+            model.train()
+            for step, batch in enumerate(train_dataloader):
+                x, y = batch
+                outputs = model(x)
+                loss = torch.nn.functional.mse_loss(outputs, y)
+                accelerator.backward(loss)
+                optimizer.step()
+                optimizer.zero_grad()
+
+        model_unwrapped = accelerator.unwrap_model(model)
+        (a1, b1) = model_unwrapped.a.item(), model_unwrapped.b.item()
+        opt_state1 = optimizer.state_dict()
+
+        # Train partially
+        accelerator.load_state(initial)
+        model_unwrapped = accelerator.unwrap_model(model)
+        (a2, b2) = model_unwrapped.a.item(), model_unwrapped.b.item()
+        opt_state2 = optimizer.state_dict()
+        assert a == a2
+        assert b == b2
+        assert opt_state == opt_state2
+
+        for epoch in range(2):
+            # Train quickly
+            model.train()
+            for step, batch in enumerate(train_dataloader):
+                x, y = batch
+                outputs = model(x)
+                loss = torch.nn.functional.mse_loss(outputs, y)
+                accelerator.backward(loss)
+                optimizer.step()
+                optimizer.zero_grad()
+        # Save everything
+        checkpoint = os.path.join(tmpdir, "checkpoint")
+        accelerator.save_state(checkpoint)
+
+        # Load everything back in and make sure all states work
+        accelerator.load_state(checkpoint)
+        for epoch in range(1):
+            # Train quickly
+            model.train()
+            for step, batch in enumerate(train_dataloader):
+                x, y = batch
+                outputs = model(x)
+                loss = torch.nn.functional.mse_loss(outputs, y)
+                accelerator.backward(loss)
+                optimizer.step()
+                optimizer.zero_grad()
+        model_unwrapped = accelerator.unwrap_model(model)
+        (a3, b3) = model_unwrapped.a.item(), model_unwrapped.b.item()
+        opt_state3 = optimizer.state_dict()
+        assert a1 == a3
+        assert b1 == b3
+        assert opt_state1 == opt_state3
body: This one being deprecated, I wouldn't pass it/create it.
diff_hunk: @@ -0,0 +1,148 @@
+# Copyright 2022 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+import argparse
+import os
+import tempfile
+
+import torch
+from torch import nn
+from torch.utils.data import DataLoader, TensorDataset
+
+from accelerate import Accelerator
+from accelerate.utils import set_seed
+
+
+def dummy_dataloaders(a=2, b=3, batch_size=16, n_train_batches: int = 10, n_valid_batches: int = 2):
+    "Generates a tuple of dummy DataLoaders to test with"
+
+    def get_dataset(n_batches):
+        x = torch.randn(batch_size * n_batches, 1)
+        return TensorDataset(x, a * x + b + 0.1 * torch.randn(batch_size * n_batches, 1))
+
+    train_dataset = get_dataset(n_train_batches)
+    valid_dataset = get_dataset(n_valid_batches)
+    train_dataloader = DataLoader(train_dataset, shuffle=True, batch_size=batch_size, num_workers=4)
+    valid_dataloader = DataLoader(valid_dataset, shuffle=False, batch_size=batch_size, num_workers=4)
+    return (train_dataloader, valid_dataloader)
+
+
+class DummyModel(nn.Module):
+    "Simple model to do y=mx+b"
+
+    def __init__(self):
+        super().__init__()
+        self.a = nn.Parameter(torch.randn(1))
+        self.b = nn.Parameter(torch.randn(1))
+
+    def forward(self, x):
+        return x * self.a + self.b
+
+
+def test_can_resume_training(args):
+    with tempfile.TemporaryDirectory() as tmpdir:
+        set_seed(42)
+        model = DummyModel()
+        optimizer = torch.optim.Adam(params=model.parameters(), lr=1e-3)
+        train_dataloader, valid_dataloader = dummy_dataloaders()
+        # Train baseline
+        accelerator = Accelerator(
+            fp16=args.fp16, cpu=args.cpu, mixed_precision=args.mixed_precision, device_placement=True
+        )
+        model, optimizer, train_dataloader, valid_dataloader = accelerator.prepare(
+            model, optimizer, train_dataloader, valid_dataloader
+        )
+        # Save initial
+        initial = os.path.join(tmpdir, "initial")
+        accelerator.save_state(initial)
+        model_unwrapped = accelerator.unwrap_model(model)
+        (a, b) = model_unwrapped.a.item(), model_unwrapped.b.item()
+        opt_state = optimizer.state_dict()
+        for epoch in range(3):
+            # Train quickly
+            model.train()
+            for step, batch in enumerate(train_dataloader):
+                x, y = batch
+                outputs = model(x)
+                loss = torch.nn.functional.mse_loss(outputs, y)
+                accelerator.backward(loss)
+                optimizer.step()
+                optimizer.zero_grad()
+
+        model_unwrapped = accelerator.unwrap_model(model)
+        (a1, b1) = model_unwrapped.a.item(), model_unwrapped.b.item()
+        opt_state1 = optimizer.state_dict()
+
+        # Train partially
+        accelerator.load_state(initial)
+        model_unwrapped = accelerator.unwrap_model(model)
+        (a2, b2) = model_unwrapped.a.item(), model_unwrapped.b.item()
+        opt_state2 = optimizer.state_dict()
+        assert a == a2
+        assert b == b2
+        assert opt_state == opt_state2
+
+        for epoch in range(2):
+            # Train quickly
+            model.train()
+            for step, batch in enumerate(train_dataloader):
+                x, y = batch
+                outputs = model(x)
+                loss = torch.nn.functional.mse_loss(outputs, y)
+                accelerator.backward(loss)
+                optimizer.step()
+                optimizer.zero_grad()
+        # Save everything
+        checkpoint = os.path.join(tmpdir, "checkpoint")
+        accelerator.save_state(checkpoint)
+
+        # Load everything back in and make sure all states work
+        accelerator.load_state(checkpoint)
+        for epoch in range(1):
+            # Train quickly
+            model.train()
+            for step, batch in enumerate(train_dataloader):
+                x, y = batch
+                outputs = model(x)
+                loss = torch.nn.functional.mse_loss(outputs, y)
+                accelerator.backward(loss)
+                optimizer.step()
+                optimizer.zero_grad()
+        model_unwrapped = accelerator.unwrap_model(model)
+        (a3, b3) = model_unwrapped.a.item(), model_unwrapped.b.item()
+        opt_state3 = optimizer.state_dict()
+        assert a1 == a3
+        assert b1 == b3
+        assert opt_state1 == opt_state3
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Simple example of training script.")
+    parser.add_argument("--fp16", action="store_true", help="If passed, will use FP16 training.")
body: I'd like some advice on how to convert my script -> a proper test. I tried looking in `test_multigpu` but still doesn't quite line up with how the testing workflow operates
diff_hunk: @@ -0,0 +1,148 @@
+# Copyright 2022 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+import argparse
+import os
+import tempfile
+
+import torch
+from torch import nn
+from torch.utils.data import DataLoader, TensorDataset
+
+from accelerate import Accelerator
+from accelerate.utils import set_seed
+
+
+def dummy_dataloaders(a=2, b=3, batch_size=16, n_train_batches: int = 10, n_valid_batches: int = 2):
+    "Generates a tuple of dummy DataLoaders to test with"
+
+    def get_dataset(n_batches):
+        x = torch.randn(batch_size * n_batches, 1)
+        return TensorDataset(x, a * x + b + 0.1 * torch.randn(batch_size * n_batches, 1))
+
+    train_dataset = get_dataset(n_train_batches)
+    valid_dataset = get_dataset(n_valid_batches)
+    train_dataloader = DataLoader(train_dataset, shuffle=True, batch_size=batch_size, num_workers=4)
+    valid_dataloader = DataLoader(valid_dataset, shuffle=False, batch_size=batch_size, num_workers=4)
+    return (train_dataloader, valid_dataloader)
+
+
+class DummyModel(nn.Module):
+    "Simple model to do y=mx+b"
+
+    def __init__(self):
+        super().__init__()
+        self.a = nn.Parameter(torch.randn(1))
+        self.b = nn.Parameter(torch.randn(1))
+
+    def forward(self, x):
+        return x * self.a + self.b
+
+
+def test_can_resume_training(args):
+    with tempfile.TemporaryDirectory() as tmpdir:
+        set_seed(42)
+        model = DummyModel()
+        optimizer = torch.optim.Adam(params=model.parameters(), lr=1e-3)
+        train_dataloader, valid_dataloader = dummy_dataloaders()
+        # Train baseline
+        accelerator = Accelerator(
+            fp16=args.fp16, cpu=args.cpu, mixed_precision=args.mixed_precision, device_placement=True
+        )
+        model, optimizer, train_dataloader, valid_dataloader = accelerator.prepare(
+            model, optimizer, train_dataloader, valid_dataloader
+        )
+        # Save initial
+        initial = os.path.join(tmpdir, "initial")
+        accelerator.save_state(initial)
+        model_unwrapped = accelerator.unwrap_model(model)
+        (a, b) = model_unwrapped.a.item(), model_unwrapped.b.item()
+        opt_state = optimizer.state_dict()
+        for epoch in range(3):
+            # Train quickly
+            model.train()
+            for step, batch in enumerate(train_dataloader):
+                x, y = batch
+                outputs = model(x)
+                loss = torch.nn.functional.mse_loss(outputs, y)
+                accelerator.backward(loss)
+                optimizer.step()
+                optimizer.zero_grad()
+
+        model_unwrapped = accelerator.unwrap_model(model)
+        (a1, b1) = model_unwrapped.a.item(), model_unwrapped.b.item()
+        opt_state1 = optimizer.state_dict()
+
+        # Train partially
+        accelerator.load_state(initial)
+        model_unwrapped = accelerator.unwrap_model(model)
+        (a2, b2) = model_unwrapped.a.item(), model_unwrapped.b.item()
+        opt_state2 = optimizer.state_dict()
+        assert a == a2
+        assert b == b2
+        assert opt_state == opt_state2
+
+        for epoch in range(2):
+            # Train quickly
+            model.train()
+            for step, batch in enumerate(train_dataloader):
+                x, y = batch
+                outputs = model(x)
+                loss = torch.nn.functional.mse_loss(outputs, y)
+                accelerator.backward(loss)
+                optimizer.step()
+                optimizer.zero_grad()
+        # Save everything
+        checkpoint = os.path.join(tmpdir, "checkpoint")
+        accelerator.save_state(checkpoint)
+
+        # Load everything back in and make sure all states work
+        accelerator.load_state(checkpoint)
+        for epoch in range(1):
+            # Train quickly
+            model.train()
+            for step, batch in enumerate(train_dataloader):
+                x, y = batch
+                outputs = model(x)
+                loss = torch.nn.functional.mse_loss(outputs, y)
+                accelerator.backward(loss)
+                optimizer.step()
+                optimizer.zero_grad()
+        model_unwrapped = accelerator.unwrap_model(model)
+        (a3, b3) = model_unwrapped.a.item(), model_unwrapped.b.item()
+        opt_state3 = optimizer.state_dict()
+        assert a1 == a3
+        assert b1 == b3
+        assert opt_state1 == opt_state3
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Simple example of training script.")
+    parser.add_argument("--fp16", action="store_true", help="If passed, will use FP16 training.")
body: ```suggestion
        Saves the current states of the model, optimizer, scaler, and RNG generators.
```
It's actually output_dir but since we document it just after, no need to mention it up there I believe.
diff_hunk: @@ -560,6 +562,36 @@ def save(self, obj, f):
         """
         save(obj, f)
 
+    def save_state(self, output_dir: str):
+        """
+        Saves the current states of the model, optimizer, scaler, and RNG generators to `folder_name`.
body: ```suggestion
        Loads the current states of the model, optimizer, scaler, and RNG generators.
```
diff_hunk: @@ -560,6 +562,36 @@ def save(self, obj, f):
         """
         save(obj, f)
 
+    def save_state(self, output_dir: str):
+        """
+        Saves the current states of the model, optimizer, scaler, and RNG generators to `folder_name`.
+
+        Args:
+            output_dir (:obj:`str` or :obj:`os.PathLike`):
+                The name of the folder to save all relevant weights and states.
+        """
+        # Check if folder exists
+        output_dir = os.path.expanduser(output_dir)
+        os.makedirs(output_dir, exist_ok=True)
+        logger.info(f"Saving current state to {output_dir}")
+        weights = [self.get_state_dict(m) for m in self._models]
+        return save_accelerator_state(output_dir, weights, self._optimizers, self.state.process_index, self.scaler)
+
+    def load_state(self, input_dir: str):
+        """
+        Loads the current states of the model, optimizer, scaler, and RNG generators from `folder_name`.
body: ```suggestion
# Copyright 2022 The HuggingFace Team. All rights reserved.
```
diff_hunk: @@ -0,0 +1,134 @@
+# Copyright 2021 The HuggingFace Team. All rights reserved.
body: ```suggestion
    Saves the current states of the models, optimizers, scaler, and RNG generators to a given directory.
```
Let's leave the argument names for the arg descriptions and be moe "English" in the description (or if you want them, you need to put them in double backticks or add an `:obj:` to have them code-formatted)
diff_hunk: @@ -0,0 +1,134 @@
+# Copyright 2021 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+import os
+import random
+from typing import List
+
+import numpy as np
+import torch
+from torch.cuda.amp import GradScaler
+
+from .state import is_tpu_available
+from .utils import MODEL_NAME, OPTIMIZER_NAME, RNG_STATE_NAME, SCALER_NAME, save
+
+
+if is_tpu_available():
+    import torch_xla.core.xla_model as xm
+
+import logging
+
+
+logger = logging.getLogger(__name__)
+
+
+def save_accelerator_state(
+    output_dir: str, model_states: List[dict], optimizers: list, process_index: int, scaler: GradScaler = None
+):
+    """
+    Saves the current states of the `models`, `optimizers`, scaler, and RNG generators to `output_dir`.
body: ```suggestion
        scaler (:obj:`torch.cuda.amp.GradScaler`, `optional`):
```
diff_hunk: @@ -0,0 +1,134 @@
+# Copyright 2021 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+import os
+import random
+from typing import List
+
+import numpy as np
+import torch
+from torch.cuda.amp import GradScaler
+
+from .state import is_tpu_available
+from .utils import MODEL_NAME, OPTIMIZER_NAME, RNG_STATE_NAME, SCALER_NAME, save
+
+
+if is_tpu_available():
+    import torch_xla.core.xla_model as xm
+
+import logging
+
+
+logger = logging.getLogger(__name__)
+
+
+def save_accelerator_state(
+    output_dir: str, model_states: List[dict], optimizers: list, process_index: int, scaler: GradScaler = None
+):
+    """
+    Saves the current states of the `models`, `optimizers`, scaler, and RNG generators to `output_dir`.
+
+    Args:
+        output_dir (:obj:`str` or :obj:`os.PathLike`):
+            The name of the folder to save all relevant weights and states.
+        model_states (:obj:`list`):
+            A list of model states
+        optimizers (:obj:`list`):
+            A list of optimizer instances
+        process_index (:obj:`int`):
+            The current process index in the Accelerator state
+        scaler (:obj:`GradScaler`), Optional:
body: ```suggestion
            An optional gradient scaler instance to save
```
(The type is already described)
diff_hunk: @@ -0,0 +1,134 @@
+# Copyright 2021 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+import os
+import random
+from typing import List
+
+import numpy as np
+import torch
+from torch.cuda.amp import GradScaler
+
+from .state import is_tpu_available
+from .utils import MODEL_NAME, OPTIMIZER_NAME, RNG_STATE_NAME, SCALER_NAME, save
+
+
+if is_tpu_available():
+    import torch_xla.core.xla_model as xm
+
+import logging
+
+
+logger = logging.getLogger(__name__)
+
+
+def save_accelerator_state(
+    output_dir: str, model_states: List[dict], optimizers: list, process_index: int, scaler: GradScaler = None
+):
+    """
+    Saves the current states of the `models`, `optimizers`, scaler, and RNG generators to `output_dir`.
+
+    Args:
+        output_dir (:obj:`str` or :obj:`os.PathLike`):
+            The name of the folder to save all relevant weights and states.
+        model_states (:obj:`list`):
+            A list of model states
+        optimizers (:obj:`list`):
+            A list of optimizer instances
+        process_index (:obj:`int`):
+            The current process index in the Accelerator state
+        scaler (:obj:`GradScaler`), Optional:
+            An optional `GradScaler` instance to save
body: This also needs to use the `save` function (especially for TPUs).
diff_hunk: @@ -0,0 +1,134 @@
+# Copyright 2021 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+import os
+import random
+from typing import List
+
+import numpy as np
+import torch
+from torch.cuda.amp import GradScaler
+
+from .state import is_tpu_available
+from .utils import MODEL_NAME, OPTIMIZER_NAME, RNG_STATE_NAME, SCALER_NAME, save
+
+
+if is_tpu_available():
+    import torch_xla.core.xla_model as xm
+
+import logging
+
+
+logger = logging.getLogger(__name__)
+
+
+def save_accelerator_state(
+    output_dir: str, model_states: List[dict], optimizers: list, process_index: int, scaler: GradScaler = None
+):
+    """
+    Saves the current states of the `models`, `optimizers`, scaler, and RNG generators to `output_dir`.
+
+    Args:
+        output_dir (:obj:`str` or :obj:`os.PathLike`):
+            The name of the folder to save all relevant weights and states.
+        model_states (:obj:`list`):
+            A list of model states
+        optimizers (:obj:`list`):
+            A list of optimizer instances
+        process_index (:obj:`int`):
+            The current process index in the Accelerator state
+        scaler (:obj:`GradScaler`), Optional:
+            An optional `GradScaler` instance to save
+    """
+    # Model states
+    for i, state in enumerate(model_states):
+        weights_name = f"{MODEL_NAME}.bin" if i == 0 else f"{MODEL_NAME}_{i}.bin"
+        output_model_file = os.path.join(output_dir, weights_name)
+        save(state, output_model_file)
+        logger.info(f"Model weights saved in {output_model_file}")
+    # Optimizer states
+    for i, opt in enumerate(optimizers):
+        state = opt.state_dict()
+        optimizer_name = f"{OPTIMIZER_NAME}.bin" if i == 0 else f"{OPTIMIZER_NAME}_{i}.bin"
+        output_optimizer_file = os.path.join(output_dir, optimizer_name)
+        torch.save(state, output_optimizer_file)
body: ```suggestion
        model_states (:obj:`List[torch.nn.Module]`):
```
diff_hunk: @@ -0,0 +1,134 @@
+# Copyright 2021 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+import os
+import random
+from typing import List
+
+import numpy as np
+import torch
+from torch.cuda.amp import GradScaler
+
+from .state import is_tpu_available
+from .utils import MODEL_NAME, OPTIMIZER_NAME, RNG_STATE_NAME, SCALER_NAME, save
+
+
+if is_tpu_available():
+    import torch_xla.core.xla_model as xm
+
+import logging
+
+
+logger = logging.getLogger(__name__)
+
+
+def save_accelerator_state(
+    output_dir: str, model_states: List[dict], optimizers: list, process_index: int, scaler: GradScaler = None
+):
+    """
+    Saves the current states of the `models`, `optimizers`, scaler, and RNG generators to `output_dir`.
+
+    Args:
+        output_dir (:obj:`str` or :obj:`os.PathLike`):
+            The name of the folder to save all relevant weights and states.
+        model_states (:obj:`list`):
body: ```suggestion
        optimizers (:obj:`List[torch.optim.Optimizer]`):
```
diff_hunk: @@ -0,0 +1,134 @@
+# Copyright 2021 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+import os
+import random
+from typing import List
+
+import numpy as np
+import torch
+from torch.cuda.amp import GradScaler
+
+from .state import is_tpu_available
+from .utils import MODEL_NAME, OPTIMIZER_NAME, RNG_STATE_NAME, SCALER_NAME, save
+
+
+if is_tpu_available():
+    import torch_xla.core.xla_model as xm
+
+import logging
+
+
+logger = logging.getLogger(__name__)
+
+
+def save_accelerator_state(
+    output_dir: str, model_states: List[dict], optimizers: list, process_index: int, scaler: GradScaler = None
+):
+    """
+    Saves the current states of the `models`, `optimizers`, scaler, and RNG generators to `output_dir`.
+
+    Args:
+        output_dir (:obj:`str` or :obj:`os.PathLike`):
+            The name of the folder to save all relevant weights and states.
+        model_states (:obj:`list`):
+            A list of model states
+        optimizers (:obj:`list`):
body: Fix those as above.
diff_hunk: @@ -0,0 +1,134 @@
+# Copyright 2021 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+import os
+import random
+from typing import List
+
+import numpy as np
+import torch
+from torch.cuda.amp import GradScaler
+
+from .state import is_tpu_available
+from .utils import MODEL_NAME, OPTIMIZER_NAME, RNG_STATE_NAME, SCALER_NAME, save
+
+
+if is_tpu_available():
+    import torch_xla.core.xla_model as xm
+
+import logging
+
+
+logger = logging.getLogger(__name__)
+
+
+def save_accelerator_state(
+    output_dir: str, model_states: List[dict], optimizers: list, process_index: int, scaler: GradScaler = None
+):
+    """
+    Saves the current states of the `models`, `optimizers`, scaler, and RNG generators to `output_dir`.
+
+    Args:
+        output_dir (:obj:`str` or :obj:`os.PathLike`):
+            The name of the folder to save all relevant weights and states.
+        model_states (:obj:`list`):
+            A list of model states
+        optimizers (:obj:`list`):
+            A list of optimizer instances
+        process_index (:obj:`int`):
+            The current process index in the Accelerator state
+        scaler (:obj:`GradScaler`), Optional:
+            An optional `GradScaler` instance to save
+    """
+    # Model states
+    for i, state in enumerate(model_states):
+        weights_name = f"{MODEL_NAME}.bin" if i == 0 else f"{MODEL_NAME}_{i}.bin"
+        output_model_file = os.path.join(output_dir, weights_name)
+        save(state, output_model_file)
+        logger.info(f"Model weights saved in {output_model_file}")
+    # Optimizer states
+    for i, opt in enumerate(optimizers):
+        state = opt.state_dict()
+        optimizer_name = f"{OPTIMIZER_NAME}.bin" if i == 0 else f"{OPTIMIZER_NAME}_{i}.bin"
+        output_optimizer_file = os.path.join(output_dir, optimizer_name)
+        torch.save(state, output_optimizer_file)
+        logger.info(f"Optimizer state saved in {output_optimizer_file}")
+    # GradScaler state
+    if scaler is not None:
+        state = scaler.state_dict()
+        output_scaler_file = os.path.join(output_dir, SCALER_NAME)
+        torch.save(state, output_scaler_file)
+        logger.info(f"Gradient scaler state saved in {output_scaler_file}")
+    # Random number generator states
+    states = {}
+    states_name = f"{RNG_STATE_NAME}_{process_index}.pkl"
+    states["random_state"] = random.getstate()
+    states["numpy_random_seed"] = np.random.get_state()
+    states["torch_manual_seed"] = torch.get_rng_state()
+    states["torch_cuda_manual_seed"] = torch.cuda.get_rng_state()
+    # ^^ safe to call this function even if cuda is not available
+    if is_tpu_available():
+        states["xm_seed"] = torch.tensor(xm.get_rng_state())
+    output_states_file = os.path.join(output_dir, states_name)
+    torch.save(states, output_states_file)
+    logger.info(f"Random states saved in {output_states_file}")
+    return output_dir
+
+
+def load_accelerator_state(input_dir, models, optimizers, process_index, scaler=None):
+    """
+    Loads states of the `models`, `optimizers`, scaler, and RNG generators from `input_dir`.
+
+    Args:
+        input_dir (:obj:`str` or :obj:`os.PathLike`):
+            The name of the folder to load all relevant weights and states.
+        model_stmodelsates (:obj:`list`):
+            A list of model instances
+        optimizers (:obj:`list`):
+            A list of optimizer instances
+        process_index (:obj:`int`):
+            The current process index in the Accelerator state
+        scaler (:obj:`GradScaler`), Optional:
body: ```suggestion
    Loads states of the models, optimizers, scaler, and RNG generators from a given directory.
```
diff_hunk: @@ -0,0 +1,134 @@
+# Copyright 2021 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+import os
+import random
+from typing import List
+
+import numpy as np
+import torch
+from torch.cuda.amp import GradScaler
+
+from .state import is_tpu_available
+from .utils import MODEL_NAME, OPTIMIZER_NAME, RNG_STATE_NAME, SCALER_NAME, save
+
+
+if is_tpu_available():
+    import torch_xla.core.xla_model as xm
+
+import logging
+
+
+logger = logging.getLogger(__name__)
+
+
+def save_accelerator_state(
+    output_dir: str, model_states: List[dict], optimizers: list, process_index: int, scaler: GradScaler = None
+):
+    """
+    Saves the current states of the `models`, `optimizers`, scaler, and RNG generators to `output_dir`.
+
+    Args:
+        output_dir (:obj:`str` or :obj:`os.PathLike`):
+            The name of the folder to save all relevant weights and states.
+        model_states (:obj:`list`):
+            A list of model states
+        optimizers (:obj:`list`):
+            A list of optimizer instances
+        process_index (:obj:`int`):
+            The current process index in the Accelerator state
+        scaler (:obj:`GradScaler`), Optional:
+            An optional `GradScaler` instance to save
+    """
+    # Model states
+    for i, state in enumerate(model_states):
+        weights_name = f"{MODEL_NAME}.bin" if i == 0 else f"{MODEL_NAME}_{i}.bin"
+        output_model_file = os.path.join(output_dir, weights_name)
+        save(state, output_model_file)
+        logger.info(f"Model weights saved in {output_model_file}")
+    # Optimizer states
+    for i, opt in enumerate(optimizers):
+        state = opt.state_dict()
+        optimizer_name = f"{OPTIMIZER_NAME}.bin" if i == 0 else f"{OPTIMIZER_NAME}_{i}.bin"
+        output_optimizer_file = os.path.join(output_dir, optimizer_name)
+        torch.save(state, output_optimizer_file)
+        logger.info(f"Optimizer state saved in {output_optimizer_file}")
+    # GradScaler state
+    if scaler is not None:
+        state = scaler.state_dict()
+        output_scaler_file = os.path.join(output_dir, SCALER_NAME)
+        torch.save(state, output_scaler_file)
+        logger.info(f"Gradient scaler state saved in {output_scaler_file}")
+    # Random number generator states
+    states = {}
+    states_name = f"{RNG_STATE_NAME}_{process_index}.pkl"
+    states["random_state"] = random.getstate()
+    states["numpy_random_seed"] = np.random.get_state()
+    states["torch_manual_seed"] = torch.get_rng_state()
+    states["torch_cuda_manual_seed"] = torch.cuda.get_rng_state()
+    # ^^ safe to call this function even if cuda is not available
+    if is_tpu_available():
+        states["xm_seed"] = torch.tensor(xm.get_rng_state())
+    output_states_file = os.path.join(output_dir, states_name)
+    torch.save(states, output_states_file)
+    logger.info(f"Random states saved in {output_states_file}")
+    return output_dir
+
+
+def load_accelerator_state(input_dir, models, optimizers, process_index, scaler=None):
+    """
+    Loads states of the `models`, `optimizers`, scaler, and RNG generators from `input_dir`.
body: gt stands for? An expanded name would be easier to read :-)
Since we have `test` below it's probably ground truth, so `ground_truth_rands` or if you wnat shorter `expected_rands`?
diff_hunk: @@ -0,0 +1,144 @@
+# Copyright 2022 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+import argparse
+import os
+import random
+import tempfile
+import unittest
+
+import torch
+from torch import nn
+from torch.utils.data import DataLoader, TensorDataset
+
+from accelerate import Accelerator
+from accelerate.utils import set_seed
+
+
+def dummy_dataloaders(a=2, b=3, batch_size=16, n_train_batches: int = 10, n_valid_batches: int = 2):
+    "Generates a tuple of dummy DataLoaders to test with"
+
+    def get_dataset(n_batches):
+        x = torch.randn(batch_size * n_batches, 1)
+        return TensorDataset(x, a * x + b + 0.1 * torch.randn(batch_size * n_batches, 1))
+
+    train_dataset = get_dataset(n_train_batches)
+    valid_dataset = get_dataset(n_valid_batches)
+    train_dataloader = DataLoader(train_dataset, shuffle=True, batch_size=batch_size, num_workers=4)
+    valid_dataloader = DataLoader(valid_dataset, shuffle=False, batch_size=batch_size, num_workers=4)
+    return (train_dataloader, valid_dataloader)
+
+
+def train(num_epochs, model, dataloader, optimizer, accelerator):
+    "Trains for `num_epochs`"
+    rands = []
+    for epoch in range(num_epochs):
+        # Train quickly
+        model.train()
+        for step, batch in enumerate(dataloader):
+            x, y = batch
+            outputs = model(x)
+            loss = torch.nn.functional.mse_loss(outputs, y)
+            accelerator.backward(loss)
+            optimizer.step()
+            optimizer.zero_grad()
+        rands.append(random.random())  # Introduce some randomness
+    return rands
+
+
+class DummyModel(nn.Module):
+    "Simple model to do y=mx+b"
+
+    def __init__(self):
+        super().__init__()
+        self.a = nn.Parameter(torch.randn(1))
+        self.b = nn.Parameter(torch.randn(1))
+
+    def forward(self, x):
+        return x * self.a + self.b
+
+
+class CheckpointTest(unittest.TestCase):
+    def test_can_resume_training(self, args):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            set_seed(42)
+            model = DummyModel()
+            optimizer = torch.optim.Adam(params=model.parameters(), lr=1e-3)
+            train_dataloader, valid_dataloader = dummy_dataloaders()
+            # Train baseline
+            accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision, device_placement=True)
+            model, optimizer, train_dataloader, valid_dataloader = accelerator.prepare(
+                model, optimizer, train_dataloader, valid_dataloader
+            )
+            # Save initial
+            initial = os.path.join(tmpdir, "initial")
+            accelerator.save_state(initial)
+            model_unwrapped = accelerator.unwrap_model(model)
+            (a, b) = model_unwrapped.a.item(), model_unwrapped.b.item()
+            opt_state = optimizer.state_dict()
+
+            gt_rands = train(3, model, train_dataloader, optimizer, accelerator)
body: No need to make this an executable, pytest will do that on its own. I think we can remove the `args` from the `test_can_resume_training` method.
diff_hunk: @@ -0,0 +1,144 @@
+# Copyright 2022 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+import argparse
+import os
+import random
+import tempfile
+import unittest
+
+import torch
+from torch import nn
+from torch.utils.data import DataLoader, TensorDataset
+
+from accelerate import Accelerator
+from accelerate.utils import set_seed
+
+
+def dummy_dataloaders(a=2, b=3, batch_size=16, n_train_batches: int = 10, n_valid_batches: int = 2):
+    "Generates a tuple of dummy DataLoaders to test with"
+
+    def get_dataset(n_batches):
+        x = torch.randn(batch_size * n_batches, 1)
+        return TensorDataset(x, a * x + b + 0.1 * torch.randn(batch_size * n_batches, 1))
+
+    train_dataset = get_dataset(n_train_batches)
+    valid_dataset = get_dataset(n_valid_batches)
+    train_dataloader = DataLoader(train_dataset, shuffle=True, batch_size=batch_size, num_workers=4)
+    valid_dataloader = DataLoader(valid_dataset, shuffle=False, batch_size=batch_size, num_workers=4)
+    return (train_dataloader, valid_dataloader)
+
+
+def train(num_epochs, model, dataloader, optimizer, accelerator):
+    "Trains for `num_epochs`"
+    rands = []
+    for epoch in range(num_epochs):
+        # Train quickly
+        model.train()
+        for step, batch in enumerate(dataloader):
+            x, y = batch
+            outputs = model(x)
+            loss = torch.nn.functional.mse_loss(outputs, y)
+            accelerator.backward(loss)
+            optimizer.step()
+            optimizer.zero_grad()
+        rands.append(random.random())  # Introduce some randomness
+    return rands
+
+
+class DummyModel(nn.Module):
+    "Simple model to do y=mx+b"
+
+    def __init__(self):
+        super().__init__()
+        self.a = nn.Parameter(torch.randn(1))
+        self.b = nn.Parameter(torch.randn(1))
+
+    def forward(self, x):
+        return x * self.a + self.b
+
+
+class CheckpointTest(unittest.TestCase):
+    def test_can_resume_training(self, args):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            set_seed(42)
+            model = DummyModel()
+            optimizer = torch.optim.Adam(params=model.parameters(), lr=1e-3)
+            train_dataloader, valid_dataloader = dummy_dataloaders()
+            # Train baseline
+            accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision, device_placement=True)
+            model, optimizer, train_dataloader, valid_dataloader = accelerator.prepare(
+                model, optimizer, train_dataloader, valid_dataloader
+            )
+            # Save initial
+            initial = os.path.join(tmpdir, "initial")
+            accelerator.save_state(initial)
+            model_unwrapped = accelerator.unwrap_model(model)
+            (a, b) = model_unwrapped.a.item(), model_unwrapped.b.item()
+            opt_state = optimizer.state_dict()
+
+            gt_rands = train(3, model, train_dataloader, optimizer, accelerator)
+
+            model_unwrapped = accelerator.unwrap_model(model)
+            (a1, b1) = model_unwrapped.a.item(), model_unwrapped.b.item()
+            opt_state1 = optimizer.state_dict()
+
+            # Train partially
+            accelerator.load_state(initial)
+            model_unwrapped = accelerator.unwrap_model(model)
+            (a2, b2) = model_unwrapped.a.item(), model_unwrapped.b.item()
+            opt_state2 = optimizer.state_dict()
+            self.assertEqual(a, a2)
+            self.assertEqual(b, b2)
+            self.assertEqual(opt_state, opt_state2)
+            set_seed(42)
+            train_dataloader, valid_dataloader = dummy_dataloaders()
+            # Ensure all numbers align
+
+            test_rands = train(2, model, train_dataloader, optimizer, accelerator)
+            # Save everything
+            checkpoint = os.path.join(tmpdir, "checkpoint")
+            accelerator.save_state(checkpoint)
+
+            # Load everything back in and make sure all states work
+            accelerator.load_state(checkpoint)
+            test_rands += train(1, model, train_dataloader, optimizer, accelerator)
+
+            model_unwrapped = accelerator.unwrap_model(model)
+            (a3, b3) = model_unwrapped.a.item(), model_unwrapped.b.item()
+            opt_state3 = optimizer.state_dict()
+            self.assertEqual(a1, a3)
+            self.assertEqual(b1, b3)
+            self.assertEqual(opt_state1, opt_state3)
+            self.assertEqual(gt_rands, test_rands)
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Simple example of training script.")
+    parser.add_argument(
+        "--mixed_precision",
+        type=str,
+        default="no",
+        choices=["no", "fp16", "bf16"],
+        help="Whether to use mixed precision. Choose"
+        "between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."
+        "and an Nvidia Ampere GPU.",
+    )
+    parser.add_argument("--cpu", action="store_true", help="If passed, will train on the CPU.")
+    args = parser.parse_args()
+    CheckpointTest().test_can_resume_training(args)
+
+
+if __name__ == "__main__":
+    main()
body: @sgugger how would we delegate the tests to be ran on cpu, mixed precision, etc then in that case? 
diff_hunk: @@ -0,0 +1,144 @@
+# Copyright 2022 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+import argparse
+import os
+import random
+import tempfile
+import unittest
+
+import torch
+from torch import nn
+from torch.utils.data import DataLoader, TensorDataset
+
+from accelerate import Accelerator
+from accelerate.utils import set_seed
+
+
+def dummy_dataloaders(a=2, b=3, batch_size=16, n_train_batches: int = 10, n_valid_batches: int = 2):
+    "Generates a tuple of dummy DataLoaders to test with"
+
+    def get_dataset(n_batches):
+        x = torch.randn(batch_size * n_batches, 1)
+        return TensorDataset(x, a * x + b + 0.1 * torch.randn(batch_size * n_batches, 1))
+
+    train_dataset = get_dataset(n_train_batches)
+    valid_dataset = get_dataset(n_valid_batches)
+    train_dataloader = DataLoader(train_dataset, shuffle=True, batch_size=batch_size, num_workers=4)
+    valid_dataloader = DataLoader(valid_dataset, shuffle=False, batch_size=batch_size, num_workers=4)
+    return (train_dataloader, valid_dataloader)
+
+
+def train(num_epochs, model, dataloader, optimizer, accelerator):
+    "Trains for `num_epochs`"
+    rands = []
+    for epoch in range(num_epochs):
+        # Train quickly
+        model.train()
+        for step, batch in enumerate(dataloader):
+            x, y = batch
+            outputs = model(x)
+            loss = torch.nn.functional.mse_loss(outputs, y)
+            accelerator.backward(loss)
+            optimizer.step()
+            optimizer.zero_grad()
+        rands.append(random.random())  # Introduce some randomness
+    return rands
+
+
+class DummyModel(nn.Module):
+    "Simple model to do y=mx+b"
+
+    def __init__(self):
+        super().__init__()
+        self.a = nn.Parameter(torch.randn(1))
+        self.b = nn.Parameter(torch.randn(1))
+
+    def forward(self, x):
+        return x * self.a + self.b
+
+
+class CheckpointTest(unittest.TestCase):
+    def test_can_resume_training(self, args):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            set_seed(42)
+            model = DummyModel()
+            optimizer = torch.optim.Adam(params=model.parameters(), lr=1e-3)
+            train_dataloader, valid_dataloader = dummy_dataloaders()
+            # Train baseline
+            accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision, device_placement=True)
+            model, optimizer, train_dataloader, valid_dataloader = accelerator.prepare(
+                model, optimizer, train_dataloader, valid_dataloader
+            )
+            # Save initial
+            initial = os.path.join(tmpdir, "initial")
+            accelerator.save_state(initial)
+            model_unwrapped = accelerator.unwrap_model(model)
+            (a, b) = model_unwrapped.a.item(), model_unwrapped.b.item()
+            opt_state = optimizer.state_dict()
+
+            gt_rands = train(3, model, train_dataloader, optimizer, accelerator)
+
+            model_unwrapped = accelerator.unwrap_model(model)
+            (a1, b1) = model_unwrapped.a.item(), model_unwrapped.b.item()
+            opt_state1 = optimizer.state_dict()
+
+            # Train partially
+            accelerator.load_state(initial)
+            model_unwrapped = accelerator.unwrap_model(model)
+            (a2, b2) = model_unwrapped.a.item(), model_unwrapped.b.item()
+            opt_state2 = optimizer.state_dict()
+            self.assertEqual(a, a2)
+            self.assertEqual(b, b2)
+            self.assertEqual(opt_state, opt_state2)
+            set_seed(42)
+            train_dataloader, valid_dataloader = dummy_dataloaders()
+            # Ensure all numbers align
+
+            test_rands = train(2, model, train_dataloader, optimizer, accelerator)
+            # Save everything
+            checkpoint = os.path.join(tmpdir, "checkpoint")
+            accelerator.save_state(checkpoint)
+
+            # Load everything back in and make sure all states work
+            accelerator.load_state(checkpoint)
+            test_rands += train(1, model, train_dataloader, optimizer, accelerator)
+
+            model_unwrapped = accelerator.unwrap_model(model)
+            (a3, b3) = model_unwrapped.a.item(), model_unwrapped.b.item()
+            opt_state3 = optimizer.state_dict()
+            self.assertEqual(a1, a3)
+            self.assertEqual(b1, b3)
+            self.assertEqual(opt_state1, opt_state3)
+            self.assertEqual(gt_rands, test_rands)
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Simple example of training script.")
+    parser.add_argument(
+        "--mixed_precision",
+        type=str,
+        default="no",
+        choices=["no", "fp16", "bf16"],
+        help="Whether to use mixed precision. Choose"
+        "between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."
+        "and an Nvidia Ampere GPU.",
+    )
+    parser.add_argument("--cpu", action="store_true", help="If passed, will train on the CPU.")
+    args = parser.parse_args()
+    CheckpointTest().test_can_resume_training(args)
+
+
+if __name__ == "__main__":
+    main()