Repository: huggingface/accelerate
Pull request: Add in checkpointing capability
Closes https://github.com/huggingface/accelerate/issues/171

This PR adds in two functions, `Accelerator.save_state` and `Accelerator.load_state`, which will go through and checkpoint the current state of everything accelerator touches.

I've included a testing script, but this needs to be properly turned into a test, as currently it is just a script with a bunch of asserts. I have tested successfully on:

- Single GPU
- Single CPU
- Multi-GPU

Problem statement:
Add utility function to save/load all models/optimizers stored in an accelerator
Hi all,
I'm not from the nlp regime, but I love the concept of accelerate. However, I would like to propose some new features.

Since every accelerator.prepare call will essentially store a reference of the object in accelerator. It should be easy to implement a save_all()/load_all() utility that saves/loads all the pytorch module and optimizer state dicts to/from a checkpoint. This utility also has the benifit to save and load the state dict of grad scaler, and potentially any other internal states (which are all agnostic to accelerator users)

I guess the singature could look something like `def save_all(self, filename: str, dict_of_additional_usr_data: Dict[str, Any])`

What do the maintainers think of this utility?

Hints from issue discussion:
Yes, this is on the roadmap and should come in the next few weeks! It will save everything you said plus all the RNG states, so that one can resume training in the exact same conditions. It will also provide an easy way to upload said checkpoint to the Hub so that you can easily resume training on another machine.

I am swamped until the end of next week but hoping to get into adding this the first week of October :-)

Is this feature still on the roadmap?
Yes, I just need the time to dig into it. I have a few things to wrap up in Transformers but plan to work on Accelerate only afterward for a couple of months.

Commit message: More unwrapping

Patch to review:
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index b841bc809f9..e22a5cc8e65 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -14,10 +14,12 @@
 
 import gc
 import os
+import random
 import warnings
 from contextlib import contextmanager
 from typing import List, Optional, Union
 
+import numpy as np
 import torch
 
 from packaging import version
@@ -25,7 +27,7 @@
 from .data_loader import prepare_data_loader
 from .kwargs_handlers import DistributedDataParallelKwargs, GradScalerKwargs, InitProcessGroupKwargs, KwargsHandler
 from .optimizer import AcceleratedOptimizer
-from .state import AcceleratorState, DistributedType, is_deepspeed_available
+from .state import AcceleratorState, DistributedType, is_deepspeed_available, is_tpu_available
 from .utils import (
     DeepSpeedPlugin,
     RNGType,
@@ -40,8 +42,12 @@
 
 if is_deepspeed_available():
     import deepspeed
+
     from .deepspeed_utils import DeepSpeedEngineWrapper, DeepSpeedOptimizerWrapper
 
+if is_tpu_available():
+    import torch_xla.core.xla_model as xm
+
 import logging
 
 
@@ -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"))
+        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")
+
     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/tests/test_state_checkpointing.py b/tests/test_state_checkpointing.py
new file mode 100644
index 00000000000..887ec769753
--- /dev/null
+++ b/tests/test_state_checkpointing.py
@@ -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.")
+    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()
+    test_can_resume_training(args)
+
+
+if __name__ == "__main__":
+    main()


Code review comments:
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}"
path: src/accelerate/accelerator.py
text: @user1:
```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}"
+            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"
path: src/accelerate/accelerator.py
text: @user1:
```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"
+            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}")
path: src/accelerate/accelerator.py
text: @user1:
```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)
path: src/accelerate/accelerator.py
text: @user1:
```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
path: src/accelerate/accelerator.py
text: @user1:
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
+        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
path: src/accelerate/accelerator.py
text: @user1:
```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
+        states = {}
+        states_name = "random_states.pkl"
path: src/accelerate/accelerator.py
text: @user1:
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"
+        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}"
path: src/accelerate/accelerator.py
text: @user1:
```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}"
+            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}"
path: src/accelerate/accelerator.py
text: @user1:
```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}"
+            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"
path: src/accelerate/accelerator.py
text: @user1:
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"
+            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"))
path: src/accelerate/accelerator.py
text: @user1:
Same ehre + one for each process.

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
path: tests/test_state_checkpointing.py
text: @user1:
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
+        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
path: tests/test_state_checkpointing.py
text: @user1:
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
+
+        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.")
path: tests/test_state_checkpointing.py
text: @user1:
This one being deprecated, I wouldn't pass it/create it.

@author:
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: @@ -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)
path: src/accelerate/accelerator.py
text: @user1:
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)
+            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}"
path: src/accelerate/accelerator.py
text: @user1:
```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}"
+            optimizer_name += ".pt"
+            output_optimizer_file = os.path.join(output_dir, optimizer_name)
+            torch.save(state, output_optimizer_file)
path: src/accelerate/accelerator.py
text: @user1:
Same here:
```suggestion
            accelerator.save(state, output_optimizer_file)
```